code: def calc_states_info_gain(A, qs_pi):
      """
      Computes the Bayesian surprise or information gain about states of a policy,
      using the observation model and the hidden state distribution expected under that policy.
  
      Parameters
      ----------
      A: ``numpy.ndarray`` of dtype object
          Sensory likelihood mapping or 'observation model', mapping from hidden states to observations. Each element ``A[m]`` of
          stores an ``numpy.ndarray`` multidimensional array for observation modality ``m``, whose entries ``A[m][i, j, k, ...]`` store
          the probability of observation level ``i`` given hidden state levels ``j, k, ...``
      qs_pi: ``list`` of ``numpy.ndarray`` of dtype object
          Predictive posterior beliefs over hidden states expected under the policy, where ``qs_pi[t]`` stores the beliefs about
          hidden states expected under the policy at time ``t``
  
      Returns
      -------
      states_surprise: float
          Bayesian surprise (about states) or salience expected under the policy in question
      """
  
      n_steps = len(qs_pi)
  
      states_surprise = 0
      for t in range(n_steps):
          states_surprise += spm_MDP_G(A, qs_pi[t])
  
      return states_surprise
  
imports: import itertools
  import numpy
  from pymdp.maths import softmax, softmax_obj_arr, spm_dot, spm_wnorm, spm_MDP_G, spm_log_single, spm_log_obj_array
  from pymdp import utils
  import copy
label_desc: Use image node_id 6 for calling a global function with example usage: calc_states_info_gain(A, qs_pi) and returns: states_surprise