class_docstr: 2-dimensional grid-world implementation with 5 actions (the 4 cardinal directions and staying put).
code: def get_transition_dist(self):
      B = np.zeros([self.n_states, self.n_states, self.n_control])
      for s in range(self.n_states):
          for a in range(self.n_control):
              ns = int(self.P[s][a])
              B[ns, s, a] = 1
      return B
  
imports: import numpy
  import matplotlib.pyplot
  import seaborn
  from pymdp.envs import Env
label_desc: Use image node_id 9 for calling the GridWorldEnv obj's underlying member method code with example usage: obj.get_transition_dist() and returns: B