Domain randomization with mjx backend #523

yardenas · 2024-07-18T17:13:55Z

yardenas
Jul 18, 2024

I have the following test file, which I use to test whether my domain randomization implementation works. The idea is that I condition (randomized) environments on the same action and initial state, and expect to get perturbed state when taking one step forward.

from typing import Any
import jax

import jax.numpy as jnp
from brax import base
from brax.io import mjcf
from brax.envs import register_environment
from brax.envs.base import PipelineEnv, State
from brax import envs

xml = """
<!-- https://github.com/google-deepmind/dm_control/blob/main/dm_control/suite/cartpole.xml -->
<mujoco model="cart-pole">
  <option timestep="0.01">
    <flag contact="disable"/>
  </option>
  <default>
    <default class="pole">
      <joint type="hinge" axis="0 1 0"  damping="2e-6"/>
      <geom type="capsule" fromto="0 0 0 0 0 1" size="0.045" mass=".1"/>
    </default>
  </default>
  <worldbody>
    <light name="light" pos="0 0 6"/>
    <camera name="fixed" pos="0 -4 1" zaxis="0 -1 0"/>
    <camera name="lookatcart" mode="targetbody" target="cart" pos="0 -2 2"/>
    <geom name="floor" pos="0 0 -.05" size="4 4 .2" type="plane"/>
    <geom name="rail1" type="capsule" pos="0  .07 1" zaxis="1 0 0" size="0.02 2"/>
    <geom name="rail2" type="capsule" pos="0 -.07 1" zaxis="1 0 0" size="0.02 2"/>
    <body name="cart" pos="0 0 1">
      <joint name="slider" type="slide" limited="true" axis="1 0 0" range="-1.8 1.8" solreflimit=".08 1" damping="5e-4"/>
      <geom name="cart" type="box" size="0.2 0.15 0.1" mass="1"/>
      <body name="pole_1" childclass="pole">
        <joint name="hinge_1"/>
        <geom name="pole_1"/>
      </body>
    </body>
  </worldbody>
  <actuator>
    <motor name="slide" joint="slider" gear="10" ctrllimited="true" ctrlrange="-1 1" />
  </actuator>
</mujoco>
"""


_PARALLEL_ENVS = 128


def domain_randomization(sys, rng):
    @jax.vmap
    def randomize(rng):
        cpole = jax.random.normal(rng) * 1. + sys.actuator.gear[0]
        mass = sys.actuator.gear.at[0].set(cpole)
        return mass, cpole

    mass, samples = randomize(rng)
    in_axes = jax.tree_map(lambda x: None, sys)
    in_axes = in_axes.tree_replace({"actuator.gear": 0})
    sys = sys.tree_replace({"actuator.gear": mass})
    return sys, in_axes, samples[:, None]


class Cartpole(PipelineEnv):
    def __init__(self, backend="mjx", **kwargs):
        sys = mjcf.loads(xml)
        self.sparse = kwargs.pop("sparse", False)
        self.swingup = kwargs.pop("swingup", False)
        super().__init__(sys=sys, backend=backend, n_frames=1, **kwargs)

    def reset(self, rng: jax.Array) -> State:
        """Resets the environment to an initial state."""
        rng, rng1, rng2 = jax.random.split(rng, 3)
        if self.swingup:
            q = self.sys.init_q + jax.random.normal(rng1, (self.sys.q_size(),)) * 0.01
            q = q.at[1].add(jnp.pi)
        else:
            q = self.sys.init_q
            q = q.at[0].set(jax.random.uniform(rng1, shape=(), minval=-1.0, maxval=1.0))
            q = q.at[1].set(
                jax.random.uniform(rng1, shape=(), minval=-0.034, maxval=0.034)
            )
        qd = jax.random.normal(rng2, (self.sys.qd_size(),)) * 0.01
        pipeline_state = self.pipeline_init(q, qd)
        obs = self._get_obs(pipeline_state)
        reward, done = jnp.zeros(2)
        metrics: dict[str, Any] = {}
        return State(pipeline_state, obs, reward, done, metrics)

    def step(self, state: State, action: jax.Array) -> State:
        """Run one timestep of the environment's dynamics."""
        # Scale action from [-1,1] to actuator limits
        action_min = self.sys.actuator.ctrl_range[:, 0]
        action_max = self.sys.actuator.ctrl_range[:, 1]
        action = (action + 1) * (action_max - action_min) * 0.5 + action_min
        pipeline_state = self.pipeline_step(state.pipeline_state, action)
        obs = self._get_obs(pipeline_state)
        done = jnp.zeros_like(state.done)
        return state.replace(
            pipeline_state=pipeline_state,
            obs=obs,
            reward=self._reward(pipeline_state, action),
            done=done,
        )

    def cart_position(self, pipeline_state: base.State) -> jax.Array:
        return pipeline_state.q[0]

    def pole_angle_components(self, pipeline_state: base.State) -> jax.Array:
        return jnp.cos(pipeline_state.q[1]), jnp.sin(pipeline_state.q[1])

    def bounded_position(self, pipeline_state: base.State) -> jax.Array:
        return jnp.hstack(
            (
                self.cart_position(pipeline_state),
                *self.pole_angle_components(pipeline_state),
            )
        )

    def _reward(self, pipeline_state: base.State, action: jax.Array) -> jax.Array:
        if self.sparse:
            cart_in_bounds = 1.0
            angle_in_bounds = 1.0
            return cart_in_bounds * angle_in_bounds
        else:
            upright = (self.pole_angle_components(pipeline_state)[0] + 1) / 2
            centered = 1.0
            centered = (1 + centered) / 2
            small_control = 1.0
            small_control = (4 + small_control) / 5
            small_velocity = 1.0
            small_velocity = (1 + small_velocity) / 2
            return upright.mean() * small_control * small_velocity * centered

    @property
    def action_size(self):
        return 1

    def _get_obs(self, pipeline_state: base.State) -> jax.Array:
        """Observe cartpole body position and velocities."""
        return jnp.concatenate(
            [self.bounded_position(pipeline_state), pipeline_state.qd]
        )


register_environment(
    "cartpole_swingup_sparse",
    lambda **kwargs: Cartpole(swingup=True, sparse=True, **kwargs),
)
register_environment(
    "cartpole_swingup", lambda **kwargs: Cartpole(swingup=True, **kwargs)
)
register_environment(
    "cartpole_balance", lambda **kwargs: Cartpole(swingup=False, **kwargs)
)


def test_parameterization():
    def policy(*_, **__):
        return jnp.ones((1,)), None

    environment = envs.get_environment(env_name="cartpole_balance")
    rng = jax.random.PRNGKey(1)
    rng = jax.random.split(rng, _PARALLEL_ENVS)
    new_sys, in_axes, samples = domain_randomization(environment.sys, rng)
    environment = envs.training.wrap(
        environment,
        action_repeat=1,
        randomization_fn=lambda *_, **__: (new_sys, in_axes),
    )
    policy = jax.vmap(policy, in_axes=(0, None))
    state = environment.reset(jnp.asarray([jax.random.PRNGKey(0)] * _PARALLEL_ENVS))
    next_state = environment.step(state, policy(state, None)[0])
    count = sum(
        jnp.allclose(next_state.obs[0, :], obs[0, :])
        for obs in jnp.split(next_state.obs[1:, :], _PARALLEL_ENVS - 1, axis=0)
    )
    assert (
        count / _PARALLEL_ENVS < 0.2
    ), "Different environment initializations should have different trajectories"


if __name__ == "__main__":
    test_parameterization()

As seen, the environment is based on deepmind control suite cartpole. Now, for some reason, when I use the mjx backend, in line 64 (Cartpole's constructor), I get the exact same next state (while I would expect different states, as the dynamics are perturbed), while if I use generalized backend, the test passes.

Any idea why this may happen? The barkour tutorial suggests that domain randomization with the mjx backend should work.

Any help would be very much appreciated!

EDIT: tried now running a similar experiment with the InvertedPendulum environment -- seems like I get the same behavior. Happy to create a standalone running code for it as well.

btaba · 2024-10-04T16:51:33Z

btaba
Oct 4, 2024
Maintainer

Hi @yardenas , actuator.gear is an attribute only used by generalized, spring and positional but not by mjx. Please follow the example exactly as shown in the tutorial! Sorry for the late reply!

0 replies

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Domain randomization with mjx backend #523

{{title}}

{{editor}}'s edit

{{editor}}'s edit

Replies: 1 comment

{{title}}

Select a reply

Domain randomization with mjx backend #523

yardenas Jul 18, 2024

Replies: 1 comment

btaba Oct 4, 2024 Maintainer

yardenas
Jul 18, 2024

btaba
Oct 4, 2024
Maintainer