Add MetricsManager for logging custom metrics during training

Adds a MetricsManager so users can log custom per-step metrics without
hacking reward functions or adding zero-weight reward terms. Metrics
terms use the same callable signature as rewards (env, **params) but
have no weight, no dt scaling, and no normalization by episode length.
Episode values are true per-step averages (sum / step_count) logged
under "Episode_Metrics/{term_name}".

Closes #584

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: kevinzakka

src/mjlab/envs/manager_based_rl_env.py

@@ -22,6 +22,11 @@
   NullCurriculumManager,
 )
 from mjlab.managers.event_manager import EventManager, EventTermCfg
+from mjlab.managers.metrics_manager import (
+  MetricsManager,
+  MetricsTermCfg,
+  NullMetricsManager,
+)
 from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationManager
 from mjlab.managers.reward_manager import RewardManager, RewardTermCfg
 from mjlab.managers.termination_manager import TerminationManager, TerminationTermCfg
@@ -114,6 +119,9 @@ class ManagerBasedRlEnvCfg:
   curriculum: dict[str, CurriculumTermCfg] = field(default_factory=dict)
   """Curriculum terms for adaptive difficulty."""
 
+  metrics: dict[str, MetricsTermCfg] = field(default_factory=dict)
+  """Custom metric terms for logging per-step values as episode averages."""
+
   is_finite_horizon: bool = False
   """Whether the task has a finite or infinite horizon. Defaults to False (infinite).
 
@@ -291,6 +299,11 @@ def load_managers(self) -> None:
     else:
       self.curriculum_manager = NullCurriculumManager()
     print_info(f"[INFO] {self.curriculum_manager}")
+    if len(self.cfg.metrics) > 0:
+      self.metrics_manager = MetricsManager(self.cfg.metrics, self)
+    else:
+      self.metrics_manager = NullMetricsManager()
+    print_info(f"[INFO] {self.metrics_manager}")
 
     # Configure spaces for the environment.
     self._configure_gym_env_spaces()
@@ -367,6 +380,7 @@ def step(self, action: torch.Tensor) -> types.VecEnvStepReturn:
     self.reset_time_outs = self.termination_manager.time_outs
 
     self.reward_buf = self.reward_manager.compute(dt=self.step_dt)
+    self.metrics_manager.compute()
 
     # Reset envs that terminated/timed-out and log the episode info.
     reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1)
@@ -485,6 +499,9 @@ def _reset_idx(self, env_ids: torch.Tensor | None = None) -> None:
     # rewards manager.
     info = self.reward_manager.reset(env_ids)
     self.extras["log"].update(info)
+    # metrics manager.
+    info = self.metrics_manager.reset(env_ids)
+    self.extras["log"].update(info)
     # curriculum manager.
     info = self.curriculum_manager.reset(env_ids)
     self.extras["log"].update(info)

src/mjlab/managers/__init__.py

@@ -18,6 +18,9 @@
 from mjlab.managers.manager_base import ManagerBase as ManagerBase
 from mjlab.managers.manager_base import ManagerTermBase as ManagerTermBase
 from mjlab.managers.manager_base import ManagerTermBaseCfg as ManagerTermBaseCfg
+from mjlab.managers.metrics_manager import MetricsManager as MetricsManager
+from mjlab.managers.metrics_manager import MetricsTermCfg as MetricsTermCfg
+from mjlab.managers.metrics_manager import NullMetricsManager as NullMetricsManager
 from mjlab.managers.observation_manager import (
   ObservationGroupCfg as ObservationGroupCfg,
 )

src/mjlab/managers/metrics_manager.py

@@ -0,0 +1,144 @@
+"""Metrics manager for logging custom per-step metrics during training."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Sequence
+
+import torch
+from prettytable import PrettyTable
+
+from mjlab.managers.manager_base import ManagerBase, ManagerTermBaseCfg
+
+if TYPE_CHECKING:
+  from mjlab.envs.manager_based_rl_env import ManagerBasedRlEnv
+
+
+@dataclass(kw_only=True)
+class MetricsTermCfg(ManagerTermBaseCfg):
+  """Configuration for a metrics term."""
+
+  pass
+
+
+class MetricsManager(ManagerBase):
+  """Accumulates per-step metric values, reports episode averages.
+
+  Unlike rewards, metrics have no weight, no dt scaling, and no
+  normalization by episode length. Episode values are true per-step
+  averages (sum / step_count), so a metric in [0,1] stays in [0,1]
+  in the logger.
+  """
+
+  _env: ManagerBasedRlEnv
+
+  def __init__(self, cfg: dict[str, MetricsTermCfg], env: ManagerBasedRlEnv):
+    self._term_names: list[str] = list()
+    self._term_cfgs: list[MetricsTermCfg] = list()
+    self._class_term_cfgs: list[MetricsTermCfg] = list()
+
+    self.cfg = deepcopy(cfg)
+    super().__init__(env=env)
+
+    self._episode_sums: dict[str, torch.Tensor] = {}
+    for term_name in self._term_names:
+      self._episode_sums[term_name] = torch.zeros(
+        self.num_envs, dtype=torch.float, device=self.device
+      )
+    self._step_count = torch.zeros(self.num_envs, dtype=torch.long, device=self.device)
+    self._step_values = torch.zeros(
+      (self.num_envs, len(self._term_names)), dtype=torch.float, device=self.device
+    )
+
+  def __str__(self) -> str:
+    msg = f"<MetricsManager> contains {len(self._term_names)} active terms.\n"
+    table = PrettyTable()
+    table.title = "Active Metrics Terms"
+    table.field_names = ["Index", "Name"]
+    table.align["Name"] = "l"
+    for index, name in enumerate(self._term_names):
+      table.add_row([index, name])
+    msg += table.get_string()
+    msg += "\n"
+    return msg
+
+  # Properties.
+
+  @property
+  def active_terms(self) -> list[str]:
+    return self._term_names
+
+  # Methods.
+
+  def reset(
+    self, env_ids: torch.Tensor | slice | None = None
+  ) -> dict[str, torch.Tensor]:
+    if env_ids is None:
+      env_ids = slice(None)
+    extras = {}
+    counts = self._step_count[env_ids].float()
+    # Avoid division by zero for envs that haven't stepped.
+    safe_counts = torch.clamp(counts, min=1.0)
+    for key in self._episode_sums:
+      episode_avg = torch.mean(self._episode_sums[key][env_ids] / safe_counts)
+      extras["Episode_Metrics/" + key] = episode_avg
+      self._episode_sums[key][env_ids] = 0.0
+    self._step_count[env_ids] = 0
+    for term_cfg in self._class_term_cfgs:
+      term_cfg.func.reset(env_ids=env_ids)
+    return extras
+
+  def compute(self) -> None:
+    self._step_count += 1
+    for term_idx, (name, term_cfg) in enumerate(
+      zip(self._term_names, self._term_cfgs, strict=False)
+    ):
+      value = term_cfg.func(self._env, **term_cfg.params)
+      self._episode_sums[name] += value
+      self._step_values[:, term_idx] = value
+
+  def get_active_iterable_terms(
+    self, env_idx: int
+  ) -> Sequence[tuple[str, Sequence[float]]]:
+    terms = []
+    for idx, name in enumerate(self._term_names):
+      terms.append((name, [self._step_values[env_idx, idx].cpu().item()]))
+    return terms
+
+  def _prepare_terms(self):
+    for term_name, term_cfg in self.cfg.items():
+      term_cfg: MetricsTermCfg | None
+      if term_cfg is None:
+        print(f"term: {term_name} set to None, skipping...")
+        continue
+      self._resolve_common_term_cfg(term_name, term_cfg)
+      self._term_names.append(term_name)
+      self._term_cfgs.append(term_cfg)
+      if hasattr(term_cfg.func, "reset") and callable(term_cfg.func.reset):
+        self._class_term_cfgs.append(term_cfg)
+
+
+class NullMetricsManager:
+  """Placeholder for absent metrics manager that safely no-ops all operations."""
+
+  def __init__(self):
+    self.active_terms: list[str] = []
+    self.cfg = None
+
+  def __str__(self) -> str:
+    return "<NullMetricsManager> (inactive)"
+
+  def __repr__(self) -> str:
+    return "NullMetricsManager()"
+
+  def get_active_iterable_terms(
+    self, env_idx: int
+  ) -> Sequence[tuple[str, Sequence[float]]]:
+    return []
+
+  def reset(self, env_ids: torch.Tensor | None = None) -> dict[str, float]:
+    return {}
+
+  def compute(self) -> None:
+    pass

tests/test_metrics_manager.py

@@ -0,0 +1,124 @@
+"""Tests for metrics manager functionality."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from mjlab.managers.metrics_manager import (
+  MetricsManager,
+  MetricsTermCfg,
+  NullMetricsManager,
+)
+
+
+class SimpleTestMetric:
+  """A class-based metric that tracks state."""
+
+  def __init__(self, cfg: MetricsTermCfg, env):
+    self.call_count = torch.zeros(env.num_envs, device=env.device)
+
+  def __call__(self, env, **kwargs):
+    self.call_count += 1
+    return torch.ones(env.num_envs, device=env.device) * 0.5
+
+  def reset(self, env_ids: torch.Tensor | None = None, env=None):
+    if env_ids is not None and len(env_ids) > 0:
+      self.call_count[env_ids] = 0
+
+
+@pytest.fixture
+def mock_env():
+  env = Mock()
+  env.num_envs = 4
+  env.device = "cpu"
+  env.scene = {"robot": Mock()}
+  return env
+
+
+def test_episode_averages_and_reset(mock_env):
+  """Compute for N steps, reset a subset, verify averages and zeroing."""
+  cfg = {
+    "term": MetricsTermCfg(
+      func=lambda env: torch.ones(env.num_envs, device=env.device) * 0.5,
+      params={},
+    )
+  }
+  manager = MetricsManager(cfg, mock_env)
+
+  for _ in range(10):
+    manager.compute()
+
+  info = manager.reset(env_ids=torch.tensor([0, 1]))
+
+  # Each env: sum=5.0, count=10, avg=0.5. Mean across 2 reset envs = 0.5.
+  assert info["Episode_Metrics/term"].item() == pytest.approx(0.5)
+  # Reset envs zeroed; non-reset envs untouched.
+  assert manager._episode_sums["term"][0] == 0.0
+  assert manager._step_count[0] == 0
+  assert manager._episode_sums["term"][2] == pytest.approx(5.0)
+  assert manager._step_count[2] == 10
+
+
+def test_early_termination_uses_per_env_step_count(mock_env):
+  """Envs with different episode lengths get correct per-step averages."""
+  step = [0]
+
+  def step_dependent_metric(env):
+    step[0] += 1
+    return torch.full((env.num_envs,), float(step[0]), device=env.device)
+
+  cfg = {"m": MetricsTermCfg(func=step_dependent_metric, params={})}
+  manager = MetricsManager(cfg, mock_env)
+
+  # 4 steps for all envs: values are 1, 2, 3, 4.
+  for _ in range(4):
+    manager.compute()
+  # Env 0: sum=10, count=4. Reset it (env 1 keeps accumulating).
+  manager.reset(env_ids=torch.tensor([0]))
+
+  # 2 more steps: values are 5, 6.
+  for _ in range(2):
+    manager.compute()
+  # Env 0: sum=11, count=2, avg=5.5.
+  # Env 1: sum=21, count=6, avg=3.5.
+  info = manager.reset(env_ids=torch.tensor([0, 1]))
+  # Mean of [5.5, 3.5] = 4.5.
+  assert info["Episode_Metrics/m"].item() == pytest.approx(4.5)
+
+
+def test_class_based_metric_reset_targets_correct_envs(mock_env):
+  """Class-based term's reset() is called with the correct env_ids."""
+  cfg = {"term": MetricsTermCfg(func=SimpleTestMetric, params={})}
+  manager = MetricsManager(cfg, mock_env)
+  term = manager._class_term_cfgs[0].func
+
+  for _ in range(10):
+    manager.compute()
+
+  manager.reset(env_ids=torch.tensor([0, 2]))
+
+  assert term.call_count[0] == 0
+  assert term.call_count[1] == 10
+  assert term.call_count[2] == 0
+  assert term.call_count[3] == 10
+
+
+def test_null_metrics_manager(mock_env):
+  """NullMetricsManager doesn't crash and returns empty dict on reset."""
+  manager = NullMetricsManager()
+  manager.compute()
+  assert manager.reset(env_ids=torch.tensor([0])) == {}
+
+
+def test_none_terms_are_skipped(mock_env):
+  """None terms in config are skipped without error."""
+  cfg: dict[str, MetricsTermCfg | None] = {
+    "valid": MetricsTermCfg(
+      func=lambda env: torch.ones(env.num_envs, device=env.device),
+      params={},
+    ),
+    "skipped": None,
+  }
+  manager = MetricsManager(cfg, mock_env)  # type: ignore[arg-type]
+  assert manager._term_names == ["valid"]