Performance differences between v1 and v2 #396

benjamin-macadam · 2023-09-22T17:33:04Z

benjamin-macadam
Sep 22, 2023

Hi all,

I'm a big fan of the Brax project and I have been trying to use it for a small project with a masters student, essentially playing around with https://diffmimic.github.io/. We've noticed some pretty big discrepancies between the behavior in v1 and v2 (using the 'positional' backend). We found Brax v1 had pretty reasonable 'rag doll' behavior when you dropped a humanoid into a scene and didn't actuate the joints - in version 2 we've noticed some fairly significant collisions during the initial drop and the character never really comes to rest.

To sum this up as a proper question:

Is this possibly due to the settings for the humanoid environment (which I've attracted at the bottom).
Are some of these changes meant to improve gradient calculations, and if so, are there any demonstrations that you can share?

I have found a few very exciting papers in character animation and robotics using Brax at the tail-end of v1. The move to v2 seems like an upgrade, but without any clear guide on how to reproduce v1 behavior in v2 and documentation on how to port code, it will become very difficult for researchers to replicate/build off of these results.

Falling in v1:

2023-09-20.11-25-34.mp4

Falling in v2:

2023-09-20.11-21-48.mp4

To reproduce the script:

import jax
import jax.dlpack
import jax.numpy as jp
import numpy as np

import streamlit.components.v1 as components
import streamlit as st
st.set_page_config(layout='wide')

import brax
from brax import envs
from brax import base
from brax.envs.base import PipelineEnv, State
from brax.io import html
from brax.io import mjcf
import brax.generalized.dynamics


class HumanoidEnv(PipelineEnv):
    def __init__(self):
        path = 'assets/amp_humanoid.xml'
        sys = mjcf.load(path)
        super().__init__(sys=sys, backend='positional', n_frames=1)
        return

    def reset(self, q: jp.ndarray, qd: jp.ndarray) -> State:
        """Resets the environment to an initial state."""
        pipeline_state = self.pipeline_init(q, qd)
        return pipeline_state

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

def main():
    env = HumanoidEnv()

    jit_env_reset = jax.jit(env.reset)
    jit_env_step = jax.jit(env.step)

    q = env.sys.init_q
    qd = jp.zeros((env.sys.qd_size(),))

    rollout = []
    state = jit_env_reset(q, qd)
    rollout.append(state)
    for i in range(2000):
        if i % 100 == 0:
            print(f'step: {i}')
        state = jit_env_step(state, np.zeros(env.action_size, dtype=float))
        rollout.append(state)

    components.html(html.render(env.sys, rollout, height=1080), height=800)
    return

if __name__ == "__main__":
    main()

  <body name="torso" pos="0 0 0.236151">
    <light name="top" pos="0 0 2" mode="trackcom"/>
    <camera name="back" pos="-3 0 1" xyaxes="0 -1 0 1 0 2" mode="trackcom"/>
    <camera name="side" pos="0 -3 1" xyaxes="1 0 0 0 1 2" mode="trackcom"/>

    <inertial pos="0 0 0" diaginertia="1 1 1" mass="12.002052"/>
    <joint name="abdomen_x" pos="0 0 0" axis="1 0 0" range="-60 60" stiffness="1000" damping="100" armature=".02"/>
    <joint name="abdomen_y" pos="0 0 0" axis="0 1 0" range="-60 90" stiffness="1000" damping="100" armature=".02"/>
    <joint name="abdomen_z" pos="0 0 0" axis="0 0 1" range="-50 50" stiffness="1000" damping="100" armature=".02"/>
    <geom name="torso" type="sphere" pos="0 0 0.12" size="0.11" density="1794"/>

    <geom name="right_clavicle" fromto="-0.0060125 -0.0457775 0.2287955 -0.016835 -0.128177 0.2376182" size=".045" density="1100"/>
    <geom name="left_clavicle" fromto="-0.0060125 0.0457775 0.2287955 -0.016835 0.128177 0.2376182" size=".045" density="1100"/>

    <body name="head" pos="0 0 0.223894">
      <inertial pos="0 0 0" diaginertia="1 1 1" mass="3.8822644"/>
      <joint name="neck_x" axis="1 0 0" range="-50 50" stiffness="100" damping="10" armature=".01"/>
      <joint name="neck_y" axis="0 1 0" range="-40 60" stiffness="100" damping="10" armature=".01"/>
      <joint name="neck_z" axis="0 0 1" range="-45 45" stiffness="100" damping="10" armature=".01"/>
      <geom name="head" type="sphere" pos="0 0 0.175" size="0.095" density="1081"/>
      <camera name="egocentric" pos=".103 0 0.175" xyaxes="0 -1 0 .1 0 1" fovy="80"/>
    </body>

    <body name="right_upper_arm" pos="-0.02405 -0.18311 0.24350">
      <inertial pos="0 0 0" diaginertia="1 1 1" mass="1.4993314"/>
      <joint name="right_shoulder_x" axis="1 0 0" range="-180 45" stiffness="400" damping="40" armature=".02"/>
      <joint name="right_shoulder_y" axis="0 1 0" range="-180 60" stiffness="400" damping="40" armature=".02"/>
      <joint name="right_shoulder_z" axis="0 0 1" range="-90 90" stiffness="400" damping="40" armature=".02"/>
      <geom name="right_upper_arm" fromto="0 0 -0.05 0 0 -0.23" size=".045" density="982"/>

      <body name="right_lower_arm" pos="0 0 -0.274788">
        <inertial pos="0 0 0" diaginertia="1 1 1" mass="1.5"/>
        <joint name="right_elbow" axis="0 1 0" range="-160 0" stiffness="300" damping="30" armature=".01"/>
        <geom name="right_lower_arm" fromto="0 0 -0.0525 0 0 -0.1875" size="0.04" density="1056"/>

        <body name="right_hand" pos="0 0 -0.258947">
          <geom name="right_hand" type="sphere" size=".04" density="1865"/>
        </body>
      </body>
    </body>

    <body name="left_upper_arm" pos="-0.02405 0.18311 0.24350">
      <inertial pos="0 0 0" diaginertia="1 1 1" mass="1.4993314"/>
      <joint name="left_shoulder_x" axis="1 0 0" range="-45 180" stiffness="400" damping="40" armature=".02"/>
      <joint name="left_shoulder_y" axis="0 1 0" range="-180 60" stiffness="400" damping="40" armature=".02"/>
      <joint name="left_shoulder_z" axis="0 0 1" range="-90 90" stiffness="400" damping="40" armature=".02"/>
      <geom name="left_upper_arm" fromto="0 0 -0.05 0 0 -0.23" size="0.045" density="982"/>

      <body name="left_lower_arm" pos="0 0 -0.274788">
        <inertial pos="0 0 0" diaginertia="1 1 1" mass="1.5"/>
        <joint name="left_elbow" axis="0 1 0" range="-160 0" stiffness="300" damping="30" armature=".01"/>
        <geom name="left_lower_arm" fromto="0 0 -0.0525 0 0 -0.1875" size="0.04" density="1056"/>

        <body name="left_hand" pos="0 0 -0.258947">
          <geom name="left_hand" type="sphere" size=".04" density="1865"/>
        </body>
      </body>
    </body>
  </body>

  <body name="right_thigh" pos="0 -0.084887 0">
    <inertial pos="0 0 0" diaginertia="1 1 1" mass="4.5022917"/>
    <joint name="right_hip_x" axis="1 0 0" range="-60 15" stiffness="500" damping="50" armature=".02"/>
    <joint name="right_hip_y" axis="0 1 0" range="-140 60" stiffness="500" damping="50" armature=".02"/>
    <joint name="right_hip_z" axis="0 0 1" range="-60 35" stiffness="500" damping="50" armature=".02"/>
    <geom name="right_thigh" fromto="0 0 -0.06 0 0 -0.36" size="0.055" density="1269"/>

    <body name="right_shin" pos="0 0 -0.421546">
      <inertial pos="0 0 0" diaginertia="1 1 1" mass="2.9997497"/>
      <joint name="right_knee" pos="0 0 0" axis="0 1 0" range="0 160" stiffness="500" damping="50" armature=".02"/>
      <geom name="right_shin" fromto="0 0 -0.045 0 0 -0.355" size=".05" density="1014"/>

      <body name="right_foot" pos="0 0 -0.409870">
        <inertial pos="0 0 0" diaginertia="1 1 1" mass="0.99968714"/>
        <joint name="right_ankle_x" pos="0 0 0" axis="1 0 0" range="-30 30" stiffness="400" damping="40" armature=".01"/>
        <joint name="right_ankle_y" pos="0 0 0" axis="0 1 0" range="-55 55" stiffness="400" damping="40" armature=".01"/>
        <joint name="right_ankle_z" pos="0 0 0" axis="0 0 1" range="-40 40" stiffness="400" damping="40" armature=".01"/>
        <geom name="right_foot" type="box" pos="0.045 0 -0.0225" size="0.0885 0.045 0.0275" density="1141"/>
      </body>
    </body>
  </body>

  <body name="left_thigh" pos="0 0.084887 0">
    <inertial pos="0 0 0" diaginertia="1 1 1" mass="4.5022917"/>
    <joint name="left_hip_x" axis="1 0 0" range="-15 60" stiffness="500" damping="50" armature=".02"/>
    <joint name="left_hip_y" axis="0 1 0" range="-140 60" stiffness="500" damping="50" armature=".02"/>
    <joint name="left_hip_z" axis="0 0 1" range="-35 60" stiffness="500" damping="50" armature=".02"/>
    <geom name="left_thigh" fromto="0 0 -0.06 0 0 -0.36" size=".055" density="1269"/>

    <body name="left_shin" pos="0 0 -0.421546">
      <inertial pos="0 0 0" diaginertia="1 1 1" mass="2.9997497"/>
      <joint name="left_knee" pos="0 0 0" axis="0 1 0" range="0 160" stiffness="500" damping="50" armature=".02"/>
      <geom name="left_shin" fromto="0 0 -0.045 0 0 -0.355" size=".05" density="1014"/>

      <body name="left_foot" pos="0 0 -0.409870">
        <inertial pos="0 0 0" diaginertia="1 1 1" mass="0.99968714"/>
        <joint name="left_ankle_x" pos="0 0 0" axis="1 0 0" range="-30 30" stiffness="400" damping="40" armature=".01"/>
        <joint name="left_ankle_y" pos="0 0 0" axis="0 1 0" range="-55 55" stiffness="400" damping="40" armature=".01"/>
        <joint name="left_ankle_z" pos="0 0 0" axis="0 0 1" range="-40 40" stiffness="400" damping="40" armature=".01"/>
        <geom name="left_foot" type="box" pos="0.045 0 -0.0225" size="0.0885 0.045 0.0275" density="1141"/>
      </body>
    </body>
  </body>
</body>

```

btaba · 2023-10-15T22:40:57Z

btaba
Oct 15, 2023
Maintainer

Hi @benjamin-macadam , In v2 we're using inertias computed from the MJCF loader, but these are pretty unstable with the positional backend without significant tuning. I expect this is likely the issue you're seeing.

We expose a parameter to "sphericalize" the inertias here:

brax/brax/envs/assets/humanoid.xml

Line 20 in 45729a4

which was needed in humanoid. Can you inspect the masses/inertias and see if this fixes your issue?

3 replies

benjamin-macadam Dec 5, 2023
Author

Hi, sorry to bring this back up, but the issue actually persists with the change.

Is there any way to load systems that were defined using the older brax format from v1? That would solve most of my problems if it is possible.

btaba Dec 5, 2023
Maintainer

Hi @benjamin-macadam , you can use brax.v1

btaba Dec 5, 2023
Maintainer

Are you not seeing stable behavior or is there another issue? Which physics backend are you trying to use and do you have an updated XML you want us to look at?

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Performance differences between v1 and v2 #396

{{title}}

Replies: 1 comment 3 replies

{{title}}

{{editor}}'s edit

{{editor}}'s edit

{{title}}

{{title}}

{{title}}

Select a reply

Performance differences between v1 and v2 #396

benjamin-macadam Sep 22, 2023

Replies: 1 comment · 3 replies

btaba Oct 15, 2023 Maintainer

benjamin-macadam Dec 5, 2023 Author

btaba Dec 5, 2023 Maintainer

btaba Dec 5, 2023 Maintainer

benjamin-macadam
Sep 22, 2023

Replies: 1 comment 3 replies

btaba
Oct 15, 2023
Maintainer

benjamin-macadam Dec 5, 2023
Author

btaba Dec 5, 2023
Maintainer

btaba Dec 5, 2023
Maintainer