Error when trying to simulate multiple bodies

hackertyper commented 1 year ago

When trying to simulate a scene with multiple bodies, e.g. one of them a free falling sphere, I receive the error: "TypeError: take requires ndarray or scalar arguments, got <class 'NoneType'> at position 1."

The simulation works in MuJoCo. Below find a small snippet to reproduce the issue:

scene_xml = """
<mujoco>
    <option timestep="0.005"/>
    <worldbody>
        <body pos="0 0 3">
            <joint type="free"/>
            <geom size="0.5" type="sphere"/>
        </body>
        <geom size="40 40 40" type="plane"/>
        <body pos="1 0 0">
            <geom size="0.5" type="sphere"/>
        </body>
    </worldbody>
</mujoco>
"""

ball = mjcf.loads(scene_xml)
qd = jp.array([0, 0, 0, 0, 0, 0])
state = jax.jit(pipeline.init)(ball, ball.init_q, qd)
// Advance physics, error occurs here
jax.jit(pipeline.step)(ball, state, None)

btaba commented 10 months ago

Hi @hackertyper , I just ran the snippet in a colab, and can't reproduce the error. Will close the issue for now

hackertyper commented 9 months ago

Hi @btaba, I also cannot reproduce it in the current version anymore, so consider it fixed. Thanks!

google / brax

Error when trying to simulate multiple bodies #377