Model unable to generate correct thrust.

[aaprasad]

Intro

Hi!

I am a graduate student at MIT in the Brain & Cognitive Sciences under Dr. Steven Flavell. I use mujoco to develop neural controllers for C. elegans biomechanical models in order to study the neural computations underlying natural behaviors in C. elegans.

My setup

mujoco 3.7.0, mjx, python, mujoco-playground, ubuntu, nvidia 5090.

My question

TLDR

I'm trying to simulate a C.elegans biomechanical model moving on agar by learning a low-level controller via imitation learning. I'm having issues with learning a realistic policy and I think its coming from issues with the body model itself as well difficulties simulating the agar environment's friction

Problem

I've been collaborating with Dr. Talmo Pereira and his group to integrate a body model into their mimic-mjx platform using imitation learning to train low-level controllers for virtual neuroscience. Although the imitation learning seems to be working, I had been seeing this issue where it seems like the controller has learned to exploit the physics of the model by oscillating the head rapidly in order to generate thrust to propel the body model. Below I've shared a couple videos.

The first is just injecting random noise into the body model's actuators via the mujoco viewer. It seems like the body model is struggling to move much despite the high noise scale.

celegans_mjx_noise-2.mp4

The second is randomly sampling actions from a standard normal distribution in my environment. It seems strange that the head of the model is sort of perpendicular to the rest of the body.

celegans_mjx_random_action.mp4

The final one is an example of the issue with the imitation I've described above. The foreground is the learned neural controller attempting to imitate the expert trajectory which is displayed as the transparent body model in the background.

29_clip.109.mp4

What I've tried

After reading the discussions I saw this conversation with @yuvaltassa and realized that although my body model is in mm (ie the full length of the model adds up to 1 mm), the physical units (e.g gravity, viscosity and density) were in MKS units. I tried scaling the body model so that its in CGS units (the body model length now is 0.1 cm) which required me to switch to the RK4 integrator and decrease armature in order to stabilize the physics but the issue still persisted. I've explored a few parts of the body model that I think could be the issue:

1. Joint armature and dampening are incorrect/too low. I've tried playing around with these values but either the physics blows up or it doesn't seem to help
2. muscle parameters are incorrect. Interestingly, the original body model (see below) I used didn't actually have any of the muscle parameters configured so I just set the lengthrange such that the simulator would converge. However, when I've tried looking in the literature for feasible values, I've had difficulty translating them to mujoco in a way that makes sense as they've usually caused the model to not be able to move at all.
3. The friction is too high. I'm trying to simulate the worm on agar and from the literature it seems like the worm makes grooves into the agar and uses anisotropic friction to propel itself. Its not clear to me what those friction values should be because most of the literature describes it as a ratio.

I've also shared below the corresponding videos for the CGS version of the model.

celegans_cgs_noise.mp4

celegans_cgs_random_action.mp4

rollout_cgs.mp4

I'd like some feedback on my body model and how I should change it so that I can generate an undulatory sinusoidal forward/backward locomotor pattern with it. Please let me know if you need any clarification, I'm also happy to share any data necessary.

Minimal model and/or code that explain my question

Model:
I've taken the C. elegans biomechanical model from this repo. It's based on the $n$-link swimmer from gymnasium and contains 25 body components with corresponding meshes. The total length of the model is 1mm. I am trying to simulate a worm on agar so I'm using air as the medium. It is constrained to 2D movement in the xy plane with 24 joints hinge joints actuated by 95 muscles. Each joint is actuated by a set of 4 muscles organized into dorsal/ventral and left/right. Except for the last one which is missing its ventral left muscle. In order to get the model to work efficiently with mjx I removed mesh contact computation and used collision geoms instead. I also only kept contacts between the floor and the body (ie no inter body contact). I've shared three versions of the body model: celegans_og.xml is the original body model taken from the above repo, The other two are the xmls for the mjx_model I've dynamically created in mujoco-playground using mjspec. celegans_mjx.xml is the model I've modified for mjx, and celegans_cgs.xml is the mjx model scaled down to cgs units.

celegans_bodymodel.zip

Code:
I'm using mujoco-playground with mjspec to dynamically change the floor and other parameters. I instantiate the model as follows:

import mujoco
def add_worm(
    spec: mujoco.MjSpec
    walker_xml_path: str,
    torque_actuators: bool = False,
    rescale_factor: float = 1.0,
    dim: int = 2,
    pos: Tuple[float, float, float] = (0, 0, 0.05),
    quat: Tuple[float, float, float, float] = (1, 0, 0, 0),
    friction: Tuple[float, ...] = (1, 1, 0.005, 0.0001, 0.0001),
    solimp: Tuple[float, ...] = (0.9, 0.95, 0.001, 0.5, 2),
    rgba: Optional[Tuple[float, float, float, float]] = None,
    suffix: str = "-worm",
) -> None:
    """Adds the c. elegans model to the environment.

    Args:
        spec: current spec
        walker_xml_path: path to walker xml
        dim: dims to constrain motion to. 2 indicates planar movement in XY, 3 indicates unconstrained movement.
        torque_actuators: Whether to convert motors to torque-mode actuators.
        rescale_factor: Factor to rescale the c. elegans body. Defaults to 1.0.
        pos: Position (x, y, z) to spawn the c. elegans. Defaults to (0, 0, 0.05).
        quat: Quaternion (w, x, y, z) for c. elegans orientation. Defaults to (1, 0, 0, 0).
        rgba: RGBA color values (red, green, blue, alpha) for recoloring the body.
            If None, no recoloring is applied. Defaults to None.
        suffix: Suffix to append to body names. Defaults to "-worm".
    """
    print(f"Loading worm from {walker_xml_path}")
    worm = mujoco.MjSpec.from_file(walker_xml_path)

    # a) Convert motors to torque‑mode if requested
    if torque_actuators and hasattr(worm, "actuator"):
        logging.info("Converting to torque actuators")
        for actuator in worm.actuators:  # type: ignore[attr-defined]
            # Set gain to max force; remove bias terms if present
            if actuator.forcerange.size >= 2:
                actuator.gainprm[0] = actuator.forcerange[1]
            # reset custom bias terms
            actuator.biastype = mujoco.mjtBias.mjBIAS_NONE
            actuator.biasprm = np.zeros((10, 1))

    if rescale_factor != 1.0:
        logging.info(f"Rescaling body tree with scale factor {rescale_factor}")
        worm = dm_scale_spec(worm, rescale_factor)

    # Recolor the body if rgba is specified
    if rgba is not None:
        for body in worm.worldbody.bodies:
            _recolour_tree(body, rgba=rgba)

    spawn_site = spec.worldbody.add_frame(
        pos=pos,
        quat=quat,
    )
    root = worm.worldbody
    spawn_body = spawn_site.attach_body(root, "", suffix=suffix)
    self._suffix = suffix
    self._n_worms += 1

    if dim == 3:
        spawn_body.add_freejoint()
    elif dim == 2:
        spawn_body.add_joint(
            axis=(1, 0, 0),
            name="rootx" + suffix,
            type=mujoco.mjtJoint.mjJNT_SLIDE,
            pos=[0, 0, 0],
        )
        spawn_body.add_joint(
            axis=(0, 1, 0),
            name="rooty" + suffix,
            type=mujoco.mjtJoint.mjJNT_SLIDE,
            pos=[0, 0, 0],
        )
        spawn_body.add_joint(
            axis=(0, 0, 1),
            name="rootz" + suffix,
            type=mujoco.mjtJoint.mjJNT_SLIDE,
            pos=[0, 0, 0],
        )
        spawn_body.add_joint(
            axis=(0, 0, 1),
            name="free_body_rot" + suffix,
            type=mujoco.mjtJoint.mjJNT_HINGE,
            pos=[0, 0, 0],
        )

    for body_name in self.body_names:
        body = worm.body(f"{body_name}{suffix}")
        for geom in body.geoms:
            if geom.type == mujoco.mjtGeom.mjGEOM_SPHERE:
                spec.add_pair(
                    name=f"{body_name}_floor",
                    geomname1=geom.name,
                    geomname2="floor",
                    condim=3,
                    friction=friction,
                )  # , solimp=solimp)
      return spec

  
def compile(spec: mujoco.MjSpec,
               noslip_iterations: int=0,
               sim_dt: float=0.005,
               iterations: int=4,
               ls_iterations: int=4,
               solver: str="cg") -> None:
  """Compile the model from the mj_spec and put models to mjx.
  
  Args:
      spec: spec to compile
  """
  if not self._compiled or forced:
      spec.option.noslip_iterations = noslip_iterations
      mj_model = spec.compile()
      mj_model.opt.timestep = sim_dt
      # Increase offscreen framebuffer size to render at higher resolutions.
      mj_model.vis.global_.offwidth = 3840
      mj_model.vis.global_.offheight = 2160
      mj_model.opt.iterations = iterations
      mj_model.opt.ls_iterations = ls_iterations
      mj_model.opt.solver = {
          "cg": mujoco.mjtSolver.mjSOL_CG,
          "newton": mujoco.mjtSolver.mjSOL_NEWTON,
      }[solver]
      mjx_model = mjx.put_model(
          mj_model
      )  # , impl=self.config.mujoco_impl)
      return spec, mj_model, mjx_model
        
 spec = mujoco.MjSpec.from_file(arena_xml_path)
 spec = add_worm(spec)
 spec, mj_model, mjx_model = compile(spec)

Confirmations

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