Partial reset advances actuator delay buffers in environments that were not reset #1216

Activity

1. lukehugh commented on Oct 8, 2026

   

   lukehugh

   on Oct 8, 2026

   Author

   More actions

   I quantified the downstream impact using mjlab's bundled cartpole balance task. Over 10 s, only environment 0 was reset; the other 24 environments received identical initial states and scripted actions. No external assets or trained policy were used.

   On #1218 (0bf9a0a), I compared no unrelated resets, scoped reset writes, and the affected whole-batch write reproduced by discarding reset env_ids at Scene.write_data_to_sim. Initialization was identical.

   Delay / reset scenario	Cart position RMSE	Wrapped pole angle RMSE	Mean absolute return difference / mean baseline return
   80 ms, explicit reset every control step	9.72 cm	52.52 degrees	6.07%
   80 ms, automatic reset every control step	25.21 cm	68.89 degrees	10.22%
   125 ms, explicit reset every control step	12.12 cm	39.61 degrees	1.91%
   125 ms, automatic reset every control step	11.46 cm	53.45 degrees	4.42%

   The 80 ms cases use timestep=10 ms, decimation=5; the 125 ms cases use 1 ms, decimation=50. Automatic resets use environment 0's timeout counter. Observation corruption is disabled. These are open-loop differences, not PPO performance losses: unstable cartpole dynamics amplify small errors. Reward changes have both signs; termination flags were unchanged.

   With #1218, all eight scenarios matched baseline state, observations, rewards, and physics-substep controls exactly. No-delay and 30 ms held-action cases were negative controls. In affected 80/125 ms cases, all 24 untouched environments had changed returns, even though immediate reset snapshots and end-of-control-step ctrl matched. The error occurs inside the decimation loop.

   This is important for our RM75_6F training with a 125 ms command delay and 1 ms physics timestep. Pending commands are part of the robot's dynamics. Unrelated resets make those dynamics depend on other environments in the batch, introducing a mismatch with single-robot deployment and affecting safety evaluation of delayed responses.

   The reproducer below runs on #1218 and saves trajectories and JSON metrics. Validated on RTX 4090, mjlab 1.6.0, MuJoCo 3.11.0, PyTorch 2.14.0+cu132.

   uv run python cartpole_benchmark.py --device cuda

   To run just the 80 ms explicit-reset case:

   uv run python cartpole_benchmark.py --device cuda --case native_80ms_every1

   Reproducer: cartpole_benchmark.py (run on PR #1218)

   """Replay identical actions through bundled cartpole with unrelated resets.
   
   Run on PR #1218's fixed implementation. The affected replay discards reset
   env_ids at Scene.write_data_to_sim, reproducing upstream's whole-batch write.
   No production files are modified. Environment zero is the reset distractor;
   the other 24 environments provide distinct initial states and action phases.
   """
   
   import argparse
   import contextlib
   import io
   import json
   import math
   import subprocess
   from importlib.metadata import version
   from pathlib import Path
   
   import numpy as np
   import torch
   
   from mjlab.actuator import XmlActuatorCfg
   from mjlab.entity import EntityArticulationInfoCfg
   from mjlab.envs import ManagerBasedRlEnv
   from mjlab.tasks.cartpole.cartpole_env_cfg import cartpole_balance_env_cfg
   
   
   def replay(env, steps, reset_every, affected, automatic):
     obs, _ = env.reset(seed=1216)
     original_write = env.scene.write_data_to_sim
     records = {key: [] for key in ("qpos", "qvel", "ctrl", "obs", "reward", "done")}
     physics_ctrl = []
   
     def tracked_write(env_ids=None):
       original_write(env_ids=None if affected else env_ids)
       if env_ids is None:
         physics_ctrl.append(env.sim.data.ctrl[1:].cpu().numpy().copy())
   
     env.scene.write_data_to_sim = tracked_write
     reset_changes = []
     phase = torch.arange(env.num_envs, device=env.device)[:, None] * 0.27
     ids = torch.tensor([0], device=env.device)
     try:
       for step in range(steps):
         if step > 0 and reset_every and step % reset_every == 0:
           if automatic:
             env.episode_length_buf[0] = env.max_episode_length - 1
           else:
             before = {
               key: getattr(env.sim.data, key)[1:].clone()
               for key in ("qpos", "qvel", "ctrl", "time")
             }
             env.reset(env_ids=ids)
             reset_changes.append(
               {
                 key: float((getattr(env.sim.data, key)[1:] - value).abs().max())
                 for key, value in before.items()
               }
             )
         t = step * env.step_dt
         action = 0.15 * torch.sin(2 * math.pi * 1.3 * t + phase)
         action += 0.08 * torch.sin(2 * math.pi * 0.4 * t + phase)
         obs, reward, terminated, truncated, _ = env.step(action)
         for key in ("qpos", "qvel", "ctrl"):
           records[key].append(getattr(env.sim.data, key)[1:].cpu().numpy().copy())
         records["obs"].append(obs["actor"][1:].cpu().numpy().copy())
         records["reward"].append(reward[1:].cpu().numpy().copy())
         records["done"].append((terminated | truncated)[1:].cpu().numpy().copy())
       delay = env.scene["cartpole"].actuators[0]._delay_buffer
       counters = None if delay is None else delay._step_count[1:].cpu().tolist()
     finally:
       env.scene.write_data_to_sim = original_write
     records["physics_ctrl"] = physics_ctrl
     return (
       {key: np.stack(value) for key, value in records.items()},
       counters,
       reset_changes,
     )
   
   
   def compare(reference, other):
     metrics = {}
     for key in reference:
       delta = other[key].astype(float) - reference[key].astype(float)
       metrics[key + "_rmse"] = float(np.sqrt(np.mean(delta**2)))
       metrics[key + "_max_abs_error"] = float(np.max(np.abs(delta)))
     angle_delta = other["qpos"][..., 1] - reference["qpos"][..., 1]
     wrapped = np.arctan2(np.sin(angle_delta), np.cos(angle_delta))
     metrics["pole_angle_wrapped_rmse_rad"] = float(np.sqrt(np.mean(wrapped**2)))
     metrics["cart_position_rmse_m"] = float(
       np.sqrt(np.mean((other["qpos"][..., 0] - reference["qpos"][..., 0]) ** 2))
     )
     returns = {
       name: data["reward"].sum(axis=0)
       for name, data in (("reference", reference), ("other", other))
     }
     metrics["reference_mean_return"] = float(returns["reference"].mean())
     metrics["other_mean_return"] = float(returns["other"].mean())
     metrics["mean_return_delta"] = float((returns["other"] - returns["reference"]).mean())
     metrics["mean_absolute_return_delta"] = float(
       np.abs(returns["other"] - returns["reference"]).mean()
     )
     metrics["max_absolute_return_delta"] = float(
       np.abs(returns["other"] - returns["reference"]).max()
     )
     metrics["affected_return_count"] = int(
       np.count_nonzero(np.abs(returns["other"] - returns["reference"]) > 1e-6)
     )
     metrics["termination_disagreement_count"] = int(
       np.count_nonzero(other["done"] != reference["done"])
     )
     return metrics
   
   
   def run_case(
     name, timestep, decimation, lag, reset_every, automatic, device, steps, output_dir
   ):
     cfg = cartpole_balance_env_cfg()
     cfg.scene.num_envs = 25
     cfg.sim.mujoco.timestep = timestep
     cfg.decimation = decimation
     cfg.episode_length_s = 1000
     cfg.auto_reset = automatic
     cfg.observations["actor"].enable_corruption = False
     cfg.scene.entities["cartpole"].articulation = EntityArticulationInfoCfg(
       actuators=(
         XmlActuatorCfg(
           target_names_expr=("slider",), delay_min_lag=lag, delay_max_lag=lag
         ),
       )
     )
     with contextlib.redirect_stdout(io.StringIO()):
       env = ManagerBasedRlEnv(cfg, device=device)
       try:
         reference, _, _ = replay(env, steps, 0, False, automatic)
         fixed, fixed_counts, fixed_jumps = replay(
           env, steps, reset_every, False, automatic
         )
         affected, affected_counts, affected_jumps = replay(
           env, steps, reset_every, True, automatic
         )
       finally:
         env.close()
     fixed_metrics = compare(reference, fixed)
     for key in ("qpos", "qvel", "ctrl", "obs", "reward", "done", "physics_ctrl"):
       assert fixed_metrics[key + "_max_abs_error"] == 0, (name, key)
     if not lag:
       for key in ("qpos", "qvel", "ctrl", "obs", "reward", "done", "physics_ctrl"):
         assert np.array_equal(reference[key], affected[key]), (name, key)
     traces = {
       f"{mode}_{key}": value
       for mode, record in (
         ("reference", reference),
         ("fixed", fixed),
         ("affected", affected),
       )
       for key, value in record.items()
     }
     np.savez_compressed(output_dir / (name + ".npz"), **traces)
     result = {
       "name": name,
       "physics_dt_s": timestep,
       "decimation": decimation,
       "lag": lag,
       "reset_every_control_steps": reset_every,
       "automatic_reset": automatic,
       "steps": steps,
       "duration_s": steps * timestep * decimation,
       "victim_envs": 24,
       "fixed_vs_reference": fixed_metrics,
       "affected_vs_reference": compare(reference, affected),
       "fixed_delay_counts": fixed_counts,
       "affected_delay_counts": affected_counts,
       "fixed_reset_max_changes": {
         key: max((x[key] for x in fixed_jumps), default=0)
         for key in ("qpos", "qvel", "ctrl", "time")
       },
       "affected_reset_max_changes": {
         key: max((x[key] for x in affected_jumps), default=0)
         for key in ("qpos", "qvel", "ctrl", "time")
       },
     }
     print(json.dumps(result), flush=True)
     return result
   
   
   if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--device", default="cpu")
     parser.add_argument("--steps", type=int, default=200)
     parser.add_argument("--output-dir", type=Path, default=Path(__file__).parent)
     parser.add_argument("--case", help="Run one named case instead of the full suite.")
     args = parser.parse_args()
     cases = [
       ("no_delay", 0.01, 5, 0, 1, False),
       ("native_30ms_every1", 0.01, 5, 3, 1, False),
       ("native_80ms_every10", 0.01, 5, 8, 10, False),
       ("native_80ms_every1", 0.01, 5, 8, 1, False),
       ("native_80ms_auto", 0.01, 5, 8, 1, True),
       ("rm75_125ms_every10", 0.001, 50, 125, 10, False),
       ("rm75_125ms_every1", 0.001, 50, 125, 1, False),
       ("rm75_125ms_auto", 0.001, 50, 125, 1, True),
     ]
     if args.case:
       cases = [case for case in cases if case[0] == args.case]
       if not cases:
         parser.error("Unknown case name")
     args.output_dir.mkdir(parents=True, exist_ok=True)
     results = [
       run_case(*case, args.device, args.steps, args.output_dir) for case in cases
     ]
     metadata = {
       "versions": {name: version(name) for name in ("mjlab", "mujoco", "torch")},
       "git_head": subprocess.check_output(
         ["git", "rev-parse", "HEAD"], text=True
       ).strip(),
       "device": args.device,
       "seed": 1216,
       "description": "Bundled cartpole with configured actuator delay; fixed-action replay.",
     }
     (args.output_dir / "cartpole_results.json").write_text(
       json.dumps({"metadata": metadata, "cases": results}, indent=2) + "\n"
     )