Merge pull request #22 from KraHsu/fix/perf-hot-path-allocations

Cut per-step allocations + CUDA syncs on training hot path

Author: KraHsu

examples/unitree/README.md

@@ -82,6 +82,24 @@ uv run genelab play Genelab-Tracking-Flat-Unitree-G1-v0 \
 `--agent` accepts `zero`, `random`, or `trained`. Without `--agent`, play defaults to
 `trained` when `--checkpoint` is set, else `zero`.
 
+## Long training runs
+
+For multi-hour runs (the G1 tasks default to `max_iterations=30_000` ≈ several wall-clock
+hours on 8×H200), set PyTorch's allocator to use growable segments **before** launching
+training. This dramatically reduces fragmentation-driven slowdown over time:
+
+```bash
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+```
+
+If training time drifts upward across iterations anyway, capture a profiler trace by
+re-running with `GENELAB_PROFILE=1` set. The trace lands under `logs/torch_profile/` and
+can be loaded by `tensorboard --logdir logs/torch_profile`. Take one trace at the start
+and one a few hours in to compare cumulative time per section. Optional env vars:
+`GENELAB_PROFILE_OUT`, `GENELAB_PROFILE_WAIT`, `GENELAB_PROFILE_WARMUP`,
+`GENELAB_PROFILE_ACTIVE`, `GENELAB_PROFILE_REPEAT` (see
+`src/genelab/rl/_profiler.py`). Only rank 0 emits a trace under distributed launches.
+
 ## Notes
 
 - The G1 MJCF and STL meshes are vendored under `assets/g1/`. They originate from Unitree's

src/genelab/entity/articulation.py

@@ -89,6 +89,7 @@ def __init__(self, cfg: ArticulationCfg, *, name: str = "robot") -> None:
         self._action_scale_tensor: torch.Tensor = torch.empty(0)
         self._actuators: dict[str, ActuatorBase] = {}
         self._joint_pos_target: torch.Tensor = torch.empty(0)
+        self._gravity_w: torch.Tensor = torch.empty(0)
         self._data: RobotState | None = None
 
     # ------------------------------------------------------------------ spawn / bind
@@ -118,6 +119,10 @@ def bind(self, num_envs: int, device: str) -> None:
             self._default_joint_pos.unsqueeze(0).expand(num_envs, -1).contiguous()
         )
         self._data = RobotState(num_envs, self._num_dofs, self._num_links, device)
+        # Cached constant gravity vector, broadcast across envs. Built once at bind to keep
+        # ``refresh`` allocation-free on the hot path (called every step + every reset).
+        self._gravity_w = torch.zeros(num_envs, 3, device=device)
+        self._gravity_w[:, 2] = -1.0
         self._assemble_actuators()
 
     # ------------------------------------------------------------------ internals
@@ -253,14 +258,14 @@ def refresh(self) -> None:
         rs = self._data
         robot = self._gs_handle
         try:
-            rs.root_pos = to_tensor(robot.get_pos(), self._device)
-            rs.root_quat = to_tensor(robot.get_quat(), self._device)
-            rs.root_lin_vel_w = to_tensor(robot.get_vel(), self._device)
-            rs.root_ang_vel_w = to_tensor(robot.get_ang(), self._device)
+            rs.root_pos.copy_(to_tensor(robot.get_pos(), self._device))
+            rs.root_quat.copy_(to_tensor(robot.get_quat(), self._device))
+            rs.root_lin_vel_w.copy_(to_tensor(robot.get_vel(), self._device))
+            rs.root_ang_vel_w.copy_(to_tensor(robot.get_ang(), self._device))
             joint_pos_full = to_tensor(robot.get_dofs_position(), self._device)
             joint_vel_full = to_tensor(robot.get_dofs_velocity(), self._device)
-            rs.joint_pos = joint_pos_full.index_select(-1, self._actuated_dof_idx)
-            rs.joint_vel = joint_vel_full.index_select(-1, self._actuated_dof_idx)
+            rs.joint_pos.copy_(joint_pos_full.index_select(-1, self._actuated_dof_idx))
+            rs.joint_vel.copy_(joint_vel_full.index_select(-1, self._actuated_dof_idx))
         except AttributeError:
             pass
         for attr, target in (
@@ -277,13 +282,14 @@ def refresh(self) -> None:
             except Exception:
                 continue
             tensor = to_tensor(value, self._device)
+            # Genesis may return either (num_envs, num_links, ...) or the bare (num_links, ...)
+            # shape on older single-env builds; ``expand`` broadcasts without allocating.
             if tensor.dim() == 2:
                 tensor = tensor.unsqueeze(0).expand(self._num_envs, -1, -1)
-            setattr(rs, target, tensor)
+            getattr(rs, target).copy_(tensor)
         rs.root_lin_vel_b = quat_rotate_inverse(rs.root_quat, rs.root_lin_vel_w)
         rs.root_ang_vel_b = quat_rotate_inverse(rs.root_quat, rs.root_ang_vel_w)
-        gravity_w = torch.tensor([0.0, 0.0, -1.0], device=self._device).expand_as(rs.root_lin_vel_w)
-        rs.projected_gravity_b = quat_rotate_inverse(rs.root_quat, gravity_w)
+        rs.projected_gravity_b = quat_rotate_inverse(rs.root_quat, self._gravity_w)
 
     def write_joint_state(
         self,

src/genelab/envs/manager_based_rl_env.py

@@ -245,7 +245,10 @@ def _reset_idx(self, env_ids: torch.Tensor) -> None:
     def step(
         self, action: torch.Tensor
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]:
-        action = action.to(self._device).clone()
+        # ``action_manager.process_action`` copies into a pre-allocated buffer, so we don't
+        # need ``.clone()`` here; the ``.to`` call is a no-op when the action is already on
+        # the right device.
+        action = action.to(self._device, non_blocking=True)
         self.action_manager.process_action(action)
         for _ in range(self._decimation):
             self.action_manager.apply_action()
@@ -269,7 +272,9 @@ def step(
             self._reset_idx(reset_ids)
             self._articulation.refresh()
         obs = self.observation_manager.compute()
-        return obs, reward, terminated, time_outs, dict(self._extras)
+        # The RSL-RL wrapper shallow-copies ``extras`` itself before mutating, so we hand back
+        # the live dict to skip an unnecessary per-step copy.
+        return obs, reward, terminated, time_outs, self._extras
 
     def close(self) -> None:
         self._scene.close()

src/genelab/managers/curriculum_manager.py

@@ -39,10 +39,19 @@ def active_terms(self) -> list[str]:
 
     def compute(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float]:
         extras: dict[str, float] = {}
+        # Collect tensor-valued curriculum terms separately so all their means come back
+        # with a single host sync (one ``.tolist()`` instead of one ``.item()`` per term).
+        tensor_names: list[str] = []
+        tensor_means: list[torch.Tensor] = []
         for name, term_cfg in zip(self._term_names, self._term_cfgs, strict=True):
             value = term_cfg.func(self._env, env_ids, **term_cfg.params)
             if isinstance(value, torch.Tensor):
-                extras[f"Curriculum/{name}"] = float(value.float().mean().item())
+                tensor_names.append(name)
+                tensor_means.append(value.float().mean())
             elif value is not None:
                 extras[f"Curriculum/{name}"] = float(value)
+        if tensor_means:
+            means_list = torch.stack(tensor_means).tolist()
+            for name, mean in zip(tensor_names, means_list):
+                extras[f"Curriculum/{name}"] = float(mean)
         return extras

src/genelab/managers/reward_manager.py

@@ -64,13 +64,18 @@ def active_terms(self) -> list[str]:
     def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float]:
         if env_ids is None:
             env_ids = slice(None)
-        extras: dict[str, float] = {}
         max_episode_length_s = getattr(self._env, "max_episode_length_s", 1.0) or 1.0
-        for name, episode_sum in self._episode_sums.items():
-            mean = torch.mean(episode_sum[env_ids]).item()
-            extras[f"Episode_Reward/{name}"] = mean / max_episode_length_s
-            episode_sum[env_ids] = 0.0
-        return extras
+        names = list(self._episode_sums.keys())
+        if not names:
+            return {}
+        # Stack per-term means into a single tensor and pull all values to host with one sync.
+        # The per-term ``.item()`` loop was a per-reset CUDA sync per reward term (6+ per step
+        # on busy reset paths in g1-velocity-flat).
+        means = torch.stack([torch.mean(self._episode_sums[name][env_ids]) for name in names])
+        means_list = (means / max_episode_length_s).tolist()
+        for name in names:
+            self._episode_sums[name][env_ids] = 0.0
+        return {f"Episode_Reward/{name}": value for name, value in zip(names, means_list)}
 
     def compute(self, dt: float) -> torch.Tensor:
         self._reward_buf.zero_()

src/genelab/managers/termination_manager.py

@@ -66,12 +66,14 @@ def terminated(self) -> torch.Tensor:
     def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float]:
         if env_ids is None:
             env_ids = slice(None)
-        extras: dict[str, float] = {}
-        for name, term_done in self._term_dones.items():
-            extras[f"Episode_Termination/{name}"] = float(
-                torch.count_nonzero(term_done[env_ids]).item()
-            )
-        return extras
+        names = list(self._term_dones.keys())
+        if not names:
+            return {}
+        # One CUDA sync for all termination terms, instead of one per term.
+        counts = torch.stack(
+            [torch.count_nonzero(self._term_dones[name][env_ids]) for name in names]
+        ).tolist()
+        return {f"Episode_Termination/{name}": float(c) for name, c in zip(names, counts)}
 
     def compute(self) -> torch.Tensor:
         self._truncated_buf.zero_()

src/genelab/rl/_profiler.py

@@ -0,0 +1,72 @@
+"""Opt-in ``torch.profiler`` wrapper for diagnosing training slowdowns.
+
+Activated by ``GENELAB_PROFILE=1`` in the environment. Disabled by default so existing
+runs are unaffected. Only rank 0 emits a trace under distributed launches, to avoid every
+worker writing simultaneously.
+
+Environment variables (all optional):
+
+* ``GENELAB_PROFILE=1`` — turn the profiler on. Any other value is treated as off.
+* ``GENELAB_PROFILE_OUT`` — directory the TensorBoard trace handler writes to. Default
+  ``logs/torch_profile`` (relative to the current working directory).
+* ``GENELAB_PROFILE_WAIT`` / ``_WARMUP`` / ``_ACTIVE`` / ``_REPEAT`` — passed straight to
+  ``torch.profiler.schedule``. Defaults: ``wait=10``, ``warmup=5``, ``active=10``,
+  ``repeat=2``. Tweak if iterations are very long or very short.
+
+The profiler's ``.step()`` must be advanced once per training iteration for the schedule
+to move forward. We wire this in by patching the active runner's ``learn`` loop one level
+up (see ``genelab.rl.runner.train_task``), which inserts a profiler-step callback. If the
+patching path isn't taken, the context manager still records a single contiguous trace.
+"""
+
+import os
+from contextlib import contextmanager
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from genelab.rl.distributed import is_main_process
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+
+def profiler_enabled() -> bool:
+    return os.environ.get("GENELAB_PROFILE", "0") == "1"
+
+
+def _read_int(name: str, default: int) -> int:
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    try:
+        return int(raw)
+    except ValueError:
+        return default
+
+
+@contextmanager
+def maybe_profile() -> "Generator[Any, None, None]":
+    """Yield a profiler-step callback. No-op when ``GENELAB_PROFILE != 1`` or not rank 0."""
+    if not profiler_enabled() or not is_main_process():
+        yield None
+        return
+    import torch
+    from torch.profiler import ProfilerActivity, profile, schedule, tensorboard_trace_handler
+
+    out_dir = Path(os.environ.get("GENELAB_PROFILE_OUT", "logs/torch_profile"))
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    activities = [ProfilerActivity.CPU]
+    if torch.cuda.is_available():
+        activities.append(ProfilerActivity.CUDA)
+
+    sched = schedule(
+        wait=_read_int("GENELAB_PROFILE_WAIT", 10),
+        warmup=_read_int("GENELAB_PROFILE_WARMUP", 5),
+        active=_read_int("GENELAB_PROFILE_ACTIVE", 10),
+        repeat=_read_int("GENELAB_PROFILE_REPEAT", 2),
+    )
+    handler = tensorboard_trace_handler(str(out_dir))
+
+    with profile(activities=activities, schedule=sched, on_trace_ready=handler) as prof:
+        yield prof.step

src/genelab/rl/runner.py

@@ -9,6 +9,7 @@
 
 from genelab.cache import ensure_project_cache
 from genelab.registry import TASKS
+from genelab.rl._profiler import maybe_profile, profiler_enabled
 from genelab.rl.config import RslRlOnPolicyRunnerCfg
 from genelab.rl.distributed import is_main_process
 from genelab.rl.rsl_rl_wrapper import RslRlVecEnvWrapper
@@ -135,7 +136,21 @@ def train_task(
     )
     if resume_from is not None:
         runner.load(str(resume_from))
-    runner.learn(num_learning_iterations=agent_cfg.max_iterations)
+    with maybe_profile() as prof_step:
+        if prof_step is not None and profiler_enabled():
+            # Advance the profiler schedule once per env step. RSL-RL doesn't expose a
+            # per-iteration hook, so the wrapper's ``step`` is the closest fire-once-per-
+            # rollout-step seam. Multiply WAIT/WARMUP/ACTIVE by ``num_steps_per_env`` if
+            # you want the schedule expressed in PPO iterations.
+            original_step = wrapped.step
+
+            def _step_with_profiler(*args: Any, **kwargs: Any) -> Any:
+                result = original_step(*args, **kwargs)
+                prof_step()
+                return result
+
+            wrapped.step = _step_with_profiler  # type: ignore[method-assign]
+        runner.learn(num_learning_iterations=agent_cfg.max_iterations)
     env.close()
     return log_dir
 

src/genelab/sensor/body_velocity.py

@@ -10,7 +10,7 @@
 import torch
 
 from genelab.sensor.sensor import Sensor, SensorCfg
-from genelab.utils.math import matrix_from_quat, quat_apply_inverse
+from genelab.utils.math import quat_apply, quat_apply_inverse
 
 if TYPE_CHECKING:
     from genelab.envs.manager_based_rl_env import ManagerBasedRlEnv
@@ -83,7 +83,6 @@ def _compute_data(self) -> torch.Tensor:
         if self._cfg_typed.measure == "ang_vel":
             return quat_apply_inverse(link_quat, link_ang_vel_w) + self._bias
         link_lin_vel_w = rs.link_lin_vel_w[:, self._link_idx]
-        R = matrix_from_quat(link_quat)
-        r_world = torch.matmul(R, self._offset_local)
+        r_world = quat_apply(link_quat, self._offset_local.expand_as(link_lin_vel_w))
         v_site_w = link_lin_vel_w + torch.cross(link_ang_vel_w, r_world, dim=-1)
         return quat_apply_inverse(link_quat, v_site_w) + self._bias

src/genelab/sensor/contact.py

@@ -117,7 +117,7 @@ def update(self, dt: float) -> None:
         force = self._read_link_forces()
         in_contact = force.norm(dim=-1) > self._cfg_typed.force_threshold
         state = self._air_state
-        # Detect just-landed / just-lifted-off transitions before mutating the timers.
+        # Capture transitions before mutating the timers (the timer values matter pre-tick).
         first_contact = in_contact & (state.current_air_time > 0)
         first_detached = (~in_contact) & (state.current_contact_time > 0)
         state.last_air_time = torch.where(
@@ -126,16 +126,13 @@ def update(self, dt: float) -> None:
         state.last_contact_time = torch.where(
             first_detached, state.current_contact_time + dt, state.last_contact_time
         )
-        state.current_air_time = torch.where(
-            in_contact,
-            torch.zeros_like(state.current_air_time),
-            state.current_air_time + dt,
-        )
-        state.current_contact_time = torch.where(
-            in_contact,
-            state.current_contact_time + dt,
-            torch.zeros_like(state.current_contact_time),
-        )
+        # Tick all timers, then mask-zero whichever phase ended. Original ``torch.where`` with
+        # ``torch.zeros_like`` allocated two (B, N) tensors per step; this in-place form is
+        # zero-allocation and removes the dominant per-step churn from the contact sensor.
+        state.current_air_time += dt
+        state.current_air_time.masked_fill_(in_contact, 0.0)
+        state.current_contact_time += dt
+        state.current_contact_time.masked_fill_(~in_contact, 0.0)
 
     def reset(self, env_ids: torch.Tensor | None = None) -> None:
         super().reset(env_ids)