Energy calculations in MJX #2314
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
|
@simeon-ned what a nice PR! Thank you for creating it. Things are still slow due to the holidays, but someone will get to this soon. |
|
Hello @erikfrey, is there anything I can do to help with this PR? I could add tests similar to those implemented in mjx src. |
| raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}') | ||
| # Check enable flags using enum pattern | ||
| if types.EnableBit(o.enableflags) not in set(types.EnableBit) and o.enableflags != 0: | ||
| raise NotImplementedError(f'{mujoco.mjtEnableBit(o.enableflags)}') |
There was a problem hiding this comment.
I don't think mujoco.mjtEnableBit(o.enableflags) works when you have multiple bits set, you will get an error like this:
ValueError: Invalid int value for mjtEnableBit: 3
mjx/mujoco/mjx/_src/sensor.py
Outdated
|
|
||
| # Add joint spring potential energy | ||
| if not m.opt.disableflags & DisableBit.PASSIVE: | ||
| for i in range(m.njnt): |
There was a problem hiding this comment.
I challenge you to figure out how to do this without this for loop iterating over joints. You should find code elsewhere that does this: look for usage of scan.flat in smooth.py for some examples.
scan is a bit gnarly and not well super well documented, but should be used wherever possible.
mjx/mujoco/mjx/_src/sensor.py
Outdated
|
|
||
| # Add tendon spring potential energy | ||
| if not m.opt.disableflags & DisableBit.PASSIVE: | ||
| for i in range(m.ntendon): |
There was a problem hiding this comment.
this for loop can be removed, and you should not need scan at all.
you can do something like:
lower = m.tendon_lengthspring[0::2]
upper = m.tendon_lengthspring[1::2]
length = d.ten_length
displacement = jp.where(length > upper, upper - length, 0)
displacement = jp.where(length < lower, lower - length, displacement)
...
mjx/mujoco/mjx/_src/sensor.py
Outdated
| """ | ||
|
|
||
| # Initialize potential energy | ||
| energy = jp.array(0.0) |
There was a problem hiding this comment.
mega nitpick: please move this variable definition to right after the if statement guard
|
|
||
| # Add tendon spring potential energy using vectorized operations | ||
| if not m.opt.disableflags & DisableBit.PASSIVE: | ||
| if not m.opt.disableflags & DisableBit.PASSIVE & m.tendon_lengthspring.size > 0: |
There was a problem hiding this comment.
Order of logical operations is probably incorrect. This evaluates to true even if m.tendon_lengthspring.size > 0 is false, which leads to access to non-existent elements of this array.
(not m.opt.disableflags & DisableBit.PASSIVE) & (m.tendon_lengthspring.size > 0) should work.
|
|
||
| elif jnt_type in (JointType.FREE, JointType.BALL): | ||
| # Convert quaternion difference to angular displacement | ||
| quat = qpos[padr:padr+4] |
There was a problem hiding this comment.
Probably slicing syntax in this form is not supported, as jax requires static start/stop/step to be used with NumPy indexing syntax.
IndexError: Array slice indices must have static start/stop/step to be used with NumPy indexing syntax. Found slice(Traced<ShapedArray(int32[1])>with<BatchTrace> with
val = Traced<ShapedArray(int32[1,1])>with<DynamicJaxprTrace>
batch_dim = 0, Traced<ShapedArray(int32[1])>with<BatchTrace> with
val = Traced<ShapedArray(int32[1,1])>with<DynamicJaxprTrace>
batch_dim = 0, None). To index a statically sized array at a dynamic position, try lax.dynamic_slice/dynamic_update_slice (JAX does not support dynamically sized arrays within JIT compiled functions).
This happens when I try to jit compile functions which take derivatives of energies.
This PR adds support for kinetic and potential energy calculations in MJX, matching MuJoCo's implementation.
Key changes: