docstr: This unit test runs the a concise version of the code in the `tmaze_demo.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.ipynb","`","tutorial","notebook","to","make","sure","it","works","if","things","are","changed"]
code: def test_tmaze_demo(self):
      """
      This unit test runs the a concise version of the code in the `tmaze_demo.ipynb` tutorial notebook to make sure it works if things are changed
      """
  
      reward_probabilities = [
          0.98,
          0.02,
      ]  # probabilities used in the original SPM T-maze demo
      env = TMazeEnv(reward_probs=reward_probabilities)
  
      """test plotting of the observation likelihood (just plot one slice)"""
      A_gp = env.get_likelihood_dist()
      plot_likelihood(A_gp[1][:, :, 0], "Reward Right")
  
      """test plotting of the transition likelihood (just plot one slice)"""
      B_gp = env.get_transition_dist()
      plot_likelihood(B_gp[1][:, :, 0], "Reward Condition Transitions")
  
      A_gm = copy.deepcopy(
          A_gp
      )  # make a copy of the true observation likelihood to initialize the observation model
      B_gm = copy.deepcopy(
          B_gp
      )  # make a copy of the true transition likelihood to initialize the transition model
  
      control_fac_idx = [0]
      agent = Agent(A=A_gm, B=B_gm, control_fac_idx=control_fac_idx)
      plot_beliefs(agent.D[0], "Beliefs about initial location")
  
      agent.C[1][1] = 3.0  # they like reward
      agent.C[1][2] = -3.0  # they don't like loss
  
      T = 5  # number of timesteps
  
      obs = (
          env.reset()
      )  # reset the environment and get an initial observation
  
      # these are useful for displaying read-outs during the loop over time
      reward_conditions = ["Right", "Left"]
      location_observations = [
          "CENTER",
          "RIGHT ARM",
          "LEFT ARM",
          "CUE LOCATION",
      ]
      reward_observations = ["No reward", "Reward!", "Loss!"]
      cue_observations = ["Cue Right", "Cue Left"]
  
      for t in range(T):
          qx = agent.infer_states(obs)
  
          q_pi, G = agent.infer_policies()
  
          action = agent.sample_action()
  
          obs = env.step(action)
  
          if int(action[0]) == 3:
              # if the reward condition is Reward on RIGHT
              if env.reward_condition == 0:
                  self.assertEqual(
                      obs[2], 0
                  )  # this tests that the cue observation is 'Cue Right' in case of 'Reward on Right' condition
  
              # if the reward condition is Reward on RIGHT
              if env.reward_condition == 1:
                  self.assertEqual(
                      obs[2], 1
                  )  # this tests that the cue observation is 'Cue Left' in case of 'Reward on Left' condition
  
      plot_beliefs(
          qx[1], "Final posterior beliefs about reward condition"
      )
  
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 2 for calling the TestDemos obj's underlying member method code with example usage: obj.test_tmaze_demo() without return types