[mjx] subtree_linvel function computation

[Rian-Jo]

Intro

Hi!

Thanks for your great works.

I'm using mjx for RL.

My setup

mjx : 3.2.4
ubuntu : 22.04

My question

Hi!

Thanks for your great works.

I'm using mjx for RL.

In 'training_apg' tutorial, when i added in step function below,

...

  def step(self, state: State, action: jax.Array) -> State:
    action = jp.clip(action, -1, 1) # Raw action

    action = self.action_loc + (action * self.action_scale)

    data = self.pipeline_step(state.pipeline_state, action)

    mjx.subtree_vel(self.sys.mj_model, data)
    subtree_linvel = data.subtree_linvel[1, :]

...

i got the error below

in step
    mjx.subtree_vel(self.sys.mj_model, data)
  File "/home/u01/.pyenv/versions/3.10.14/lib/python3.10/site-packages/mujoco/mjx/_src/smooth.py", line 511, in subtree_vel
    subtree_angmom, * = scan.body*tree(
  File "/home/u01/.pyenv/versions/3.10.14/lib/python3.10/site-packages/mujoco/mjx/_src/scan.py", line 479, in body_tree
    key_y[key] = *nvmap(f, carry, fargs)
  File "/home/u01/.pyenv/versions/3.10.14/lib/python3.10/site-packages/mujoco/mjx/_src/scan.py", line 116, in _nvmap
    raise RuntimeError(f'numpy arg elements do not match: {arg}')
RuntimeError: numpy arg elements do not match: [2.915 2.889 2.895 2.895]

And subtree_linvel calculated in the first part of subtree_vel is zeroes.

1. Is above code usage correct?
2. Is there any other function executed befored 'mjx.subtree_vel'?

Thank you!

Minimal model and/or code that explain my question

No response

Confirmations

• I searched the latest documentation thoroughly before posting.
• I searched previous Issues and Discussions, I am certain this has not been raised before.