Materialize SceneEntityCfg ids as device tensors to avoid per-step CUDA syncs

scripts/benchmarks/measure_synthetic_throughput.py

@@ -0,0 +1,151 @@
+"""Stress-test SceneEntityCfg ID indexing with a synthetic many-term env.
+
+Builds a cartpole env inflated with N extra observation and reward terms, each
+using a single-joint ``SceneEntityCfg`` so ``joint_ids`` resolves to a real
+``list[int]`` (not the optimized ``slice(None)``). This is a focused microbench
+for the ``tensor[:, cfg.joint_ids]`` implicit-CUDA-sync question raised in
+issue #1019: light physics keeps manager overhead as the dominant signal, and
+every extra term hits the suspected hot path.
+
+Reuses the timing helpers from ``measure_throughput`` for apples-to-apples
+comparison with the regression benchmark.
+"""
+
+from __future__ import annotations
+
+import statistics
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import torch
+import tyro
+
+import mjlab
+from mjlab.envs import ManagerBasedRlEnv
+from mjlab.envs.mdp import joint_pos_rel, joint_vel_rel
+from mjlab.managers.observation_manager import ObservationTermCfg
+from mjlab.managers.reward_manager import RewardTermCfg
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.tasks.cartpole.cartpole_env_cfg import (
+  cartpole_balance_env_cfg,
+  pole_angle_cos_sin,
+)
+
+sys.path.insert(0, str(Path(__file__).parent))
+from measure_throughput import measure_env_sps, measure_physics_sps  # noqa: E402
+
+
+def _synthetic_reward(
+  env: ManagerBasedRlEnv,
+  asset_cfg: SceneEntityCfg,
+) -> torch.Tensor:
+  """Indexing-heavy reward that mirrors the joint_pos_rel access pattern."""
+  asset = env.scene[asset_cfg.name]
+  jp = asset.data.joint_pos[:, asset_cfg.joint_ids]
+  jv = asset.data.joint_vel[:, asset_cfg.joint_ids]
+  return -(jp.pow(2).sum(-1) + 0.1 * jv.pow(2).sum(-1))
+
+
+def _inflate(cfg, extra_terms: int) -> None:
+  """Add ``extra_terms`` obs + reward terms that all do list[int] indexing.
+
+  Alternates between the slider and hinge joints; each picks exactly one joint
+  so ``joint_ids`` cannot collapse to ``slice(None)``.
+  """
+  if extra_terms <= 0:
+    return
+  joint_names = ("slider", "hinge_1")
+  obs_funcs = (joint_pos_rel, joint_vel_rel, pole_angle_cos_sin)
+  for i in range(extra_terms):
+    joint = joint_names[i % 2]
+    entity_cfg = SceneEntityCfg("cartpole", joint_names=(joint,))
+    obs_func = obs_funcs[i % len(obs_funcs)]
+    cfg.observations["actor"].terms[f"synth_obs_{i}"] = ObservationTermCfg(
+      func=obs_func,
+      params={"asset_cfg": entity_cfg},
+    )
+    cfg.observations["critic"].terms[f"synth_obs_{i}"] = ObservationTermCfg(
+      func=obs_func,
+      params={"asset_cfg": entity_cfg},
+    )
+    cfg.rewards[f"synth_reward_{i}"] = RewardTermCfg(
+      func=_synthetic_reward,
+      weight=0.0,
+      params={"asset_cfg": entity_cfg},
+    )
+
+
+@dataclass
+class SyntheticConfig:
+  num_envs: int = 4096
+  num_steps: int = 200
+  warmup_steps: int = 50
+  reps: int = 5
+  device: str = "cuda:0"
+  extra_terms: list[int] = field(default_factory=lambda: [0, 10, 50, 100])
+  """Sweep of extra-term counts to benchmark."""
+
+
+def benchmark_one(cfg: SyntheticConfig, extra_terms: int) -> dict:
+  print(f"\nSynthetic cartpole with {extra_terms} extra terms...")
+  env_cfg = cartpole_balance_env_cfg()
+  env_cfg.scene.num_envs = cfg.num_envs
+  _inflate(env_cfg, extra_terms)
+
+  env = ManagerBasedRlEnv(cfg=env_cfg, device=cfg.device)
+  env.reset()
+
+  action_dim = sum(env.action_manager.action_term_dim)
+  action = torch.zeros(env.num_envs, action_dim, device=env.device)
+  for _ in range(cfg.warmup_steps):
+    env.step(action)
+  torch.cuda.synchronize()
+
+  physics_samples: list[float] = []
+  env_samples: list[float] = []
+  for _ in range(cfg.reps):
+    physics_samples.append(measure_physics_sps(env, cfg.num_steps))
+    env.reset()
+    torch.cuda.synchronize()
+    env_samples.append(measure_env_sps(env, cfg.num_steps))
+    env.reset()
+    torch.cuda.synchronize()
+
+  env.close()
+
+  physics_med = statistics.median(physics_samples)
+  env_med = statistics.median(env_samples)
+  return {
+    "extra_terms": extra_terms,
+    "physics_sps": physics_med,
+    "env_sps": env_med,
+    "overhead_pct": 100 * (1 - env_med / physics_med),
+    "env_sps_min": min(env_samples),
+    "env_sps_max": max(env_samples),
+  }
+
+
+def main(cfg: SyntheticConfig) -> None:
+  print("Synthetic Many-Term Throughput Benchmark (cartpole base)")
+  print(f"  Envs: {cfg.num_envs}  Steps: {cfg.num_steps}  Reps: {cfg.reps}")
+
+  results = [benchmark_one(cfg, n) for n in cfg.extra_terms]
+
+  print("\n" + "=" * 72)
+  print(
+    f"{'Extra terms':>12} {'Physics SPS':>14} {'Env SPS':>14} "
+    f"{'Env min/max':>20} {'Overhead':>10}"
+  )
+  print("-" * 72)
+  for r in results:
+    minmax = f"{r['env_sps_min']:,.0f}/{r['env_sps_max']:,.0f}"
+    print(
+      f"{r['extra_terms']:>12} {r['physics_sps']:>14,.0f} "
+      f"{r['env_sps']:>14,.0f} {minmax:>20} {r['overhead_pct']:>9.1f}%"
+    )
+
+
+if __name__ == "__main__":
+  cfg = tyro.cli(SyntheticConfig, config=mjlab.TYRO_FLAGS)
+  main(cfg)

scripts/benchmarks/measure_throughput.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import json
+import statistics
 import subprocess
 import time
 from dataclasses import asdict, dataclass, field
@@ -35,13 +36,23 @@ class BenchmarkResult:
   physics_sps: float
   env_sps: float
   overhead_pct: float
+  reps: int = 1
+  physics_sps_min: float | None = None
+  physics_sps_max: float | None = None
+  env_sps_min: float | None = None
+  env_sps_max: float | None = None
 
   def __str__(self) -> str:
-    return (
-      f"{self.task} (dec={self.decimation}):\n"
-      f"  Physics SPS: {self.physics_sps:,.0f}\n"
-      f"  Env SPS:     {self.env_sps:,.0f}\n"
-      f"  Overhead:    {self.overhead_pct:.1f}%"
+    physics_line = f"  Physics SPS: {self.physics_sps:,.0f}"
+    env_line = f"  Env SPS:     {self.env_sps:,.0f}"
+    if self.reps > 1:
+      physics_line += (
+        f"  [min {self.physics_sps_min:,.0f} / max {self.physics_sps_max:,.0f}]"
+      )
+      env_line += f"  [min {self.env_sps_min:,.0f} / max {self.env_sps_max:,.0f}]"
+    header = f"{self.task} (dec={self.decimation}, reps={self.reps}):"
+    return "\n".join(
+      [header, physics_line, env_line, f"  Overhead:    {self.overhead_pct:.1f}%"]
     )
 
   def to_dict(self) -> dict:
@@ -61,6 +72,10 @@ class ThroughputConfig:
   warmup_steps: int = 50
   """Number of warmup steps before measuring."""
 
+  reps: int = 1
+  """Number of measurement reps per task. Reported value is the median; min/max
+  are recorded alongside. Warmup runs once at the start of each task."""
+
   device: str = "cuda:0"
   """Device to run on."""
 
@@ -76,6 +91,10 @@ class ThroughputConfig:
   tracking_motion: str = "rll_humanoid/wandb-registry-Motions/lafan_cartwheel:latest"
   """W&B artifact path for tracking task motion (entity/project/name:alias)."""
 
+  tracking_motion_file: str | None = None
+  """Optional local motion.npz path. When set, skips W&B and uses this file
+  for any MotionCommandCfg in the task config."""
+
   output_dir: Path | None = None
   """Output directory for JSON results. If None, results are only printed."""
 
@@ -130,10 +149,13 @@ def benchmark_task(task: str, cfg: ThroughputConfig) -> BenchmarkResult:
   if len(env_cfg.commands) > 0:
     motion_cmd = env_cfg.commands.get("motion")
     if isinstance(motion_cmd, MotionCommandCfg):
-      api = wandb.Api()
-      artifact = api.artifact(cfg.tracking_motion)
-      motion_dir = artifact.download()
-      motion_cmd.motion_file = str(Path(motion_dir) / "motion.npz")
+      if cfg.tracking_motion_file is not None:
+        motion_cmd.motion_file = cfg.tracking_motion_file
+      else:
+        api = wandb.Api()
+        artifact = api.artifact(cfg.tracking_motion)
+        motion_dir = artifact.download()
+        motion_cmd.motion_file = str(Path(motion_dir) / "motion.npz")
 
   env = ManagerBasedRlEnv(cfg=env_cfg, device=cfg.device)
   env.reset()
@@ -146,13 +168,21 @@ def benchmark_task(task: str, cfg: ThroughputConfig) -> BenchmarkResult:
   torch.cuda.synchronize()
 
   decimation = env.cfg.decimation
-  physics_sps = measure_physics_sps(env, cfg.num_steps)
-
-  env.reset()
-  torch.cuda.synchronize()
-
-  env_sps = measure_env_sps(env, cfg.num_steps)
 
+  physics_samples: list[float] = []
+  env_samples: list[float] = []
+  for rep in range(cfg.reps):
+    if cfg.reps > 1:
+      print(f"  rep {rep + 1}/{cfg.reps}")
+    physics_samples.append(measure_physics_sps(env, cfg.num_steps))
+    env.reset()
+    torch.cuda.synchronize()
+    env_samples.append(measure_env_sps(env, cfg.num_steps))
+    env.reset()
+    torch.cuda.synchronize()
+
+  physics_sps = statistics.median(physics_samples)
+  env_sps = statistics.median(env_samples)
   overhead_pct = 100 * (1 - env_sps / physics_sps)
 
   env.close()
@@ -165,6 +195,11 @@ def benchmark_task(task: str, cfg: ThroughputConfig) -> BenchmarkResult:
     physics_sps=physics_sps,
     env_sps=env_sps,
     overhead_pct=overhead_pct,
+    reps=cfg.reps,
+    physics_sps_min=min(physics_samples) if cfg.reps > 1 else None,
+    physics_sps_max=max(physics_samples) if cfg.reps > 1 else None,
+    env_sps_min=min(env_samples) if cfg.reps > 1 else None,
+    env_sps_max=max(env_samples) if cfg.reps > 1 else None,
   )