[Question] How to get the action space ranges?

[ZaberKo]

Hello,

I'm writing a wrapper for MjxEnv to translate it to another Env API. In brax, we have env.sys.actuator.ctrl_range to get the action range. How to get the action space in MjxEnv? Or, according to #19, should we just use [-1,1]?

[logan-dunbar]

Is this documented anywhere? I can't seem to find any wrapper or test_fn that checks if the actions are within range.

I also found the following in the arxiv doc that seems to suggest this might not always be the case:

B.25 Finetuning
Joystick policy

1. Train for 100 M timesteps with a command range of {1.5,0.8,1.2}.
2. Finetune for 50 M timesteps with a command range of {1.5,0.8,2⁢π}.
3. Finetune on rough terrain for 100 M timesteps.

And the section under B.21 Observation and Action doesn't make any range explicit either.

[logan-dunbar]

Digging further, it seems all the dm_control_suite envs have a similar line in their xml doc: <motor ctrllimited="true" ctrlrange="-1 1"/>, but I don't see any checking or enforcement, the action is just passed straight through to mujoco.

class Run(mjx_env.MjxEnv):
  """Cheetah running environment."""
  ...
  def step(self, state: mjx_env.State, action: jax.Array) -> mjx_env.State:
    data = mjx_env.step(self.mjx_model, state.data, action, self.n_substeps)

...

# mjx_env.py
def step(
    model: mjx.Model,
    data: mjx.Data,
    action: jax.Array,
    n_substeps: int = 1,
) -> mjx.Data:
  def single_step(data, _):
    data = data.replace(ctrl=action)
    data = mjx.step(model, data)

However, the locomotion envs seem to have arbitrary command ranges:
Apollo:

      command_config=config_dict.create(
          min=[-1.5, -0.8, -1.5],
          max=[1.5, 0.8, 1.5],
          zero_prob=[0.9, 0.25, 0.5],
      ),

Barkour:

  def sample_command(self, rng: jax.Array) -> jax.Array:
    rng1, rng2, rng3 = jax.random.split(rng, 3)

    lin_vel_x = jax.random.uniform(
        rng1, minval=self._config.lin_vel_x[0], maxval=self._config.lin_vel_x[1]
    )
    lin_vel_y = jax.random.uniform(
        rng2, minval=self._config.lin_vel_y[0], maxval=self._config.lin_vel_y[1]
    )
    ang_vel_yaw = jax.random.uniform(
        rng3,
        minval=self._config.ang_vel_yaw[0],
        maxval=self._config.ang_vel_yaw[1],
    )

    return jp.hstack([lin_vel_x, lin_vel_y, ang_vel_yaw])

# which has as default config
def default_config() -> config_dict.ConfigDict:
  return config_dict.create(
      ...
      lin_vel_x=[-0.6, 1.5],
      lin_vel_y=[-0.8, 0.8],
      ang_vel_yaw=[-0.7, 0.7],
      ...

Is there a consistent way to access this information like env.action_space from gymnasium style environments?