class_docstr: Wrapper class for Brax environments.
class_docstr_tok: ["Wrapper","class","for","Brax","environments","."]
code: def __init__(self, env_name: str, seed: Optional[int] = None):
      self.env = envs.create(env_name=env_name)
      self.game_over = False
      seed = int(time.time() * 1e6) if seed is None else seed
      self._rng = jax.random.PRNGKey(seed)
      self._rng, rng = jax.random.split(self._rng)
      self._state = self.env.reset(rng=rng)
  
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 1 to create a new BraxEnv object from inherited base classes: object with example: obj = BraxEnv(env_name, seed)