Validating Sophisticated Inference (SI) Planning Algorithm using the T-Maze Task¶

import sys
if "google.colab" in sys.modules:
    %pip install "inferactively-pymdp[nb]" -q

%load_ext autoreload
%autoreload 2

import jax.numpy as jnp
import jax.random as jr

import matplotlib.pyplot as plt
import numpy as np
import mediapy

from pymdp.envs import SimplifiedTMaze, rollout
from pymdp.agent import Agent
from pymdp.planning.si import si_policy_search

# setting the parameters for the environment
reward_condition = 0 # 0 is reward in left arm, 1 is reward in right arm, None is random allocation
cue_validity = 0.95 # 95% valid cues

reward_probability = 1.0 # 100% chance of reward in the correct arm
dependent_outcomes = True # if True, punishment occurs as a function of reward probability (i.e., if reward probability is 0.8, then 20% punishment). If False, punishment occurs with set probability (i.e., 20% no outcome and punishment will only occur in the other (non-rewarding) arm depending on the punishment_probability parameter)

# initializing the environment. see si_tmaze.py in pymdp/envs for the implementation details
env = SimplifiedTMaze(
    reward_condition=reward_condition,
    cue_validity=cue_validity,  
    reward_probability=reward_probability,     
    dependent_outcomes=dependent_outcomes,
)

# creating C tensors filled with zeros for [location], [reward], [cue] based on A shapes for the Agent
C = [jnp.zeros(a.shape[0], dtype=jnp.float32) for a in env.A] 

# setting preferences for outcomes only
C[1] = C[1].at[1].set(2.0)    # prefer reward
C[1] = C[1].at[2].set(-6.0)   # avoid punishment

# slight cost of observing a cue
# C[0] = C[0].at[3].set(-1.0).at[4].set(-1.0)

# flat D tensors [location], [reward] based on B shapes for the agent
D = [jnp.ones(b.shape[0], dtype=jnp.float32) / b.shape[0] for b in env.B]

key = jr.PRNGKey(0) 
T = 3

si_search = si_policy_search(
    horizon=policy_len, # plans 2 timesteps ahead
    max_nodes=5000, # maximum number of nodes allowed in the tree
    max_branching=10, # maximum number of children allowed per node (moderating the branching factor)
    policy_prune_threshold=0.0, # no pruning of unlikely policies
    observation_prune_threshold=0.0, # no pruning of unlikely observations
    entropy_stop_threshold=0.0, # disabling halting of expansion if agent is certain enough
    neg_efe_stop_threshold=1e10, # disabling efe value based halting of expansion
    kl_threshold=-1, # disabling node reuse if agent is in similar states after an action
    prune_penalty=512, # default value for prune penalty
    gamma=3, # temperature parameter; lower value (--> 1) prunes policies less aggressively as probabilities are flattened while higher value (--> 16) prunes more aggressively
    topk_obsspace=10000, # max number of top observation combinations - this default value just means we want to consider all the observation combinations
)

_, info_vanilla = rollout(agent_vanilla, env, num_timesteps=T, rng_key=key) # default policy search is vanilla
_, info_si = rollout(agent_si, env, num_timesteps=T, rng_key=key, policy_search=si_search)

def make_gif(info):
    frames = []
    for t in range(info["observation"][0].shape[1]):  # iterate over timesteps
        # get observations for this timestep
        observations_t = [
            info["observation"][0][:, t, :],
            info["observation"][1][:, t, :],  
        ]
        
        frame = env.render(mode="rgb_array", observations=observations_t) # render the environment using the observations for this timestep
        frame = np.asarray(frame, dtype=np.uint8)
        plt.close()  # close the figure to prevent memory leak
        frames.append(frame)

    frames = np.array(frames, dtype=np.uint8)
    mediapy.show_video(frames, fps=1)

# qpi is a posterior over whole policies (action sequences).
# To get the probability of the *current* action, we marginalize over policies
# that share the same first action and sum their qpi values.
# helper functions for:
# - printing out policies and respective probabilities of selecting those policies
# - printing out action and observation info for each timestep

np.set_printoptions(precision=2, suppress=True)

def print_qpi(agent, info, print_t1=False):
    qpi_values = info["qpi"]

    action_names = {
        0: "move to center",
        1: "move to left arm",
        2: "move to right arm",
        3: "move to cue",
    }
    max_timesteps = 1 if print_t1 else qpi_values.shape[1]

    # unique_multiactions returns the unique first-step actions across policies
    # for a single control factor, this is just a list of action indices
    unique_actions = agent.unique_multiactions[:, 0]

    for t in range(max_timesteps):
        print(f"Timestep {t}:")
        action_probs = agent.multiaction_probabilities(qpi_values[:, t, :])[0]

        for action_idx, total_prob in zip(unique_actions.tolist(), action_probs.tolist()):
            if action_idx < 0:
                continue
            action_name = action_names.get(action_idx, f"action_{action_idx}")
            print(f"  {action_name}: {total_prob:.3f}")
        print()

def print_agent_behavior(info):

    action_names = {0: "move to center", 1: "move to left arm", 2: "move to right arm", 3: "move to cue"}

    location_obs = {0: "center loc", 1: "left arm loc", 2: "right arm loc", 3: "cue-left-arm", 4: "cue-right-arm"}
    outcome_obs = {0: "no_outcome", 1: "reward", 2: "punishment"}

    actions = info["action"]
    observations = info["observation"]

    num_timesteps = actions.shape[1]

    for t in range(num_timesteps):
        action_idx = int(actions[0, t, 0])  # [batch, timestep, action_dim]
        action_name = action_names.get(action_idx, f"action_{action_idx}")

        location_obs_idx = int(observations[0][0, t, 0])  # [modality][batch, timestep, obs_dim]
        outcome_obs_idx = int(observations[1][0, t, 0])

        location_name = location_obs.get(location_obs_idx)
        outcome_name = outcome_obs.get(outcome_obs_idx)

        print(f"t={t}: observed=({location_name}, {outcome_name}) -> action={action_name}")

from pymdp.envs import TMaze

# setting the parameters for the environment
reward_condition = 0 # 0 is reward in left arm, 1 is reward in right arm, None is random allocation
cue_validity = 1.0 # 100% valid cues (cue probability)

reward_probability = 1.0 # 100% chance of reward in the correct arm
dependent_outcomes = True # if True, punishment occurs as a function of reward probability (i.e., if reward probability is 0.8, then 20% punishment). If False, punishment occurs with set probability (i.e., 20% no outcome and punishment will only occur in the other (non-rewarding) arm depending on the punishment_probability parameter)
punishment_probability = 1.0 # 100% chance of punishment in the other arm

# initializing the environment. see tmaze.py in pymdp/envs for the implementation details
env = TMaze(
    reward_condition=reward_condition,
    cue_validity=cue_validity,  
    reward_probability=reward_probability,
    punishment_probability=punishment_probability,     
    dependent_outcomes=dependent_outcomes,
)

# creating C tensors filled with zeros for [location], [reward], [cue] based on A shapes for the Agent
C = [jnp.zeros(a.shape[0], dtype=jnp.float32) for a in env.A] 

# setting preferences for outcomes only
C[1] = C[1].at[1].set(2.0)    # prefer reward
C[1] = C[1].at[2].set(-6.0)   # avoid punishment

# slight cost of observing a cue
C[2] = C[2].at[1].set(-0.5) 
C[2] = C[2].at[2].set(-0.5)

# D tensors [location], [reward] based on B shapes for the agent
# - agent starts in the center location
# - reward location prior is uniform
D_loc = jnp.zeros(env.B[0].shape[0], dtype=jnp.float32)
D_loc = D_loc.at[0].set(1.0)

D_reward = jnp.ones(env.B[1].shape[0], dtype=jnp.float32)
D_reward = D_reward / jnp.sum(D_reward, axis=0, keepdims=True)

D = [D_loc, D_reward]

key = jr.PRNGKey(0) 
T = 5

si_search = si_policy_search(
    horizon=policy_len, # plans 4 timesteps ahead
    max_nodes=5000, # maximum number of nodes allowed in the tree
    max_branching=45, # maximum number of children allowed per node (moderating the branching factor)
    policy_prune_threshold=0.0, # no pruning of unlikely policies
    observation_prune_threshold=1e-4, # no pruning of unlikely observations
    entropy_stop_threshold=0.0, # disabling halting of expansion if agent is certain enough
    neg_efe_stop_threshold=1e10, # disabling efe value based halting of expansion
    kl_threshold=-1, # disabling node reuse if agent is in similar states after an action
    prune_penalty=512, # default value for prune penalty
    gamma=gamma, # temperature parameter; lower value (---> 1) prunes policies less aggressively as probabilities are flattened while higher value (---> 16) prunes more aggressively
    topk_obsspace=10000, # max number of top observation combinations - this default value just means we want to consider all the observation combinations
)

_, info_vanilla = rollout(agent_vanilla, env, num_timesteps=T, rng_key=key) # default policy search is vanilla
_, info_si = rollout(agent_si, env, num_timesteps=T, rng_key=key, policy_search=si_search)

def make_gif(info):
    frames = []
    for t in range(info["observation"][0].shape[1]):  # iterate over timesteps
        # get observations for this timestep
        observations_t = [
            info["observation"][0][:, t, :],
            info["observation"][1][:, t, :],  
            info["observation"][2][:, t, :],
        ]
        
        frame = env.render(mode="rgb_array", observations=observations_t) # render the environment using the observations for this timestep
        frame = np.asarray(frame, dtype=np.uint8)
        plt.close()  # close the figure to prevent memory leak
        frames.append(frame)

    frames = np.array(frames, dtype=np.uint8)
    mediapy.show_video(frames, fps=1)

# qpi is a posterior over whole policies (action sequences).
# To get the probability of the *current* action, we marginalize over policies
# that share the same first action and sum their qpi values.
# helper functions for:
# - printing out policies and respective probabilities of selecting those policies
# - printing out action and observation info for each timestep

np.set_printoptions(precision=2, suppress=True)

def print_qpi(agent, info, print_t1=False):
    qpi_values = info["qpi"]

    action_names = {
        0: "move to center",
        1: "move to left arm",
        2: "move to right arm",
        3: "move to cue",
        4: "move to middle",
    }
    max_timesteps = 1 if print_t1 else qpi_values.shape[1]

    # unique_multiactions returns the unique first-step actions across policies
    # for a single control factor, this is just a list of action indices
    unique_actions = agent.unique_multiactions[:, 0]

    for t in range(max_timesteps):
        print(f"Timestep {t}:")
        action_probs = agent.multiaction_probabilities(qpi_values[:, t, :])[0]

        for action_idx, total_prob in zip(unique_actions.tolist(), action_probs.tolist()):
            if action_idx < 0:
                continue
            action_name = action_names.get(action_idx, f"action_{action_idx}")
            print(f"  {action_name}: {total_prob:.3f}")
        print()

def print_agent_behavior(info):

    action_names = {
        0: "move to center",
        1: "move to left arm",
        2: "move to right arm",
        3: "move to cue",
        4: "move to middle",
    }
    
    location_obs = {
        0: "center loc",
        1: "left arm loc",
        2: "right arm loc",
        3: "cue loc",
        4: "middle loc",
    }
    outcome_obs = {0: "no_outcome", 1: "reward", 2: "punishment"}
    cue_obs = {0: "no cue", 1: "cue-left", 2: "cue-right"}
    
    actions = info["action"]
    observations = info["observation"]
    
    num_timesteps = actions.shape[1]
    
    for t in range(num_timesteps):
        action_idx = int(actions[0, t, 0])  # [batch, timestep, action_dim]
        action_name = action_names.get(action_idx, f"action_{action_idx}")
        
        location_obs_idx = int(observations[0][0, t, 0])  # [modality][batch, timestep, obs_dim]
        outcome_obs_idx = int(observations[1][0, t, 0])
        cue_obs_idx = int(observations[2][0, t, 0])
        
        location_name = location_obs.get(location_obs_idx)
        outcome_name = outcome_obs.get(outcome_obs_idx)
        cue_name = cue_obs.get(cue_obs_idx)
        
        print(f"t={t}: observed=({location_name}, {outcome_name}, {cue_name}) -> action={action_name}")