class_docstr: Wrapper class for Brax environments.
class_docstr_tok: ["Wrapper","class","for","Brax","environments","."]
code: def step(
      self, action
  ) -> Tuple[onp.ndarray, onp.ndarray, bool, Mapping[Any, Any]]:
      self._state = jax.jit(self.env.step)(self._state, action)
      self.game_over = self._state.done
      return (
          onp.array(self._state.obs),
          onp.array(self._state.reward),
          self._state.done,
          self._state.info,
      )
  
imports: import time
  from typing import Any, Mapping, Optional, Tuple
  from brax import envs
  from dopamine.continuous_domains import run_experiment
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.jax.agents.sac import sac_agent
  from flax.metrics import tensorboard
  import gin
  import jax
  import numpy
label_desc: Use image node_id 7 for calling the BraxEnv obj's underlying member method code with example usage: obj.step(action) and returns: onp,  onp,  self,  self