Question about use of brax in tutorials

[shubhankarKulk]

Intro

Hi!

I am a Masters Student using MuJoCo for research purposes

My setup

MuJoCo 3.2.7, Python, Ubuntu 20.04

My question

I had a look at the example https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb and was wondering if brax has anything to do with the motion of the humanoid. I have tried different implementations of RL algorithms using torch using the gymnasium environments, and during locomotion of the humanoid, it seems to not move one foot forward when considering torch based algorithms. It usually drags the other foot and does not seem to flex the knees. But when brax based algorithm is used in this tutorial, it bolts and flexes knees. It learns the locomotion pretty quickly. Is this usually the case, or might there be an implementation error for me?

Minimal model and/or code that explain my question

No response

Confirmations

• I searched the latest documentation thoroughly before posting.
• I searched previous Issues and Discussions, I am certain this has not been raised before.

[Balint-H]

Hello! Can you confirm you use the same observation and reward spaces, as well as the hyperparameters for the learning?

[shubhankarKulk]

Is there a certain way I can circumvent this whilst writing code in PyTorch?

[Balint-H]

Check out the make_ppo_networks function in brax's ppo/networks.py. There you can see what activations, action distributions, etc are used. Also consider that in the brax examples the observations are normalized, which can have a huge effect.

As far as reconciling differences: You can supply your own network_factory to ppo.train, or you can try equivalent features in PyTorch to implement the same networks.

[shubhankarKulk]

I will definitely look onto it. What about creating an environment from scratch like the one in mujoco jax tutorial. Is pipelining an essential step for the program to work, or can it be written without pipelining?

[Balint-H]

If you have not seen it already, check out https://playground.mujoco.org/ . They have implemented many MJX environments without explicitly inheriting from brax's PipelineEnv. (Although, in most examples they end up wrapping the framework agnostic environments and using brax for training). Let me know if I misunderstood what you meant by pipelining!

[shubhankarKulk]

By pipelining I meant this line of code in the tutorials:

def step(self, state: State, action: jp.ndarray) -> State:
    """Runs one timestep of the environment's dynamics."""
    data0 = state.pipeline_state
    data = self.pipeline_step(data0, action)

    com_before = data0.subtree_com[1]
    com_after = data.subtree_com[1]
    velocity = (com_after - com_before) / self.dt
    forward_reward = self._forward_reward_weight * velocity[0]

    min_z, max_z = self._healthy_z_range
    is_healthy = jp.where(data.q[2] < min_z, 0.0, 1.0)
    is_healthy = jp.where(data.q[2] > max_z, 0.0, is_healthy)
    if self._terminate_when_unhealthy:
      healthy_reward = self._healthy_reward
    else:
      healthy_reward = self._healthy_reward * is_healthy

    ctrl_cost = self._ctrl_cost_weight * jp.sum(jp.square(action))

    obs = self._get_obs(data, action)
    reward = forward_reward + healthy_reward - ctrl_cost
    done = 1.0 - is_healthy if self._terminate_when_unhealthy else 0.0
    state.metrics.update(
        forward_reward=forward_reward,
        reward_linvel=forward_reward,
        reward_quadctrl=-ctrl_cost,
        reward_alive=healthy_reward,
        x_position=com_after[0],
        y_position=com_after[1],
        distance_from_origin=jp.linalg.norm(com_after),
        x_velocity=velocity[0],
        y_velocity=velocity[1],
    )

    return state.replace(
      
 pipeline_state=data, obs=obs, reward=reward, done=done
    )

I did not understand the clear use of pipeline here and if this can be implemented without the use of brax.