class_docstr: Environment wrapper for Multigrid environments.
  
  Note: the main difference with vanilla GymWrapper is that reward_spec() is
  overridden and rewards are cast to np.arrays in step()
docstr: Steps the environment.
docstr_tok: ["Steps","the","environment","."]
code: def step(self, action: types.NestedArray) -> dm_env.TimeStep:
      """Steps the environment."""
      if self._reset_next_step:
          return self.reset()
  
      observation, reward, done, info = self._environment.step(action)
      observation = self.process_obs(observation)
  
      self._reset_next_step = done
      self._last_info = info
  
      def _map_reward_spec(x, t):
          if np.isscalar(x):
              return t.dtype.type(x)
          return np.asarray(x, dtype=t.dtype)
  
      reward = tree.map_structure(
          _map_reward_spec, reward, self.reward_spec()
      )
  
      if done:
          truncated = info.get("TimeLimit.truncated", False)
          if truncated:
              return dm_env.truncation(reward, observation)
          return dm_env.termination(reward, observation)
      return dm_env.transition(reward, observation)
  
imports: from typing import Any, Dict, List, Optional
  import warnings
  from acme import specs
  from acme import types
  from acme import wrappers
  from acme.multiagent import types
  from acme.wrappers import multiagent_dict_key_wrapper
  import dm_env
  import gym
  from gym import spaces
  import jax
  import numpy
  import tree
label_desc: Use image node_id 4 for calling the MultigridWrapper obj's underlying member method code with example usage: obj.step(action) and returns: dm_env, self, np, dm_env, t, dm_env