Body Model Cannot Move Backwards?

[aaprasad]

Intro

Hi!

I am a graduate student at MIT, I use MuJoCo for my research on celegans neuromechanics.

My setup

Mujoco/mjx 3.4.1 python, linux 5090 cuda 12.

My question

I'm trying to train a model of neuromechanical model of c.elegans to imitate real animals behaviors and I'm able to get it to successfully imitate forward locomotion, however, backwards locomotion has been impossible! This is quite strange to me because as far as I can tell, the model is pretty much symmetric (slightly different lengthranges at the muscles and the tail is missing 1 muscle as in the real animal) and the assymetries that do exist shouldn't impair it this much right? I feel like it has to do with some small thing in the body model that I'm not understanding. I tried reversing the actions from a forward locomotion rollout as well as reverseing the actions and flipping the left and right muscle actions and it seems to push the model forwards from the tail rather than pull it backwards along the tail. I've attached a few videos below demonstrating the issue.
successful forward locomotion:

forword_locomotion.mp4

failed learning of backwards locomotion:

media_videos_videos_rollout_2_563878cc95f30509b5e0.mp4

time Reversed rollout:

reversed_locomotion.mp4

time reversed + directionally flipped rollout:

reversed_flipped_locomotion.mp4

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 (as in the real animal). 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 also scaled the model to CGS units so its 0.1 cm etc.

celegans_body_model.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.