PERF: G1Flat training throughput — contact sensor pooled reset + event sync elimination

embodichain/lab/gym/envs/managers/event_manager.py

@@ -310,6 +310,24 @@ def apply(
                     self._call_event_functor(
                         mode, functor_name, functor_cfg, self._env, None
                     )
+                elif functor_cfg.interval_step == 1:
+                    # Every environment is due on every step. Pass the manager-
+                    # owned all-row ID tensor (functors may call len(env_ids));
+                    # it is created once during preparation and treated as
+                    # read-only (~0.03 µs reuse vs ~5 µs fresh arange).
+                    if getattr(self, "_all_row_ids", None) is None:
+                        self._all_row_ids = torch.arange(
+                            self._interval_functor_step_count[index].numel(),
+                            device=self._interval_functor_step_count[index].device,
+                            dtype=torch.long,
+                        )
+                    self._call_event_functor(
+                        mode,
+                        functor_name,
+                        functor_cfg,
+                        self._env,
+                        self._all_row_ids,
+                    )
                 else:
                     valid_env_ids = (
                         (

embodichain/lab/sim/sensors/contact_history.py

@@ -77,13 +77,21 @@ def __init__(
         self.include_unknown_counterpart = include_unknown_counterpart
         self.force_threshold = float(force_threshold)
         shape = actor_ids.shape
-        self.contact = torch.zeros(shape, device=actor_ids.device, dtype=torch.bool)
-        self.found = torch.zeros_like(self.contact)
-        self.first_contact = torch.zeros_like(self.contact)
-        self.force = torch.zeros((*shape, 3), device=actor_ids.device)
-        self.peak_force = torch.zeros_like(self.force)
-        self.current_air_time = torch.zeros(shape, device=actor_ids.device)
-        self.last_air_time = torch.zeros_like(self.current_air_time)
+        # Pooled allocations: one storage per dtype/shape family so a reset
+        # fuses into a handful of launches instead of one CUDA write per
+        # field. The public attributes remain per-field contiguous views.
+        self._bool_pool = torch.zeros(
+            (3, *shape), device=actor_ids.device, dtype=torch.bool
+        )
+        self.contact = self._bool_pool[0]
+        self.found = self._bool_pool[1]
+        self.first_contact = self._bool_pool[2]
+        self._force_pool = torch.zeros((2, *shape, 3), device=actor_ids.device)
+        self.force = self._force_pool[0]
+        self.peak_force = self._force_pool[1]
+        self._air_pool = torch.zeros((2, *shape), device=actor_ids.device)
+        self.current_air_time = self._air_pool[0]
+        self.last_air_time = self._air_pool[1]
         self.contact_count = torch.zeros(shape[0], device=actor_ids.device)
         self._hits = torch.zeros(shape, device=actor_ids.device, dtype=torch.int32)
         self._env_hits = torch.zeros(
@@ -214,17 +222,24 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | None = None) -> None:
         Args:
             env_ids: Rows to clear. None selects every environment.
         """
-        ids = slice(None) if env_ids is None else env_ids
-        for value in (
-            self.contact,
-            self.found,
-            self.first_contact,
-            self.force,
-            self.peak_force,
-            self.current_air_time,
-            self.last_air_time,
-            self.contact_count,
-            self._hits,
-            self._env_hits,
-        ):
-            value[ids] = 0
+        if env_ids is None:
+            # One fused multi-tensor zero for the full-reset case.
+            torch._foreach_zero_(
+                [
+                    self._bool_pool,
+                    self._force_pool,
+                    self._air_pool,
+                    self.contact_count,
+                    self._hits,
+                    self._env_hits,
+                ]
+            )
+            return
+        ids = env_ids
+        # Pooled rows: one advanced-indexing write per dtype/shape family.
+        self._bool_pool[:, ids] = False
+        self._force_pool[:, ids] = 0
+        self._air_pool[:, ids] = 0
+        self.contact_count[ids] = 0
+        self._hits[ids] = 0
+        self._env_hits[ids] = 0

tests/gym/envs/managers/test_event_manager_interval_fastpath.py

@@ -0,0 +1,50 @@
+"""interval_step==1 快路径：functor 收到显式全量 ID 张量（非 None）。"""
+import torch
+
+from embodichain.lab.gym.envs.managers.event_manager import EventManager
+
+
+class _RecordingFunctor:
+    """记录 env_ids 的哑 functor。"""
+
+    def __init__(self):
+        self.calls = []
+
+    def __call__(self, env, env_ids, **kwargs):
+        self.calls.append(
+            None if env_ids is None else env_ids.clone()
+        )
+
+
+def test_interval_one_fastpath_passes_explicit_all_row_ids(monkeypatch):
+    """interval_step==1 且非 global：functor 收到全量 ID 张量（非 None）。"""
+    manager = EventManager.__new__(EventManager)
+    manager._env = object()
+    manager._seed = None
+
+    num_envs = 4096
+    counter = torch.zeros(num_envs, dtype=torch.long, device="cpu")
+    manager._interval_functor_step_count = [counter]
+
+    class _Cfg:
+        is_global = False
+        interval_step = 1
+
+    manager._mode_functor_names = {"interval": ["push_robot"]}
+    manager._mode_functor_cfgs = {"interval": [_Cfg()]}
+
+    # 拦截 _call_event_functor，记录传入的 env_ids
+    received = []
+    monkeypatch.setattr(
+        manager, "_call_event_functor",
+        lambda mode, name, cfg, env, ids: received.append(ids),
+    )
+
+    manager.apply(mode="interval")
+
+    assert len(received) == 1
+    ids = received[0]
+    assert ids is not None, "快路径传了 None——functor 的 len(env_ids) 会崩"
+    assert ids.shape[0] == num_envs
+    assert ids.dtype == torch.long
+    assert torch.equal(ids, torch.arange(num_envs, dtype=torch.long))

tests/sim/sensors/test_contact_history.py

@@ -146,3 +146,26 @@ def test_sparse_rows_use_device_count_and_environment_qualified_actor_ids(device
     )
     assert history.contact.tolist() == [[True, False], [False, True]]
     assert history.contact_count.tolist() == [1.0, 1.0]
+
+
+def test_full_reset_clears_every_field_in_all_pools():
+    """env_ids=None 路径：foreach 融合清零必须覆盖全部字段与全部行。"""
+    history = ContactHistory(torch.tensor([[10], [20]]))
+    history.update(sample([[[0, 10]], [[0, 20]]]), 0.1)
+    history.update(sample([[[0, 10]], [[0, 20]]]), 0.1)
+    fields = (
+        "contact",
+        "found",
+        "first_contact",
+        "force",
+        "peak_force",
+        "current_air_time",
+        "last_air_time",
+        "contact_count",
+        "_hits",
+        "_env_hits",
+    )
+    history.reset(None)
+    for name in fields:
+        value = getattr(history, name)
+        assert not value.any(), f"{name} 未被清零"