class_docstr: Tasks from the Brax simulator.
class_docstr_tok: ["Tasks","from","the","Brax","simulator","."]
code: def __init__(
      self,
      env_name: str,
      max_steps: int = 1000,
      legacy_spring: bool = True,
      bd_extractor: Optional[BDExtractor] = None,
      test: bool = False,
  ):
      self.max_steps = max_steps
      self.bd_extractor = bd_extractor
      self.test = test
      brax_env = create(
          env_name=env_name,
          episode_length=max_steps,
          legacy_spring=legacy_spring,
      )
      self.obs_shape = tuple(
          [
              brax_env.observation_size,
          ]
      )
      self.act_shape = tuple(
          [
              brax_env.action_size,
          ]
      )
  
      def detect_feet_contact(state, action):
          _, info = brax_env.sys.step(state.qp, action)
          has_contacts = (
              jnp.abs(
                  jnp.clip(info.contact.vel, a_min=-1.0, a_max=1.0)
              ).sum(axis=-1)
              > 0
          )
          return has_contacts[2::2].astype(jnp.int32)
  
      def set_feet_contact_points():
          # TODO: Got values through trail and error, need to verify if actually correct.
          contact_point_mapping = {
              "ant": 4,
              "halfcheetah": 3,
              "humanoid": 5,
              "fetch": 6,
          }
          return contact_point_mapping[env_name]
  
      def reset_fn(key):
          state = brax_env.reset(key)
          if bd_extractor is not None:
              state = get_state_dataclass(
                  state=state,
                  obs=state.obs,
                  feet_contact=jnp.zeros(
                      set_feet_contact_points(), dtype=jnp.int32
                  ),
              )
              state = bd_extractor.init_extended_state(state)
          else:
              state = get_state_dataclass(state=state, obs=state.obs)
          return state
  
      self._reset_fn = jax.jit(jax.vmap(reset_fn))
  
      def step_fn(state, action):
          if bd_extractor is not None:
              feet_contact = detect_feet_contact(state.state, action)
              brax_state = brax_env.step(state.state, action)
              state = state.replace(
                  state=brax_state,
                  obs=brax_state.obs,
                  feet_contact=feet_contact,
              )
              state = bd_extractor.update(
                  state, action, brax_state.reward, brax_state.done
              )
          else:
              brax_state = brax_env.step(state.state, action)
              state = state.replace(
                  state=brax_state, obs=brax_state.obs
              )
          return state, brax_state.reward, brax_state.done
  
      self._step_fn = jax.jit(jax.vmap(step_fn))
  
imports: import sys
  from typing import Optional
  from typing import Tuple
  from typing import Union
  import jax
  import jax.numpy
  from flax.struct import dataclass
  from evojax.task.base import VectorizedTask
  from evojax.task.base import TaskState
  from evojax.task.base import BDExtractor
label_desc: Use image node_id 1 to create a new BraxTask object from inherited base classes: VectorizedTask with example: obj = BraxTask(env_name, max_steps, legacy_spring, bd_extractor, test)