docstr: Test updating prior Dirichlet parameters over initial hidden states (pD) with the ``Agent`` class,
  in the case that you're using "vanilla" inference mode.
docstr_tok: ["Test","updating","prior","Dirichlet","parameters","over","initial","hidden","states","(","pD",")","with","the","``","Agent","``","class",",","in","the","case","that","you","'re","using","``","vanilla","''","inference","mode","."]
code: def test_agent_with_D_learning_vanilla(self):
      """
      Test updating prior Dirichlet parameters over initial hidden states (pD) with the ``Agent`` class,
      in the case that you're using "vanilla" inference mode.
      """
  
      num_obs = [2, 4]
      num_states = [2, 3]
      num_controls = [1, 1]  # HMM mode
      A = utils.random_A_matrix(num_obs, num_states)
      B = utils.random_B_matrix(num_states, num_controls)
      pD = utils.dirichlet_like(
          utils.random_single_categorical(num_states)
      )
  
      # 1. Test that the updating works when `save_belief_hist` is True, and you don't need to pass in the beliefs about first hidden states
      agent = Agent(
          A=A,
          B=B,
          pD=pD,
          inference_algo="VANILLA",
          save_belief_hist=True,
      )
  
      T = 10
  
      qs_history = []
  
      for t in range(T):
          # get some random hidden state distribution
          p_states = utils.random_single_categorical(num_states)
          observation = [
              utils.sample(maths.spm_dot(A_g, p_states)) for A_g in A
          ]
  
          qs_t = agent.infer_states(observation)
  
          qs_history.append(qs_t)
  
      pD_validation = learning.update_state_prior_dirichlet(
          agent.pD,
          qs_history[0],
          agent.lr_pD,
          factors=agent.factors_to_learn,
      )
  
      pD_test = agent.update_D()
  
      for factor in range(len(num_states)):
          self.assertTrue(
              np.allclose(pD_test[factor], pD_validation[factor])
          )
  
      # 2. Test that the updating works when `save_belief_hist` is False, and you do have to pass in the beliefs about first hidden states
      agent = Agent(
          A=A,
          B=B,
          pD=pD,
          inference_algo="VANILLA",
          save_belief_hist=False,
      )
  
      T = 10
  
      qs_history = []
  
      for t in range(T):
          # get some random hidden state distribution
          p_states = utils.random_single_categorical(num_states)
          observation = [
              utils.sample(maths.spm_dot(A_g, p_states)) for A_g in A
          ]
  
          qs_t = agent.infer_states(observation)
  
          qs_history.append(qs_t)
  
      pD_validation = learning.update_state_prior_dirichlet(
          agent.pD,
          qs_history[0],
          agent.lr_pD,
          factors=agent.factors_to_learn,
      )
  
      pD_test = agent.update_D(qs_t0=qs_history[0])
  
      for factor in range(len(num_states)):
          self.assertTrue(
              np.allclose(pD_test[factor], pD_validation[factor])
          )
  
      # 3. Same as test #1, except with learning on only certain hidden state factors. Also passed in a different learning rate
      agent = Agent(
          A=A,
          B=B,
          pD=pD,
          lr_pD=2.0,
          inference_algo="VANILLA",
          save_belief_hist=True,
          factors_to_learn=[0],
      )
  
      T = 10
  
      qs_history = []
  
      for t in range(T):
          # get some random hidden state distribution
          p_states = utils.random_single_categorical(num_states)
          observation = [
              utils.sample(maths.spm_dot(A_g, p_states)) for A_g in A
          ]
  
          qs_t = agent.infer_states(observation)
  
          qs_history.append(qs_t)
  
      pD_validation = learning.update_state_prior_dirichlet(
          agent.pD,
          qs_history[0],
          agent.lr_pD,
          factors=agent.factors_to_learn,
      )
  
      pD_test = agent.update_D()
  
      for factor in range(len(num_states)):
          self.assertTrue(
              np.allclose(pD_test[factor], pD_validation[factor])
          )
  
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 8 for calling the TestAgent obj's underlying member method code with example usage: obj.test_agent_with_D_learning_vanilla() without return types