code: def create_brax_agent(
      environment: BraxEnv,
      agent_name: str = "sac_brax",
      summary_writer: Optional[tensorboard.SummaryWriter] = None,
  ) -> dqn_agent.JaxDQNAgent:
      """Creates an agent for Brax."""
      assert agent_name is not None
      if agent_name == "sac_brax":
          return sac_agent.SACAgent(
              action_shape=(environment.action_space,),
              action_limits=(
                  -1 * onp.ones(environment.action_space),
                  onp.ones(environment.action_space),
              ),
              observation_shape=environment.observation_space,
              action_dtype=onp.float32,
              observation_dtype=onp.float64,
              summary_writer=summary_writer,
          )
      else:
          raise ValueError(f"Unknown agent: {agent_name}")
  
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 2 for calling a global function with example usage: create_brax_agent(environment, agent_name, summary_writer) and returns: sac_agent