perf: env_step sub-stage timing panel + optional numba TBB layer (+23% throughput on many-core)

docs/source/en/user_guide/getting_started/installation.md

@@ -92,6 +92,7 @@ training backend extra on top of them, and each option only selects or appends t
 | `--gpu` | `cuda`<br>`rocm` | Select the torch wheel flavor explicitly, overriding auto-detection |
 | `--skrl-jax` | — | SKRL on JAX backend, Linux only |
 | `--rslrl` | — | RSL-RL on PyTorch backend |
+| `--tbb` | — | Install Intel TBB for the numba parallel kernels; recommended for training on many-core servers (64+ cores), auto-activates once installed |
 | `--docs` | — | Add the toolchain (sphinx) needed to build the documentation locally |
 | `-h`, `--help` | — | Show the help message |
 

docs/source/zh_CN/user_guide/getting_started/installation.md

@@ -114,6 +114,7 @@ sh install.sh --all
 | `--gpu` | `cuda`<br>`rocm` | 指定 torch wheel 来源，覆盖自动探测 |
 | `--skrl-jax` | — | SKRL（JAX）训练后端，仅 Linux |
 | `--rslrl` | — | RSL-RL（PyTorch）训练后端 |
+| `--tbb` | — | 为 numba 并行 kernel 安装 Intel TBB 线程层；多核服务器（64 核以上）训练推荐启用，安装后自动生效 |
 | `--docs` | — | 追加本地构建文档所需的工具链（Sphinx） |
 | `-h`、`--help` | — | 显示帮助 |
 

install.sh

@@ -11,6 +11,8 @@
 #
 # Training backends are enabled with --skrl-torch (default), --skrl-jax or
 # --rslrl, named after the extras that install them; multiple flags combine.
+# --tbb installs Intel TBB so numba parallel kernels use its threading layer
+# (avoids the OpenMP per-region thread-wakeup storm on many-core hosts).
 
 set -e
 
@@ -25,6 +27,7 @@ usage() {
     echo "  --skrl-torch           Training backend: SKRL on PyTorch (default)"
     echo "  --skrl-jax             Training backend: SKRL on JAX (Linux only)"
     echo "  --rslrl                Training backend: RSL-RL on PyTorch"
+    echo "  --tbb                  Install Intel TBB for numba parallel kernels (recommended on many-core hosts)"
     echo "  -h, --help             Show this help"
 }
 
@@ -34,6 +37,7 @@ SKRL_TORCH=""
 SKRL_JAX=""
 RSLRL=""
 DOCS=""
+TBB=""
 
 while [ $# -gt 0 ]; do
     case "$1" in
@@ -57,6 +61,9 @@ while [ $# -gt 0 ]; do
         --rslrl)
             RSLRL=1
             ;;
+        --tbb)
+            TBB=1
+            ;;
         -h|--help)
             usage
             exit 0
@@ -162,6 +169,7 @@ else
     [ -n "$SKRL_JAX" ] && EXTRAS="$EXTRAS --extra skrl-jax"
     [ -n "$RSLRL" ] && EXTRAS="$EXTRAS --extra rslrl"
     [ -n "$DOCS" ] && EXTRAS="$EXTRAS --extra docs"
+    [ -n "$TBB" ] && EXTRAS="$EXTRAS --extra tbb"
     set -x
     uv sync --all-packages --no-default-groups --extra "$GPU"$EXTRAS
 fi

motrix_env_core/pyproject.toml

@@ -21,3 +21,10 @@ dependencies = [
     "omegaconf>=2.3,<2.4",
     "typing-extensions>=4.1",
 ]
+
+[project.optional-dependencies]
+# TBB threading layer for numba parallel kernels: on many-core machines the
+# default OpenMP layer pays a per-region thread-wakeup storm that dominates
+# short kernels. The wheel drops libtbb.so.12 into <venv>/lib, so the runtime
+# preloads it by full path (see motrix_env_core.numba.threading).
+tbb = ["tbb==2021.13.0"]

motrix_env_core/src/motrix_env_core/numba/manager/env.py

@@ -52,6 +52,7 @@
     TerminationTermCfg,
 )
 from motrix_env_core.numba.program import NumbaTaskProgram
+from motrix_env_core.numba.threading import preload_tbb
 from motrix_env_core.sim import (
     ModelQuery,
     PhysicsReadProgram,
@@ -343,6 +344,12 @@ def __init__(self, cfg: EnvCfgType, num_envs: int = 1, backend: str | None = Non
         """Construct a manager environment using the selected simulator backend."""
         if not isinstance(cfg, ManagerBasedEnvCfg):
             raise TypeError(f"{type(cfg).__name__} must inherit ManagerBasedEnvCfg.")
+        # Prefer the TBB threading layer for the parallel step kernels when the
+        # tbb extra is installed: the default OpenMP layer pays a thread-wakeup
+        # storm per kernel call that dominates short kernels on many-core
+        # machines. Must happen before the first parallel kernel executes.
+        if preload_tbb():
+            logger.info("Manager env %r: numba TBB threading layer available", type(self).__name__)
         super().__init__(cfg, num_envs)
         self._rand_seed = 1 if seed is None else seed
         from motrix_env_core.sim.registry import create_sim_backend, default_sim_backend_name

motrix_env_core/src/motrix_env_core/numba/threading.py

@@ -0,0 +1,59 @@
+# Copyright Motphys Technology Co., Ltd. 2025, 2026
+# SPDX-License-Identifier: Apache-2.0
+
+"""Runtime threading-layer helpers for numba parallel kernels.
+
+On many-core machines the default OpenMP layer pays a per-region thread-wakeup
+storm (``OMP_WAIT_POLICY=PASSIVE`` parks every worker between calls, and each
+``prange`` region must wake them all before any work starts). For short kernels
+called every environment step this overhead dominates the kernel itself. The
+TBB layer keeps a resident work-stealing pool and does not have this problem.
+
+The PyPI ``tbb`` wheel installs ``libtbb.so.12`` into ``<venv>/lib`` instead of
+``site-packages``, so numba's plain ``CDLL("libtbb.so.12")`` lookup cannot find
+it without ``LD_LIBRARY_PATH``. Preloading the library by full path with
+``RTLD_GLOBAL`` registers its SONAME in the process, and numba's later lookup by
+the same SONAME reuses the loaded handle.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import glob
+import sys
+from pathlib import Path
+
+_TBB_SONAME = "libtbb.so.12"
+_tbb_preload_result: bool | None = None
+
+
+def preload_tbb() -> bool:
+    """Preload ``libtbb.so.12`` from the active environment, if present.
+
+    Returns ``True`` when the library was loaded (or was already loaded), which
+    lets numba select the TBB threading layer for parallel kernels. Returns
+    ``False`` when the ``tbb`` extra is not installed — callers keep the numba
+    default layer in that case. The probe runs once and is cached: every
+    ``ManagerEnv`` construction calls this, and the glob + CDLL work should
+    not repeat per instance.
+    """
+    global _tbb_preload_result
+    if _tbb_preload_result is None:
+        _tbb_preload_result = _probe_tbb()
+    return _tbb_preload_result
+
+
+def _probe_tbb() -> bool:
+    if not sys.platform.startswith("linux"):
+        return False
+    candidates = sorted(glob.glob(str(Path(sys.prefix) / "lib" / f"{_TBB_SONAME}*")), reverse=True)
+    for path in candidates:
+        try:
+            ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)
+        except OSError:
+            continue
+        return True
+    return False
+
+
+__all__ = ["preload_tbb"]

motrix_env_core/src/motrix_env_core/perf.py

@@ -187,6 +187,27 @@ def snapshot(self) -> tuple[PerfNode, ...]:
             raise RuntimeError("Cannot snapshot Perf while a scope is active.")
         return tuple(root.snapshot() for root in self._roots.values())
 
+    def stage_mean_ms(self, root: str) -> dict[str, float]:
+        """Per-call mean milliseconds of every sub-scope under ``root``.
+
+        Keys are dotted paths relative to the root's children (``stage`` or
+        ``stage.sub``), so nested scope names — which may themselves contain
+        underscores — stay unambiguously separable. Returns ``{}`` when the
+        root has not run.
+        """
+        for node in self._roots.values():
+            if node.name == root:
+                means: dict[str, float] = {}
+
+                def emit(subtree: _MutablePerfNode, path: tuple[str, ...]) -> None:
+                    for child in subtree.children.values():
+                        means[".".join((*path, child.name))] = child.total_ns / 1e6 / max(child.count, 1)
+                        emit(child, (*path, child.name))
+
+                emit(node, ())
+                return means
+        return {}
+
     def call(self, name: str, function: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs) -> _R:
         """Call a function inside a named scope."""
         if not self._enabled:

motrix_env_core/tests/test_numba_threading.py

@@ -0,0 +1,10 @@
+# Copyright Motphys Technology Co., Ltd. 2025, 2026
+# SPDX-License-Identifier: Apache-2.0
+
+from motrix_env_core.numba.threading import preload_tbb
+
+
+def test_preload_tbb_returns_bool_without_raising() -> None:
+    # Whether the library exists depends on the tbb extra; the contract is
+    # only that discovery is side-effect-safe and reports success as bool.
+    assert isinstance(preload_tbb(), bool)

motrix_env_core/tests/test_perf.py

@@ -139,3 +139,29 @@ def test_active_perf_scope_contributes_to_activated_profiler() -> None:
 def test_active_perf_scope_is_disabled_without_an_activated_profiler() -> None:
     with active_perf_scope("backend"):
         pass
+
+
+def test_stage_mean_ms_flattens_nested_scopes_to_dotted_paths() -> None:
+    # two step calls; each scope entry/exit consumes one clock tick (8 per
+    # iteration). physics runs 3ms per call, transition 9ms with a 2ms read.
+    base = (0, 2_000_000, 5_000_000, 6_000_000, 10_000_000, 12_000_000, 15_000_000, 26_000_000)
+    ticks = base + tuple(t + 26_000_000 for t in base)
+    perf = Perf(enabled=True, clock=_Clock(*ticks))
+    for _ in range(2):
+        with perf.scope("step"):
+            with perf.scope("physics"):
+                pass
+            with perf.scope("transition"):
+                with perf.scope("read"):
+                    pass
+
+    means = perf.stage_mean_ms("step")
+
+    assert means["physics"] == pytest.approx(3.0)
+    assert means["transition"] == pytest.approx(9.0)
+    assert means["transition.read"] == pytest.approx(2.0)
+    assert "step" not in means  # the root itself is not part of the paths
+
+
+def test_stage_mean_ms_returns_empty_for_unrun_root() -> None:
+    assert Perf(enabled=True).stage_mean_ms("step") == {}

motrix_env_motrixsim/src/motrix_env_motrixsim/sim_data.py

@@ -11,6 +11,7 @@
 import numpy as np
 import numpy.typing as npt
 
+from motrix_env_core.perf import active_perf_scope
 from motrix_env_core.sim import PhysicsReadProgram, SimDataQuery, SimDataQueryCompiler
 
 FloatArray: TypeAlias = npt.NDArray[np.float32]
@@ -370,8 +371,10 @@ def execute(self, env_ids: npt.NDArray[np.int64] | None = None) -> None:
         if env_ids is not None and (env_ids.dtype != np.int64 or env_ids.ndim != 1):
             raise TypeError("Partial simulator read env_ids must be a one-dimensional int64 ndarray.")
         # One native call writes the authoritative arena — the full batch, or
-        # only the rows selected by env_ids.
-        self._read.program.execute(self._source, env_ids=env_ids)
+        # only the rows selected by env_ids. The perf scope nests under the
+        # caller's stage (e.g. transition/reset in the console timing tree).
+        with active_perf_scope("read"):
+            self._read.program.execute(self._source, env_ids=env_ids)
 
 
 def compile_read_program(

motrix_env_motrixsim/src/motrix_env_motrixsim/torch_env.py

@@ -12,6 +12,7 @@
 from gymnasium import Space, spaces
 
 from motrix_env_core.base import ABEnv, EnvCfg, ObsSpace
+from motrix_env_core.perf import Perf, perf_root
 from motrix_env_core.sim.backend import (
     RenderConfig,
     SimRenderer,
@@ -126,6 +127,7 @@ def __init__(
             raise ValueError("EnvCfg.scene must be configured")
         self._model = MotrixSimSceneCompiler().compile(cfg.scene, cfg.sim)
         self._render_spacing = cfg.render_spacing
+        self.perf = Perf()
 
     @property
     def model(self) -> mtx.SceneModel:
@@ -298,18 +300,29 @@ def _prev_physics_step(self):
         state.terminated.zero_()
         state.truncated.zero_()
 
+    @perf_root("step")
     def step(self, actions: torch.Tensor) -> TorchEnvState:
+        """Advance one control step.
+
+        Timing: ``apply_action`` -> ``physics`` -> ``transition`` (state update)
+        -> ``reset(env_ids)`` for done rows, mirroring ArrayEnv's scope names so
+        the console panel renders the same env_step sub-stage tree.
+        """
         if self._state is None:
             self.init_state()
         assert self._state is not None
 
         self._prev_physics_step()
-        self._state = self.apply_action(actions, self._state)
-        assert self._state is not None, "apply_action must return a valid TorchEnvState"
-        self.physics_step()
-        self._state = self.update_state(self._state)
+        with self.perf.scope("apply_action"):
+            self._state = self.apply_action(actions, self._state)
+            assert self._state is not None, "apply_action must return a valid TorchEnvState"
+        with self.perf.scope("physics"):
+            self.physics_step()
+        with self.perf.scope("transition"):
+            self._state = self.update_state(self._state)
         self._state = self._state.replace(obs=_as_obs(self._state.obs))
         self._state.episode_steps += 1
         self._update_truncate()
-        self._reset_done_envs()
+        with self.perf.scope("reset"):
+            self._reset_done_envs()
         return self._state

motrix_rl/src/motrix_rl/console.py

@@ -169,13 +169,27 @@ def _format_metric_items(items: Mapping[str, Any], *, precision: int = 3, signed
     return [f"{k} {_format_value(v, precision=precision, signed=signed)}" for k, v in items.items()]
 
 
+def _format_duration(seconds: float) -> str:
+    """Compact elapsed/remaining time: ``3h05m``, ``12m07s`` or ``0m45s``."""
+    t = max(0, int(seconds))
+    h, m, sec = t // 3600, (t % 3600) // 60, t % 60
+    return f"{h}h{m:02d}m" if h else f"{m}m{sec:02d}s"
+
+
+def _eta_seconds(stats: TrainingPanelStats) -> float | None:
+    """Remaining-time estimate from the cumulative iteration rate."""
+    if stats.iteration <= 0 or stats.elapsed_seconds <= 0.0:
+        return None
+    rate = stats.iteration / stats.elapsed_seconds
+    if rate <= 0.0:
+        return None
+    return (stats.total_iterations - stats.iteration) / rate
+
+
 def format_training_panel(stats: TrainingPanelStats, *, title: str = "rl") -> str:
     """Render a plain-text RL training panel from backend-provided scalar stats."""
 
-    def hms(t: float) -> str:
-        t = int(t)
-        h, m, sec = t // 3600, (t % 3600) // 60, t % 60
-        return f"{h}h{m:02d}m" if h else f"{m}m{sec:02d}s"
+    hms = _format_duration
 
     def si(n: float) -> str:
         for unit in ("", "k", "M"):
@@ -186,9 +200,11 @@ def si(n: float) -> str:
 
     width = 84
     pct = 100.0 * stats.iteration / max(stats.total_iterations, 1)
+    eta = _eta_seconds(stats)
+    eta_text = f" - eta ~{hms(eta)}" if eta is not None else ""
     header = (
         f" {title} - iter {stats.iteration}/{stats.total_iterations} ({pct:.1f}%) - "
-        f"{stats.steps_per_second:.0f} env-steps/s - {hms(stats.elapsed_seconds)}"
+        f"{stats.steps_per_second:.0f} env-steps/s - {hms(stats.elapsed_seconds)}{eta_text}"
     )
     lines = [
         "-" * width,
@@ -478,6 +494,18 @@ def _format_memory(memory: MemoryUsage | None) -> str:
     return f"{memory.used_bytes / gib:.1f}/{memory.total_bytes / gib:.1f} GiB"
 
 
+def _run_progress_time_text(stats: TrainingPanelStats):
+    """One ``elapsed · ~remaining`` line for the Run progress card."""
+    from rich.text import Text
+
+    elapsed = _format_duration(stats.elapsed_seconds)
+    eta = _eta_seconds(stats)
+    line = Text(f"{elapsed} elapsed", style="dim")
+    if eta is not None:
+        line.append(f" · ~{_format_duration(eta)} left", style="cyan")
+    return line
+
+
 def render_training_panel(stats: TrainingPanelStats, *, title: str = "rl", detail: bool = False):
     if not _RICH:
         raise RuntimeError("rich is not available")
@@ -537,7 +565,11 @@ def memory_style(memory: MemoryUsage | None) -> str:
     summary.add_row(
         card(
             f"Run progress ({progress * 100:.1f}%)",
-            Group(Text(f"{stats.iteration:,}/{stats.total_iterations:,} iters", style="white"), progress_row),
+            Group(
+                Text(f"{stats.iteration:,}/{stats.total_iterations:,} iters", style="white"),
+                progress_row,
+                _run_progress_time_text(stats),
+            ),
         ),
         card(
             "Episode stats",