docstr: This unit test runs the a concise version of the code in the `tmaze_demo_learning.ipynb` tutorial notebook to make sure it works if things are changed
docstr_tok: ["This","unit","test","runs","the","a","concise","version","of","the","code","in","the","`","tmaze_demo_learning.ipynb","`","tutorial","notebook","to","make","sure","it","works","if","things","are","changed"]
code: def test_tmaze_learning_demo(self):
      """
      This unit test runs the a concise version of the code in the `tmaze_demo_learning.ipynb` tutorial notebook to make sure it works if things are changed
      """
  
      reward_probabilities = [
          0.85,
          0.15,
      ]  # the 'true' reward probabilities
      env = TMazeEnvNullOutcome(reward_probs=reward_probabilities)
      A_gp = env.get_likelihood_dist()
      B_gp = env.get_transition_dist()
  
      pA = utils.dirichlet_like(A_gp, scale=1e16)
  
      pA[1][1:, 1:3, :] = 1.0
  
      A_gm = utils.norm_dist_obj_arr(pA)
  
      B_gm = copy.deepcopy(B_gp)
  
      controllable_indices = [
          0
      ]  # this is a list of the indices of the hidden state factors that are controllable
      learnable_modalities = [
          1
      ]  # this is a list of the modalities that you want to be learn-able
  
      agent = Agent(
          A=A_gm,
          pA=pA,
          B=B_gm,
          control_fac_idx=controllable_indices,
          modalities_to_learn=learnable_modalities,
          lr_pA=0.25,
          use_param_info_gain=True,
      )
  
      agent.D[0] = utils.onehot(0, agent.num_states[0])
      agent.C[1][1] = 2.0
      agent.C[1][2] = -2.0
  
      T = 1000  # number of timesteps
  
      obs = (
          env.reset()
      )  # reset the environment and get an initial observation
  
      for t in range(T):
          agent.infer_states(obs)
          agent.infer_policies()
          action = agent.sample_action()
          agent.update_A(obs)
          obs = env.step(action)
  
      # make sure they are learning the reward contingencies in the right general direction
  
      REWARD_ON_RIGHT = 0
      REWARD_ON_LEFT = 1
  
      REWARD = 1
      PUNISHMENT = 2
  
      RIGHT_ARM = 1
      LEFT_ARM = 2
  
      # in case the reward condition is 'Reward on RIGHT'
  
      if env.reward_condition == REWARD_ON_RIGHT:
          prob_reward_right = agent.A[1][
              REWARD, RIGHT_ARM, REWARD_ON_RIGHT
          ]
          prob_punishment_right = agent.A[1][
              PUNISHMENT, RIGHT_ARM, REWARD_ON_RIGHT
          ]
  
          self.assertGreater(prob_reward_right, prob_punishment_right)
  
      # in case the reward condition is 'Reward on LEFT'
  
      elif env.reward_condition == REWARD_ON_LEFT:
          prob_reward_left = agent.A[1][
              REWARD, LEFT_ARM, REWARD_ON_LEFT
          ]
          prob_punishment_left = agent.A[1][
              PUNISHMENT, LEFT_ARM, REWARD_ON_LEFT
          ]
          self.assertGreater(prob_reward_left, prob_punishment_left)
  
file_path: pymdp/test/test_demos.py
imports: import unittest
  import numpy
  import copy
  import seaborn
  import matplotlib.pyplot
  from pymdp.agent import Agent
  from pymdp.utils import plot_beliefs, plot_likelihood
  from pymdp import utils, maths, default_models
  from pymdp import control
  from pymdp.envs import TMazeEnv, TMazeEnvNullOutcome
  from copy import deepcopy
label_desc: Use image node_id 3 for calling the TestDemos obj's underlying member method code with example usage: obj.test_tmaze_learning_demo() without return types