Getting a custom URDF (A1 robot) loaded in brax

[Velythyl]

Hey all,

I have been trying to load a few different URDFs in Brax on my own time. Brax seems to load them fine, but then the simulation seems to explode when I actually run it. No errors, but the renderer visibly shows model explosions.

Now, I have a research project involving the A1 robot. And I know there must be a way to load it, as shown on Brax's gif previews.

But I can't find how.

I did follow the steps shown here: https://github.com/google/brax/discussions/325

But I ran into some more issues (the cylinders used in the model could not be loaded, so I changed them to capsules).

Finally, I managed to load the model, but again there's some clear model explosion.

spring: https://github.com/google/brax/assets/31974070/04442729-b9e8-4f22-aa34-0e40c1727d61

positional: https://github.com/google/brax/assets/31974070/98cf7888-faf1-4435-a04e-d7ba406a5eef

Generalized seems to be doing better, but it still acts weird: https://github.com/google/brax/assets/31974070/6ba01e8c-bf58-4b14-b867-537753349076

Any help in getting brax to load and simulate the A1 correctly would be greatly appreciated.

[Velythyl]

Here's another example of the issue in gif form

[Velythyl]

Progress. Using the generalized API, and the Go1 model found here https://github.com/deepmind/mujoco_menagerie/blob/main/unitree_go1/go1.xml

I replaced all cylinders with capsules, removed the geom_priority tag from the feet, and I removed self-collisions and visual meshes.

[btaba]

Hi @Velythyl ! Apologies for the late reply. I copied the XML in https://github.com/google/brax/discussions/325 and loaded it up, removed cylinders, and the visual meshes for rendering speed: https://gist.github.com/btaba/916ef532cf958f3420c082aa7cabc3ed

Then I ran:

from brax.generalized import pipeline
from brax.io import mjcf

m = mjcf.loads(config, asset_path='mujoco_menagerie/unitree_a1')

jit_env_reset = jax.jit(pipeline.init)
jit_env_step = jax.jit(pipeline.step)

rollout = []
rng = jax.random.PRNGKey(seed=1)
state = jit_env_reset(m, m.init_q, jp.zeros(m.qd_size()))
for i in range(500):
  rollout.append(state)
  act = 10 * jp.sin(i / 100) * jp.ones(m.act_size())
  state = jit_env_step(m, state, act)

HTML(html.render(m, rollout))

And it seems to work just fine.

For positional, a few more parameters would need tuning. Take a look at examples in envs, such as humanoid https://github.com/google/brax/blob/main/brax/envs/assets/humanoid.xml#L10-L23 Positional parameters that are relevant to tune are joint_scale_pos/ang, spring_mass/inertia_scale (large/small inertias are typically problematic), constraint_ang/vel_damping. Let us know if you have any further issues

[btaba]

These custom parameters seem to work ok for positional:

  <custom>
      <numeric data="10" name="constraint_ang_damping"/>
      <numeric data="0.5" name="joint_scale_pos"/>
      <numeric data="0.1" name="joint_scale_ang"/>
      <numeric data="0.9" name="spring_inertia_scale"/>
  </custom>