mj_forward modifies jacobian when called after mj_step

Open wualbert opened this issue 1 year ago • 0 comments

Summary

Calling mj_forward after a mj_step call unexpectedly modifies the Jacobian returned by mj_jacBody. By extension, other derived quantities may be affected as well. The generalized coordinates qpos, qvel are not modified by the call, which is expected.

Expected behavior

According to the documentation, mj_forward performs the same computations as mj_step but without the integration. The Jacobian, which only depends on qpos, should not change after calling mj_forward since qpos is the same and the Jacobian should have already been computed by mj_step.

Code example

The following examples showcases how mj_forward changes the Jacobian. The code was tested on a M1 mac with macOS 14.5, Python 3.10.14, Mujoco 3.1.5.

# Adapted from https://github.com/tayalmanan28/MuJoCo-Tutorial
import mujoco as mj
import numpy as np
from mujoco.glfw import glfw
import sys

xml = """
<mujoco>
	<option gravity="0 0 -9.81">
	</option>

	<worldbody>
		<body pos="0 0 2" euler="0 180 0">
			<joint name="pin" type="hinge" axis="0 -1 0" pos="0 0 0.5"/>
			<geom type="cylinder" size=".05 .5" rgba="0 .9 0 1" mass="1"/>
		</body>
	</worldbody>
	
	<actuator>
		<motor joint="pin" name="torque" gear="1" ctrllimited="true" ctrlrange="-100 100"/>
		<position name="position_servo" joint="pin" kp="10"/>
		<velocity name="velocity_servo" joint="pin" kv="0"/>
	</actuator>
</mujoco>
"""


def get_jac(model, data):
    jacp = np.zeros((3, model.nv))
    mj.mj_jacBody(model, data, jacp, None, 1)
    return jacp


def main(test_jacobian):
    model = mj.MjModel.from_xml_string(xml)
    data = mj.MjData(model)  # MuJoCo data
    cam = mj.MjvCamera()  # Abstract camera
    opt = mj.MjvOption()
    data.qpos[0] = np.pi / 2
    # Set camera configuration
    cam.azimuth = 90.0
    cam.distance = 5.0
    cam.lookat = np.array([0.012768, -0.000000, 1.254336])
    # Init GLFW, create window, make OpenGL context current, request v-sync
    glfw.init()
    window = glfw.create_window(1200, 900, "Demo", None, None)
    glfw.make_context_current(window)
    glfw.swap_interval(1)

    scene = mj.MjvScene(model, maxgeom=10000)
    context = mj.MjrContext(model, mj.mjtFontScale.mjFONTSCALE_150.value)
    while not glfw.window_should_close(window):
        mj.mj_step(model, data)
        # Compute the jacobian
        jac0 = get_jac(model, data)
        q0 = np.copy(data.qpos)
        v0 = np.copy(data.qvel)
        # Call mj_forward
        mj.mj_forward(model, data)
        if test_jacobian:
            # This test fails
            np.testing.assert_allclose(jac0, get_jac(model, data), atol=1e-4)
        np.testing.assert_allclose(q0, data.qpos, atol=1e-7)
        np.testing.assert_allclose(v0, data.qvel, atol=1e-7)
        viewport_width, viewport_height = glfw.get_framebuffer_size(window)
        viewport = mj.MjrRect(0, 0, viewport_width, viewport_height)
        # Update scene and render
        mj.mjv_updateScene(
            model,
            data,
            opt,
            None,
            cam,
            mj.mjtCatBit.mjCAT_ALL.value,
            scene,
        )
        mj.mjr_render(viewport, scene, context)
        glfw.swap_buffers(window)
        glfw.poll_events()

    glfw.terminate()


if __name__ == "__main__":
    test_jacobian = "--test-jacobian" in sys.argv
    if test_jacobian:
        print("Jacobian testing is enabled.")
    else:
        print("Jacobian testing is not enabled.")
    main(test_jacobian)

May 15 '24 23:05 wualbert