docstr: Test `infer_states` method of the Agent() class
docstr_tok: ["Test","`","infer_states","`","method","of","the","Agent","(",")","class"]
code: def test_agent_infer_states(self):
      """
      Test `infer_states` method of the Agent() class
      """
  
      """ VANILLA method (fixed point iteration) with one hidden state factor and one observation modality """
      num_obs = [5]
      num_states = [3]
      num_controls = [1]
      A = utils.random_A_matrix(num_obs, num_states)
      B = utils.random_B_matrix(num_states, num_controls)
  
      agent = Agent(A=A, B=B, inference_algo="VANILLA")
  
      o = tuple([np.random.randint(obs_dim) for obs_dim in num_obs])
      qs_out = agent.infer_states(o)
  
      qs_validation = inference.update_posterior_states(
          A, o, prior=agent.D
      )
  
      for f in range(len(num_states)):
          self.assertTrue(np.isclose(qs_validation[f], qs_out[f]).all())
  
      """ VANILLA method (fixed point iteration) with multiple hidden state factors and multiple observation modalities """
      num_obs = [2, 4]
      num_states = [2, 3]
      num_controls = [2, 3]
      A = utils.random_A_matrix(num_obs, num_states)
      B = utils.random_B_matrix(num_states, num_controls)
  
      agent = Agent(A=A, B=B, inference_algo="VANILLA")
  
      o = tuple([np.random.randint(obs_dim) for obs_dim in num_obs])
      qs_out = agent.infer_states(o)
  
      qs_validation = inference.update_posterior_states(
          A, o, prior=agent.D
      )
  
      for f in range(len(num_states)):
          self.assertTrue(np.isclose(qs_validation[f], qs_out[f]).all())
  
      """ Marginal message passing inference with one hidden state factor and one observation modality """
      num_obs = [5]
      num_states = [3]
      num_controls = [1]
      A = utils.random_A_matrix(num_obs, num_states)
      B = utils.random_B_matrix(num_states, num_controls)
  
      agent = Agent(A=A, B=B, inference_algo="MMP")
  
      o = tuple([np.random.randint(obs_dim) for obs_dim in num_obs])
      qs_pi_out = agent.infer_states(o)
  
      policies = control.construct_policies(
          num_states, num_controls, policy_len=1
      )
  
      qs_pi_validation, _ = inference.update_posterior_states_full(
          A, B, [o], policies, prior=agent.D, policy_sep_prior=False
      )
  
      for p_idx in range(len(policies)):
          for f in range(len(num_states)):
              self.assertTrue(
                  np.isclose(
                      qs_pi_validation[p_idx][0][f],
                      qs_pi_out[p_idx][0][f],
                  ).all()
              )
  
      """ Marginal message passing inference with multiple hidden state factors and multiple observation modalities """
      num_obs = [2, 4]
      num_states = [2, 2]
      num_controls = [2, 2]
      A = utils.random_A_matrix(num_obs, num_states)
      B = utils.random_B_matrix(num_states, num_controls)
  
      planning_horizon = 3
      backwards_horizon = 1
      agent = Agent(
          A=A,
          B=B,
          inference_algo="MMP",
          policy_len=planning_horizon,
          inference_horizon=backwards_horizon,
      )
      o = [0, 2]
      qs_pi_out = agent.infer_states(o)
  
      policies = control.construct_policies(
          num_states, num_controls, policy_len=planning_horizon
      )
  
      qs_pi_validation, _ = inference.update_posterior_states_full(
          A, B, [o], policies, prior=agent.D, policy_sep_prior=False
      )
  
      for p_idx in range(len(policies)):
          for t in range(planning_horizon + backwards_horizon):
              for f in range(len(num_states)):
                  self.assertTrue(
                      np.isclose(
                          qs_pi_validation[p_idx][t][f],
                          qs_pi_out[p_idx][t][f],
                      ).all()
                  )
  
imports: import os
  import unittest
  import numpy
  from copy import deepcopy
  from pymdp.agent import Agent
  from pymdp import utils, maths
  from pymdp import inference, control, learning
  from pymdp.default_models import generate_grid_world_transitions
label_desc: Use image node_id 5 for calling the TestAgent obj's underlying member method code with example usage: obj.test_agent_infer_states() without return types