Control Inputs Not Affecting Robot Movement in MuJoCo MJX Simulation

[wgu93]

Intro

Hi!

I am a researcher using MuJoCo for my research on robotics.

My setup

MuJoCo: 3.2.5
Python: 3.10
OS and arch: Apple M3 & Ubuntu 20.04 (x86 + Nvidia GPU)
Jax: 0.4.35

My question

I am currently using MJX for simulating a mobile differential drive robot. However, I have encountered an issue where setting control inputs via mjx_d.ctrl (specifically the wheel rotation) does not appear to propagate to the robot’s state, particularly its position (mjx_d.qpos[:3]).

I attempted to assign control inputs in two different ways, as shown in the attached code snippet. In both cases, I assigned a scalar control input of 1 (with control limits set to ±10). The results from the first attempt are shown below:

[0. 0. 0.1 1. 0. 0. 0. 0. 0. 0. 0. ] <class 'numpy.ndarray'> [0. 0. 0.1 1. 0. 0. 0. 0. 0. 0. 0. ] <class 'jaxlib.xla_extension.ArrayImpl'> {CpuDevice(id=0)} mjx_d.qpos: [-5.6603234e-08 8.0123148e-13 9.9960759e-02] mjx_d.ctrl: [0. 0. 1. 0.] mjx_d.qpos: [-1.6980951e-07 1.6024636e-12 9.9882282e-02] mjx_d.ctrl: [0. 0. 1. 0.] mjx_d.qpos: [-3.3961811e-07 3.1448915e-12 9.9764563e-02] mjx_d.ctrl: [0. 0. 1. 0.] mjx_d.qpos: [-5.6603244e-07 4.8413469e-12 9.9607602e-02] mjx_d.ctrl: [0. 0. 1. 0.] mjx_d.qpos: [-8.4904326e-07 6.4331097e-12 9.9411398e-02] mjx_d.ctrl: [0. 0. 1. 0.] mjx_d.qpos: [-1.1886342e-06 7.8106089e-12 9.9175960e-02] mjx_d.ctrl: [0. 0. 1. 0.] mjx_d.qpos: [-1.5848274e-06 1.0071862e-11 9.8901279e-02] mjx_d.ctrl: [0. 0. 1. 0.] mjx_d.qpos: [-2.0376197e-06 1.3492342e-11 9.8587357e-02] mjx_d.ctrl: [0. 0. 1. 0.] mjx_d.qpos: [-2.5470022e-06 1.6222758e-11 9.8234199e-02] mjx_d.ctrl: [0. 0. 1. 0.] mjx_d.qpos: [-3.1130867e-06 1.8820948e-11 9.7841799e-02] mjx_d.ctrl: [0. 0. 1. 0.]

As shown above, the control input is successfully assigned to mjx_d.ctrl. However, the robot’s position remains nearly constant, with only minimal movements near the origin observed. I repeated the process using the second method of assigning control inputs and observed similar behavior.

Additionally, I increased the control input from 1 to higher values (e.g., 10) and from below, it can be seen that the robot's position (mjx_d.qpos[:3]) changed very slightly, which suggests that the system is still not acting as expected. Moreover, further increasing the control input would lead to NaN values in the simulation.

[0. 0. 0.1 1. 0. 0. 0. 0. 0. 0. 0. ] <class 'numpy.ndarray'> [0. 0. 0.1 1. 0. 0. 0. 0. 0. 0. 0. ] <class 'jaxlib.xla_extension.ArrayImpl'> {CpuDevice(id=0)} mjx_d.qpos: [-5.660324e-07 8.012315e-12 9.996076e-02] mjx_d.ctrl: [ 0. 0. 10. 0.] mjx_d.qpos: [-1.6979701e-06 1.6958584e-11 9.9882282e-02] mjx_d.ctrl: [ 0. 0. 10. 0.] mjx_d.qpos: [-3.3957811e-06 2.6614141e-11 9.9764563e-02] mjx_d.ctrl: [ 0. 0. 10. 0.] mjx_d.qpos: [-5.6602307e-06 4.2671311e-11 9.9607602e-02] mjx_d.ctrl: [ 0. 0. 10. 0.] mjx_d.qpos: [-8.490393e-06 4.315451e-11 9.941140e-02] mjx_d.ctrl: [ 0. 0. 10. 0.] mjx_d.qpos: [-1.1885814e-05 1.1306803e-10 9.9175945e-02] mjx_d.ctrl: [ 0. 0. 10. 0.] mjx_d.qpos: [-1.5844797e-05 1.6720930e-10 9.8901242e-02] mjx_d.ctrl: [ 0. 0. 10. 0.] mjx_d.qpos: [-2.0369707e-05 1.6501520e-10 9.8587289e-02] mjx_d.ctrl: [ 0. 0. 10. 0.] mjx_d.qpos: [-2.5454458e-05 2.4332972e-10 9.8234080e-02] mjx_d.ctrl: [ 0. 0. 10. 0.] mjx_d.qpos: [-3.1100888e-05 3.5169670e-10 9.7841613e-02] mjx_d.ctrl: [ 0. 0. 10. 0.]

Could anyone please kindly provide insight into why the control inputs are not resulting in state propagation for the robot? Are there any settings, configurations, or additional parameters that I may have overlooked in the simulation code? I would greatly appreciate your comments on this issue. Thank you for your time and assistance!

Minimal model and/or code that explain my question

The MuJoCo model of a differential drive wheeled robot:

<mujoco model="differential_drive_robot">
    <!-- Assets -->
    <asset>
        <material name="wheel_material" rgba="0.2 0.2 0.2 1"/>
        <material name="body_material" rgba="0.8 0.1 0.1 1"/>
        <!-- For environment -->
        <texture builtin="gradient" height="100" rgb1="1 1 1" rgb2="0 0 0" type="skybox" width="100"/>
        <texture builtin="flat" height="1278" mark="cross" markrgb="1 1 1" name="texgeom" random="0.01" rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" type="cube" width="127"/>
        <texture builtin="checker" height="100" name="texplane" rgb1="0 0 0" rgb2="0.8 0.8 0.8" type="2d" width="100"/>
        <material name="MatPlane" reflectance="0.5" shininess="1" specular="1" texrepeat="60 60" texture="texplane"/>
        <material name="geom" texture="texgeom" texuniform="true"/>
    </asset>

    <!-- World -->
    <worldbody>
        <!-- Floor -->
        <body name="floor">
            <geom conaffinity="1" condim="3" material="MatPlane" name="floor" pos="0 0 0" rgba="0.8 0.9 0.8 1" size="40 40 40" type="plane"/>
        </body>

        <!-- Robot base -->
        <body name="robot_base" pos="0 0 0.1">
            <geom name="base_geom" type="box" size="0.2 0.1 0.05" material="body_material" mass="10"/>
            <joint name="base_joint" type="free"/>
            
            <!-- Left front wheel -->
            <body name="left_front_wheel" pos="0.15 0.12 -0.05">
                <geom name="lf_wheel_geom" type="cylinder" size="0.03 0.02" material="wheel_material" euler="90 0 0"/>
                <joint name="lf_wheel_joint" type="hinge" axis="0 1 0"/>
            </body>
            
            <!-- Left rear wheel -->
            <body name="left_rear_wheel" pos="-0.15 0.12 -0.05">
                <geom name="lr_wheel_geom" type="cylinder" size="0.03 0.02" material="wheel_material" euler="90 0 0"/>
                <joint name="lr_wheel_joint" type="hinge" axis="0 1 0"/>
            </body>
            
            <!-- Right front wheel -->
            <body name="right_front_wheel" pos="0.15 -0.12 -0.05">
                <geom name="rf_wheel_geom" type="cylinder" size="0.03 0.02" material="wheel_material" euler="90 0 0"/>
                <joint name="rf_wheel_joint" type="hinge" axis="0 1 0"/>
            </body>
            
            <!-- Right rear wheel -->
            <body name="right_rear_wheel" pos="-0.15 -0.12 -0.05">
                <geom name="rr_wheel_geom" type="cylinder" size="0.03 0.02" material="wheel_material" euler="90 0 0"/>
                <joint name="rr_wheel_joint" type="hinge" axis="0 1 0"/>
            </body>
        </body>
    </worldbody>

    <!-- Actuators -->
    <actuator>
        <!-- Left wheels -->
        <motor name="motor_left_front" joint="lf_wheel_joint" ctrlrange="-10 10" ctrllimited="true"/>
        <motor name="motor_left_rear" joint="lr_wheel_joint" ctrlrange="-10 10" ctrllimited="true"/>
        
        <!-- Right wheels -->
        <motor name="motor_right_front" joint="rf_wheel_joint" ctrlrange="-10 10" ctrllimited="true"/>
        <motor name="motor_right_rear" joint="rr_wheel_joint" ctrlrange="-10 10" ctrllimited="true"/>
    </actuator>
</mujoco>

The Python code that reproduces the results in the description (please save the above xml file and replace the placeholder xml_path with the actual directory containing the model below):

import mujoco
import mujoco.viewer
import os
import sys
import numpy as np
import jax
from mujoco import mjx


xml_path = " "

m = mujoco.MjModel.from_xml_path(xml_path)
d = mujoco.MjData(m)

mjx_m = mjx.put_model(m)
mjx_d = mjx.put_data(m, d)
print(d.qpos, type(d.qpos))
print(mjx_d.qpos, type(mjx_d.qpos), mjx_d.qpos.devices())

jit_step = jax.jit(mjx.step)
for i in range(10):
    # Attempt 1
    ctrl = mjx_d.ctrl.at[2].set(10)
    mjx_d = mjx_d.replace(ctrl=ctrl)

    # Attempt 2
    # d.ctrl = 1
    # mjx_d = mjx.put_data(m, d)

    mjx_d = jit_step(mjx_m, mjx_d)
    print(f"mjx_d.qpos: {mjx_d.qpos[:3]}")
    print(f"mjx_d.ctrl: {mjx_d.ctrl}")

    # For attempt 2 only
    # d = mjx.get_data(m, mjx_d)

Confirmations

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

[wgu93]

I just realized that the code is actually functioning correctly by extending the iterations, so I closed this discussion.