Add MetricsManager for custom metric logging

[kevinzakka]

Summary

Adds a MetricsManager so users can log custom per-step metrics during training without hacking reward functions or adding zero-weight reward terms. Closes #584.

• New: MetricsManager, MetricsTermCfg, NullMetricsManager in managers/metrics_manager.py
• Integration: wired into ManagerBasedRlEnv config, step(), and _reset_idx()
• Terms use the same callable signature as rewards (env, **params) → Tensor[num_envs]
• No weight, no dt scaling — metrics are observational, not reward signals
• Episode values are true per-step averages (sum / step_count), so a metric in [0,1] stays in [0,1] in wandb
• Empty config means zero overhead (NullMetricsManager)

Example usage

  from mjlab.managers import MetricsTermCfg

  def joint_velocity_magnitude(env, asset_cfg):
    """L1 norm of joint velocities."""
    return env.scene[asset_cfg.name].data.joint_vel.abs().sum(dim=-1)  # (num_envs,)

  @dataclass(kw_only=True)
  class MyEnvCfg(ManagerBasedRlEnvCfg):
    metrics = {
      "joint_vel_mag": MetricsTermCfg(
        func=joint_velocity_magnitude,
        params={"asset_cfg": SceneEntityCfg(name="robot")},
      ),
    }

On episode reset, Episode_Metrics/joint_vel_mag appears in extras["log"] and flows to wandb/tensorboard automatically.

Test plan

• uv run pytest tests/test_metrics_manager.py — 6 targeted tests
• uv run pytest tests/test_rewards.py — no regression
• uv run ty check / uv run pyright — clean
• uv run ruff check && uv run ruff format — clean

🤖 Generated with Claude Code

[brentyi]

+100, seems super useful!

I'm wondering if any if these things are possible in a simple/not overengineered way:

Does it make sense to allow customization of how metrics are "reduced" between steps and before logging? In the extreme case it'd be nice, for example, to be able to specify logging for std of metrics, stds of per-episode means, histograms, histograms of per-episode stds, etc.

Can we implement any of the existing things that are logged (eg rewards) as default terms in the metrics manager?

[kevinzakka]

Thanks for the review @brentyi!

Does it make sense to allow customization of how metrics are "reduced" between steps and before logging?

The rsl_rl logger only supports scalars. Everything in extras["log"] goes through torch.mean() then add_scalar() (logger.py:171). So histograms/distributions aren't possible without changing the logger. We could try submitting a PR to rsl_rl. In the meantime, for std or other reductions, users can just add a second metric term that computes it directly (e.g. joint_vel_std alongside joint_vel_mean). Both log as scalars and it works today with no extra machinery.

Can we implement any of the existing things that are logged (eg rewards) as default terms in the metrics manager?

In principle yes. Note however that rewards use dt-scaled sums divided by max_episode_length_s while metrics use raw sums divided by step_count. We'd have to add machinery to the metric manager to support this.

[brentyi]

In the meantime, for std or other reductions, users can just add a second metric term that computes it directly (e.g. joint_vel_std alongside joint_vel_mean). Both log as scalars and it works today with no extra machinery.

Makes sense!

To check my understanding: we wouldn't be able to reuse intermediates right? And we could compute std across episodes within a single timestep but not across timesteps within an episode? These are a bit annoying but fine.

Note however that rewards use dt-scaled sums divided by max_episode_length_s while metrics use raw sums divided by step_count.

Makes sense. It seems kind of nice to consolidate logic but I don't feel strongly about this.