class_docstr: 2-dimensional grid-world implementation with 5 actions (the 4 cardinal directions and staying put).
docstr: Initialization function for 2-D grid world
  
  Parameters
  ----------
  shape: ``list`` of ``int``, where ``len(shape) == 2``
      The dimensions of the grid world, stored as a list of integers, storing the discrete dimensions of the Y (vertical) and X (horizontal) spatial dimensions, respectively.
  init_state: ``int`` or ``None``
      Initial state of the environment, i.e. the location of the agent in grid world. If not ``None``, must be a discrete index  in the range ``(0, (shape[0] * shape[1])-1)``. It is thus a "linear index" of the initial location of the agent in grid world.
      If ``None``, then an initial location will be randomly sampled from the grid.
docstr_tok: ["Initialization","function","for","2-D","grid","world","Parameters","--","--","--","--","--","shape",":","``","list","``","of","``","int","``",",","where","``","len","(","shape",")","==","2","``","The","dimensions","of","the","grid","world",",","stored","as","a","list","of","integers",",","storing","the","discrete","dimensions","of","the","Y","(","vertical",")","and","X","(","horizontal",")","spatial","dimensions",",","respectively",".","init_state",":","``","int","``","or","``","None","``","Initial","state","of","the","environment",",","i.e",".","the","location","of","the","agent","in","grid","world",".","If","not","``","None","``",",","must","be","a","discrete","index","in","the","range","``","(","0",",","(","shape","[","0","]","*","shape","[","1","]",")","-1",")","``",".","It","is","thus","a","``","linear","index","''","of","the","initial","location","of","the","agent","in","grid","world",".","If","``","None","``",",","then","an","initial","location","will","be","randomly","sampled","from","the","grid","."]
code: def __init__(self, shape=[2, 2], init_state=None):
      """
      Initialization function for 2-D grid world
  
      Parameters
      ----------
      shape: ``list`` of ``int``, where ``len(shape) == 2``
          The dimensions of the grid world, stored as a list of integers, storing the discrete dimensions of the Y (vertical) and X (horizontal) spatial dimensions, respectively.
      init_state: ``int`` or ``None``
          Initial state of the environment, i.e. the location of the agent in grid world. If not ``None``, must be a discrete index  in the range ``(0, (shape[0] * shape[1])-1)``. It is thus a "linear index" of the initial location of the agent in grid world.
          If ``None``, then an initial location will be randomly sampled from the grid.
      """
  
      self.shape = shape
      self.n_states = np.prod(shape)
      self.n_observations = self.n_states
      self.n_control = 5
      self.max_y = shape[0]
      self.max_x = shape[1]
      self._build()
      self.set_init_state(init_state)
      self.last_action = None
  
imports: import numpy
  import matplotlib.pyplot
  import seaborn
  from pymdp.envs import Env
label_desc: Use image node_id 1 to create a new GridWorldEnv object from inherited base classes: Env with example: obj = GridWorldEnv(shape, init_state)