feat: redesign FastSAC console rich panel (#15)

THIRD_PARTY_NOTICES.md

@@ -28,7 +28,8 @@ separately downloaded optional asset.
 ## Runtime dependencies
 
 The following are installed from external distributions and remain under their
-own licenses: MotrixSim, MuJoCo, Gymnasium, NumPy, Numba, Hydra/OmegaConf,
+own licenses: MotrixSim, MuJoCo, Gymnasium, NumPy, Numba, nvidia-ml-py (NVML
+bindings used for training-panel GPU metrics), Hydra/OmegaConf,
 SKRL, RSL-RL, PyTorch/JAX, ONNX Runtime, TensorBoard, and the Unitree SDK2
 Python package. Their licenses are not replaced by the MotrixLab license.
 The release process should generate a dependency license report from the final

motrix_rl/pyproject.toml

@@ -20,6 +20,7 @@ dependencies = [
     "motrix-deploy",
     "numpy>=1.26",
     "omegaconf>=2.3,<2.4",
+    "nvidia-ml-py>=13.610.43",
 ]
 
 [project.optional-dependencies]

motrix_rl/src/motrix_rl/fastsac/async_impl/train.py

@@ -198,6 +198,7 @@ def _drain_child_errors() -> list[tuple[str, str]]:
                 logging_interval,
                 save_interval,
                 str(self._context.run_dir),
+                self._env_name,
                 str(self._context.checkpoint_dir),
                 self._context.checkpoint_format,
                 self._resume_from,

motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py

@@ -39,7 +39,7 @@
 from motrix_rl.fastsac.wrap import FastSacEnvWrap
 from motrix_rl.fastsac.wrap_np import FastSacNpEnvWrap
 from motrix_rl.fastsac.wrap_torch import FastSacTorchEnvWrap
-from motrix_rl.system_metrics import CpuLoadSampler
+from motrix_rl.system_metrics import CpuLoadSampler, GpuMemoryUsageSampler, GpuUtilizationSampler, MemoryUsageSampler
 
 
 def _timing_mean(values: list[float]) -> float:
@@ -190,6 +190,7 @@ def run_learner_process(
     logging_interval: int,
     save_interval: int,
     run_dir: str,
+    env_name: str,
     checkpoint_dir: str,
     checkpoint_format: str,
     resume_from: str | None,
@@ -233,12 +234,17 @@ def run_learner_process(
         last_metrics = None
         next_log = ((resume_step // logging_interval) + 1) * logging_interval if logging_interval > 0 else 0
         next_save = ((resume_step // save_interval) + 1) * save_interval if save_interval > 0 else 0
-        t_learn_win = 0.0  # wall-clock spent in gradient updates this log window
+        t_learn_win = 0.0  # wall-clock spent in learner train calls this log window
+        learner_train_samples_ms: list[float] = []
         learner_drain_samples_ms: list[float] = []
         learner_breakdown_samples_ms: dict[str, list[float]] = {}
         learner_ring_wait_samples_ms: list[float] = []
         learner_gate_wait_samples_ms: list[float] = []
         cpu_sampler = CpuLoadSampler()
+        gpu_sampler = GpuUtilizationSampler()
+        memory_sampler = MemoryUsageSampler()
+        gpu_memory_sampler = GpuMemoryUsageSampler()
+        last_checkpoint_path: str | None = None
 
         def _drain_stats():
             nonlocal last_stats
@@ -259,6 +265,10 @@ def _drain_stats():
                 last_metrics = metrics
                 elapsed_learn_s = time.perf_counter() - t_l
                 t_learn_win += elapsed_learn_s
+                learner_train_samples_ms.append(elapsed_learn_s * 1000.0)
+                # Keep the raw total cost of one agent.update(n) call. The
+                # timing tree uses total-cost semantics, not per-gradient-step
+                # normalization.
                 for key, value in learner.agent._last_update_timing_ms.items():
                     learner_breakdown_samples_ms.setdefault(key, []).append(value)
                 learner_breakdown_samples_ms.setdefault("publish", []).append(learner._last_publish_ms)
@@ -288,23 +298,22 @@ def _drain_stats():
                 # these do NOT sum to 100% like the sync panel):
                 #   timing_ms[collect] — collector's avg ms per env-step batch (from queue)
                 #   timing_ms[wait] — avg ring-backpressure wait per batch
-                #   learn_ms   — learner's avg ms per gradient update
+                #   learn_ms   — learner's avg ms per train call (one UTD
+                #                execution, which may run several gradient
+                #                updates); idle waits are not included
                 #   learn_pct  — fraction of learner wall-clock spent updating vs
                 #                idle/starved (≈100% when GPU-bound, lower if the
                 #                collector can't keep the buffer fed)
                 # Window means only: live-panel percentiles are noise at these
                 # sample counts (benchmarks own the tail statistics).
-                updates_delta = updates - last_update_idx
-                learn_ms = t_learn_win * 1000.0 / max(updates_delta, 1)
                 learn_pct = 100.0 * t_learn_win / max(now - last_log_time, 1e-9)
                 collector_timing_ms = last_stats.get("timing_ms", {})
                 collector_timing_detail_ms = {
                     key: value for key, value in collector_timing_ms.items() if key != "collect"
                 }
-                # Panel tree is per-process; the headline per-batch/per-update
-                # means are folded into the group titles, sub-stages nest under
-                # "sync" / "update" branches.
-                collect_ms = collector_timing_ms.get("collect", 0.0)
+                # Panel tree is per-process; the headline collect/learn means
+                # live on TrainingPanelStats, sub-stages nest under "sync" /
+                # "update" branches.
                 collector_items: dict[str, Any] = {}
                 sync_items: dict[str, float] = {}
                 for key, value in collector_timing_detail_ms.items():
@@ -314,20 +323,27 @@ def _drain_stats():
                         collector_items[key] = value
                 if sync_items:
                     collector_items["sync"] = {"total": collector_items.pop("sync", 0.0), **sync_items}
-                timing_groups = {f"collector [yellow]{collect_ms:.1f}[/]ms": collector_items}
+                timing_groups = {"collector": collector_items}
                 learner_items: dict[str, Any] = {}
+                drain_ms = _timing_mean(learner_drain_samples_ms) if learner_drain_samples_ms else 0.0
+                ring_wait_ms = _timing_mean(learner_ring_wait_samples_ms) if learner_ring_wait_samples_ms else 0.0
+                gate_wait_ms = _timing_mean(learner_gate_wait_samples_ms) if learner_gate_wait_samples_ms else 0.0
                 if learner_drain_samples_ms:
-                    learner_items["drain"] = _timing_mean(learner_drain_samples_ms)
+                    learner_items["drain"] = drain_ms
                 if learner_ring_wait_samples_ms:
-                    learner_items["ring wait"] = _timing_mean(learner_ring_wait_samples_ms)
+                    learner_items["ring wait"] = ring_wait_ms
                 if learner_gate_wait_samples_ms:
-                    learner_items["gate wait"] = _timing_mean(learner_gate_wait_samples_ms)
-                if learner_breakdown_samples_ms:
-                    learner_items["update"] = {
-                        key: _timing_mean(values) for key, values in learner_breakdown_samples_ms.items()
-                    }
+                    learner_items["gate wait"] = gate_wait_ms
+                update_items = {key: _timing_mean(values) for key, values in learner_breakdown_samples_ms.items()}
+                if update_items:
+                    # publish is a child stage of the learner update in the
+                    # panel, so include it in the displayed update total too.
+                    if "publish" in update_items and "total" in update_items:
+                        update_items["total"] += update_items["publish"]
+                    learner_items["update"] = update_items
                 if learner_items:
-                    timing_groups[f"learner [magenta]{learn_ms:.1f}[/]ms"] = learner_items
+                    timing_groups["learner"] = learner_items
+                learn_ms = _timing_mean(learner_train_samples_ms) if learner_train_samples_ms else 0.0
                 stats = TrainingPanelStats(
                     iteration=step,
                     total_iterations=num_iterations,
@@ -348,8 +364,12 @@ def _drain_stats():
                     timing_groups=timing_groups,
                     diagnostics={"UTD": utd},
                     cpu_load=cpu_sampler.sample(),
+                    gpu_utilization_percent=gpu_sampler.sample(),
+                    memory_usage=memory_sampler.sample(),
+                    gpu_memory_usage=gpu_memory_sampler.sample(),
+                    checkpoint_path=last_checkpoint_path,
                 )
-                emit_training_panel(live, stats, title="motrix.fastsac (async)")
+                emit_training_panel(live, stats, title=f"{env_name}/motrix.fastsac")
                 if writer is not None:
                     writer.add_scalar("rollout/mean_return", last_stats["return"], step)
                     writer.add_scalar("rollout/mean_ep_len", last_stats["ep_len"], step)
@@ -364,7 +384,7 @@ def _drain_stats():
                     writer.add_scalar("perf/collect_ms_per_batch", collector_timing_ms.get("collect", 0.0), step)
                     for k, v in collector_timing_detail_ms.items():
                         writer.add_scalar(f"perf/collector_{k}_ms", v, step)
-                    writer.add_scalar("perf/learn_ms_per_update", learn_ms, step)
+                    writer.add_scalar("perf/learn_ms_total", learn_ms, step)
                     writer.add_scalar("perf/learn_pct", learn_pct, step)
                     for k, v in last_stats["env_metrics"].items():
                         writer.add_scalar(f"metrics/{k}", v, step)
@@ -375,6 +395,7 @@ def _drain_stats():
                             writer.add_scalar(f"train/{k}", v, step)
                 last_log_time, last_log_step, last_update_idx = now, step, updates
                 t_learn_win = 0.0
+                learner_train_samples_ms = []
                 learner_drain_samples_ms = []
                 learner_breakdown_samples_ms = {}
                 learner_ring_wait_samples_ms = []
@@ -392,7 +413,10 @@ def _drain_stats():
                     checkpoints.TRAINING_STATE,
                     checkpoint_format=checkpoint_format,
                 )
-                (console.print if console else print)(f"[motrix.fastsac async] saved checkpoint {path}")
+                if console is not None:
+                    last_checkpoint_path = str(path)
+                else:
+                    print(f"[motrix.fastsac async] saved checkpoint {path}")
                 next_save += save_interval
 
         # final checkpoint (identical structure to sync fastsac)

motrix_rl/src/motrix_rl/fastsac/sync/train.py

@@ -20,7 +20,7 @@
 from motrix_rl.fastsac.wrap_np import FastSacNpEnvWrap
 from motrix_rl.fastsac.wrap_torch import FastSacTorchEnvWrap
 from motrix_rl.frameworks import TrainerBase, TrainerContext
-from motrix_rl.system_metrics import CpuLoadSampler
+from motrix_rl.system_metrics import CpuLoadSampler, GpuMemoryUsageSampler, GpuUtilizationSampler, MemoryUsageSampler
 
 # Enable TF32 matmul on Ampere+ GPUs. SAC training has no precision concern with
 # TF32 (10 mantissa bits), and the speedup is meaningful when AMP is off.
@@ -215,6 +215,10 @@ def _run_loop(
         t_collect = 0.0
         t_learn = 0.0
         cpu_sampler = CpuLoadSampler()
+        gpu_sampler = GpuUtilizationSampler()
+        memory_sampler = MemoryUsageSampler()
+        gpu_memory_sampler = GpuMemoryUsageSampler()
+        last_checkpoint_path: str | None = None
 
         # optional live console that refreshes one panel in place
         console, live = open_training_live()
@@ -321,8 +325,12 @@ def emit_msg(msg: str) -> None:
                         reward_terms=term_means,
                         env_metrics=env_metrics,
                         cpu_load=cpu_sampler.sample(),
+                        gpu_utilization_percent=gpu_sampler.sample(),
+                        memory_usage=memory_sampler.sample(),
+                        gpu_memory_usage=gpu_memory_sampler.sample(),
+                        checkpoint_path=last_checkpoint_path,
                     )
-                    emit_training_panel(live, stats, title="motrix.fastsac (sync)")
+                    emit_training_panel(live, stats, title=f"{self._env_name}/motrix.fastsac")
                     if self._writer is not None:
                         self._writer.add_scalar("rollout/mean_return", mean_ret, agent.global_step)
                         self._writer.add_scalar("rollout/mean_ep_len", mean_len, agent.global_step)
@@ -347,9 +355,10 @@ def emit_msg(msg: str) -> None:
                         path = Path(self._context.checkpoint_dir) / f"model_{agent.global_step:07d}.pt"
                         torch.save(agent.state_dict(), path)
                         record_checkpoint(path)
-                        emit_msg(
-                            f"[green]✓[/] saved checkpoint [dim]{path}[/]" if console else f"saved checkpoint {path}"
-                        )
+                        if console is not None:
+                            last_checkpoint_path = str(path)
+                        else:
+                            emit_msg(f"saved checkpoint {path}")
 
                 local += 1
                 agent.global_step += 1

motrix_rl/src/motrix_rl/system_metrics.py

@@ -1,13 +1,22 @@
 # Copyright Motphys Technology Co., Ltd. 2025, 2026
 # SPDX-License-Identifier: Apache-2.0
 
-"""Low-overhead host metrics sampled at training-panel refresh boundaries."""
+"""Low-overhead host metrics sampled at training-panel refresh boundaries.
+
+CPU samplers read Linux ``/proc`` interfaces and return ``None`` where they
+are unavailable, so panels degrade to ``n/a`` fields; memory sampling also
+supports Windows via ``GlobalMemoryStatusEx``. The GPU samplers use NVML,
+which works on any platform with an NVIDIA driver.
+"""
 
 from __future__ import annotations
 
+import ctypes
 import os
+import sys
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Any
 
 
 @dataclass(frozen=True)
@@ -121,3 +130,116 @@ def _read_physical_core_count(self) -> int | None:
         except (OSError, ValueError):
             return None
         return len(cores)
+
+
+@dataclass(frozen=True)
+class MemoryUsage:
+    """Memory usage in bytes for a host or accelerator device."""
+
+    used_bytes: int
+    total_bytes: int
+
+
+class _MemoryStatusEx(ctypes.Structure):
+    """``MEMORYSTATUSEX`` layout for the Windows ``GlobalMemoryStatusEx`` call."""
+
+    _fields_ = [
+        ("dwLength", ctypes.c_ulong),
+        ("dwMemoryLoad", ctypes.c_ulong),
+        ("ullTotalPhys", ctypes.c_ulonglong),
+        ("ullAvailPhys", ctypes.c_ulonglong),
+        ("ullTotalPageFile", ctypes.c_ulonglong),
+        ("ullAvailPageFile", ctypes.c_ulonglong),
+        ("ullTotalVirtual", ctypes.c_ulonglong),
+        ("ullAvailVirtual", ctypes.c_ulonglong),
+        ("ullAvailExtendedVirtual", ctypes.c_ulonglong),
+    ]
+
+
+def _windows_memory_status() -> MemoryUsage | None:
+    """Physical memory usage via ``GlobalMemoryStatusEx``, mirroring the /proc semantics."""
+    status = _MemoryStatusEx()
+    status.dwLength = ctypes.sizeof(_MemoryStatusEx)
+    if not ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)):
+        return None
+    return MemoryUsage(used_bytes=status.ullTotalPhys - status.ullAvailPhys, total_bytes=status.ullTotalPhys)
+
+
+class MemoryUsageSampler:
+    """Read host memory usage from Linux ``/proc/meminfo`` or the Windows memory API."""
+
+    def __init__(self, *, meminfo_path: str | Path = "/proc/meminfo") -> None:
+        self._meminfo_path = Path(meminfo_path)
+
+    def sample(self) -> MemoryUsage | None:
+        if sys.platform == "win32":
+            return _windows_memory_status()
+        try:
+            values: dict[str, int] = {}
+            for line in self._meminfo_path.read_text().splitlines():
+                key, value, *_ = line.split()
+                if key in {"MemTotal:", "MemAvailable:"}:
+                    values[key] = int(value) * 1024
+        except (OSError, ValueError):
+            return None
+        total = values.get("MemTotal:")
+        available = values.get("MemAvailable:")
+        if total is None or available is None:
+            return None
+        return MemoryUsage(used_bytes=max(0, total - available), total_bytes=total)
+
+
+# NVML reads the same counters nvidia-smi reports, but in-process at
+# microsecond cost instead of a subprocess spawn per query.
+_nvml_state: tuple[Any, list[Any]] | tuple[()] | None = None  # None: untried; (): unavailable
+
+
+def _nvml() -> tuple[Any, list[Any]] | None:
+    """Lazily initialize NVML and return ``(module, device_handles)``, or ``None``."""
+    global _nvml_state
+    if _nvml_state is None:
+        try:
+            import pynvml
+
+            pynvml.nvmlInit()
+            devices = [pynvml.nvmlDeviceGetHandleByIndex(index) for index in range(pynvml.nvmlDeviceGetCount())]
+            _nvml_state = (pynvml, devices)
+        except Exception:  # ImportError (pynvml missing) or NVML init failure (no driver/GPU)
+            _nvml_state = ()
+    return _nvml_state or None
+
+
+class GpuMemoryUsageSampler:
+    """Read aggregate NVIDIA memory usage across all visible GPUs via NVML."""
+
+    def sample(self) -> MemoryUsage | None:
+        session = _nvml()
+        if session is None:
+            return None
+        pynvml, devices = session
+        used = total = 0
+        try:
+            for device in devices:
+                info = pynvml.nvmlDeviceGetMemoryInfo(device)
+                used += info.used
+                total += info.total
+        except pynvml.NVMLError:
+            return None
+        return MemoryUsage(used_bytes=used, total_bytes=total) if total > 0 else None
+
+
+class GpuUtilizationSampler:
+    """Read aggregate NVIDIA GPU utilization across all visible GPUs via NVML."""
+
+    def sample(self) -> float | None:
+        session = _nvml()
+        if session is None:
+            return None
+        pynvml, devices = session
+        values: list[int] = []
+        try:
+            for device in devices:
+                values.append(pynvml.nvmlDeviceGetUtilizationRates(device).gpu)
+        except pynvml.NVMLError:
+            return None
+        return sum(values) / len(values) if values else None