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
This PR adds support for kinetic and potential energy calculations in MJX, matching MuJoCo's implementation.
Key changes: