Inconsistent parameter gradient behaviour across pipelines #555

JoeMWatson · 2024-10-22T16:22:20Z

JoeMWatson
Oct 22, 2024

Hi,

I'm using brax 0.11.0 and am confused about parameter gradients. In this issue from last year, pipelines appeared to be consistent, however I'm looking at parameter gradients w.r.t. link mass and they no longer appear consistent across pipelines

If you run the following script,

import jax
from jax import numpy as jp
from brax.positional import pipeline as positional_pipeline
from brax.generalized import pipeline as generalized_pipeline
from brax.mjx import pipeline as mjx_pipeline
from brax.io import mjcf
ball = """
<mujoco>
  <option gravity="0 0 -9.81" timestep="0.002"/>
  <worldbody>
    <geom name="floor" pos="0 0 0" size="40 40 40" type="plane" friction="1.0"/>
    <body pos="0 0 0.2">
      <joint type="free"/>
      <geom size=".2" mass="1.0" friction="1.0"/>
    </body>
  </worldbody>
</mujoco>
"""
for name, pipeline in {
    "pos": positional_pipeline,
    "gen": generalized_pipeline,
    "mjx": mjx_pipeline,
}.items():
    print(name)
    try:
        def simulation(pipeline, sys, params):
            init_qd = jp.zeros(6).at[0].set(1.0)  # 1m/s in +x
            mass = params
            sys = sys.replace(
                link=sys.link.replace(inertia=sys.link.inertia.replace(mass=mass))
            )
            state = jax.jit(pipeline.init)(sys, sys.init_q, init_qd)
            for i in range(1):
                state = jax.jit(pipeline.step)(sys, state, None)
            return state.qd[0]

        sys = mjcf.loads(ball)
        mass = sys.link.inertia.mass
        params = jp.ones_like(mass)
        x = simulation(pipeline, sys, params)
        grad_x = jax.grad(simulation, argnums=2)(pipeline, sys, params)
        print(f"{x=}, {grad_x=}")
    except Exception as exc:
        print(exc)

the output is

pos
x=Array(0.99439585, dtype=float32), grad_x=Array([-0.00160096], dtype=float32)
gen
reshape total size must be unchanged, got new_sizes (6, 3) (of total size 18) for shape (6,) (of total size 6).
mjx
x=Array(1., dtype=float32), grad_x=Array([0.], dtype=float32)

Should we expect to be able to compute parameter gradients for generalized and MJX pipelines?

btaba · 2024-11-16T07:30:30Z

btaba
Nov 16, 2024
Maintainer

Hi @JoeMWatson , MJX doesn't use sys.link.inertia, which is presumably why the grad is 0. MJX uses classic MuJoCo fields, so mjx.Model.body_inertia. However, keep in mind that if mass/inertia is changed on the fly, you'll want to recompile the entire model, since there will be other derived fields in the model.

0 replies

JoeMWatson · 2024-11-19T14:35:24Z

JoeMWatson
Nov 19, 2024
Author

Thanks! Would you be able to provide a minimal working example of this functionality? I updated the script with your suggestion

import traceback
import jax
from jax import numpy as jp
from brax.mjx import pipeline as mjx_pipeline
from brax.io import mjcf


ball = """
<mujoco>
  <option gravity="0 0 -9.81" timestep="0.002" solver="Newton" iterations="1" ls_iterations="4"/>
  <worldbody>
    <geom name="floor" pos="0 0 0" size="40 40 40" type="plane" friction="1.0"/>
    <body pos="0 0 0.2">
      <joint type="free"/>
      <geom size=".2" mass="1.0" friction="1.0"/>
    </body>
  </worldbody>
</mujoco>
"""
for name, pipeline in {
    "mjx": mjx_pipeline,
}.items():
    print(name)
    try:

        def simulation(pipeline, system, params, replace=True):
            init_qd = jp.zeros(6).at[0].set(1.0)  # 1m/s in +x
            if replace:
                system = system.replace(body_mass=params)
            state = jax.jit(pipeline.init)(system, system.init_q, init_qd)
            for i in range(1):
                state = jax.jit(pipeline.step)(system, state, None)
            return state.qd[0]

        sys = mjcf.loads(ball)
        params = sys.body_mass
        x = simulation(pipeline, sys, params, replace=False)
        grad_x = jax.grad(simulation, argnums=2)(pipeline, sys, params, replace=False)
        x = simulation(pipeline, sys, params, replace=True)
        grad_x = jax.grad(simulation, argnums=2)(pipeline, sys, params, replace=True)
        print(f"{x=}, {grad_x=}")
    except Exception as exc:
        print(traceback.format_exc())

The traceback is

Traceback (most recent call last):
  File "mjx_example.py", line 40, in <module>
    x = simulation(pipeline, sys, params, replace=True)
  File "mjx_example.py", line 32, in simulation
    state = jax.jit(pipeline.step)(system, state, None)
  File ".../brax/mjx/pipeline.py", line 104, in step
    data = mjx.step(sys, data)
  File ".../mujoco/mjx/_src/forward.py", line 57, in wrapper
    res = fn(*args, **kwargs)
  File ".../mujoco/mjx/_src/forward.py", line 414, in step
    d = forward(m, d)
  File ".../mujoco/mjx/_src/forward.py", line 57, in wrapper
    res = fn(*args, **kwargs)
  File ".../mujoco/mjx/_src/forward.py", line 406, in forward
    d = named_scope(solver.solve)(m, d)
  File ".../mujoco/mjx/_src/forward.py", line 57, in wrapper
    res = fn(*args, **kwargs)
  File ".../mujoco/mjx/_src/solver.py", line 558, in solve
    ctx = body(ctx)
  File ".../mujoco/mjx/_src/solver.py", line 533, in body
    ctx = _update_constraint(m, d, ctx)
  File ".../mujoco/mjx/_src/solver.py", line 275, in _update_constraint
    efc_force = d.efc_D * -ctx.Jaref * active + floss_force
  File ".../jax/_src/numpy/array_methods.py", line 736, in op
    return getattr(self.aval, f"_{name}")(self, *args)
  File "/opt/miniconda3/envs/sim2gp311/lib/python3.11/site-packages/jax/_src/numpy/array_methods.py", line 264, in deferring_binary_op
    return binary_op(*args)
jax._src.source_info_util.JaxStackTraceBeforeTransformation: TypeError: mul got incompatible shapes for broadcasting: (2,), (4,).

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "mjx_example.py", line 41, in <module>
    grad_x = jax.grad(simulation, argnums=2)(pipeline, sys, params, replace=True)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: mul got incompatible shapes for broadcasting: (2,), (4,).
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

It appears the gradient computation introduces a broadcasting error, but it is not obvious to me why this happens looking at the stack trace.

Could you also be specific about what you mean by "recompile the entire model"? Do you mean that pipeline.init must be run or does the System pytree need to be updated somehow? I noticed that the reference mj_model in System contains numpy array fields, so it cannot be updated inside a jax function without callbacks.

1 reply

JoeMWatson Nov 20, 2024
Author

I'm guessing that it's not a good idea to use grad on a while loop, because it works if I replace jax.grad with jax.jacfwd!

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Inconsistent parameter gradient behaviour across pipelines #555

{{title}}

Replies: 2 comments 1 reply

{{title}}

{{title}}

{{editor}}'s edit

{{editor}}'s edit

{{title}}

{{editor}}'s edit

{{editor}}'s edit

Select a reply

Inconsistent parameter gradient behaviour across pipelines #555

JoeMWatson Oct 22, 2024

Replies: 2 comments · 1 reply

btaba Nov 16, 2024 Maintainer

JoeMWatson Nov 19, 2024 Author

JoeMWatson Nov 20, 2024 Author

JoeMWatson
Oct 22, 2024

Replies: 2 comments 1 reply

btaba
Nov 16, 2024
Maintainer

JoeMWatson
Nov 19, 2024
Author

JoeMWatson Nov 20, 2024
Author