44tensorboard via ``env.extras["log"]["Episode_Metrics/<name>"]``. Useful for surface
55indicators (action smoothness, sensor health, etc.) where a weight-zero reward
66term is the wrong tool because it bakes the metric into the reward bookkeeping.
7+
8+ mjlab logs the corresponding diagnostics as ``Metrics/<name>_mean`` directly from
9+ inside the reward functions, then rsl-rl averages them across a rollout. GeneLab's
10+ :class:`~genelab.managers.metrics_manager.MetricsManager` instead averages over
11+ an episode and reports at env reset under ``Episode_Metrics/<name>``. The
12+ per-step scalar is the same; the time horizon (rollout vs episode) differs by
13+ a constant factor for diagnostic comparison.
714"""
815
916from typing import TYPE_CHECKING
1017
1118import torch
1219
20+ from genelab .managers .reward_manager import RewardTermCfg
21+ from genelab .managers .scene_entity_cfg import SceneEntityCfg
22+ from genelab .mdp ._helpers import (
23+ contact_sensor as _contact_sensor ,
24+ link_ids as _link_ids ,
25+ site_lin_vel_w as _site_lin_vel_w ,
26+ site_pos_w as _site_pos_w ,
27+ )
28+
1329if TYPE_CHECKING :
1430 from genelab .envs .manager_based_rl_env import ManagerBasedRlEnv
1531
@@ -29,3 +45,127 @@ def mean_action_acc(env: "ManagerBasedRlEnv") -> torch.Tensor:
2945 a_prev_prev = env .action_manager .prev_prev_action
3046 accel = a - 2.0 * a_prev + a_prev_prev
3147 return torch .mean (torch .abs (accel ), dim = - 1 )
48+
49+
50+ # --------------------------------------------------------------------- gait diagnostics
51+ #
52+ # The five metrics below mirror mjlab's ``Metrics/{air_time,peak_height,
53+ # slip_velocity,landing_force,angular_momentum}_mean`` writes (see
54+ # ``mjlab/tasks/velocity/mdp/rewards.py:205, 229, 317, 352, 373``). Each
55+ # computes mjlab's exact cross-(env, foot) scalar each step, then broadcasts
56+ # to ``(B,)`` so :class:`MetricsManager` can accumulate it uniformly across
57+ # envs — the broadcast collapses cleanly at reset because every env sees the
58+ # same value.
59+
60+
61+ def _broadcast_scalar (value : torch .Tensor , env : "ManagerBasedRlEnv" ) -> torch .Tensor :
62+ """Expand a 0-dim tensor to ``(B,)`` so MetricsManager can accumulate per-env."""
63+ return value .expand (env .num_envs )
64+
65+
66+ def angular_momentum_mean (env : "ManagerBasedRlEnv" , sensor_name : str ) -> torch .Tensor :
67+ """``||L||`` of the root-frame angular momentum.
68+
69+ mjlab logs ``torch.mean(angmom_magnitude)`` (a cross-env scalar) inside the
70+ ``angular_momentum_penalty`` reward; per-env ``||L||`` averaged over the
71+ episode gives the same per-batch value when MetricsManager reduces across
72+ envs at reset.
73+ """
74+ angmom = env .sensors [sensor_name ].data
75+ return torch .norm (angmom , dim = - 1 )
76+
77+
78+ def air_time_mean (env : "ManagerBasedRlEnv" , sensor_name : str ) -> torch .Tensor :
79+ """Mean ``current_air_time`` across (env, foot) pairs that are in the air.
80+
81+ Replicates mjlab's ``Metrics/air_time_mean`` in ``feet_air_time``: sum of
82+ air times divided by number of mid-air feet (clamped to 1). Broadcasted to
83+ ``(B,)`` for MetricsManager.
84+ """
85+ data = _contact_sensor (env , sensor_name ).data
86+ current_air_time = data .current_air_time
87+ in_air = current_air_time > 0
88+ in_air_f = in_air .float ()
89+ num_in_air = in_air_f .sum ().clamp (min = 1.0 )
90+ mean = (current_air_time * in_air_f ).sum () / num_in_air
91+ return _broadcast_scalar (mean , env )
92+
93+
94+ def slip_velocity_mean (
95+ env : "ManagerBasedRlEnv" , sensor_name : str , asset_cfg : SceneEntityCfg
96+ ) -> torch .Tensor :
97+ """Mean ``||v_xy||`` across grounded (env, foot) pairs.
98+
99+ Replicates mjlab's ``Metrics/slip_velocity_mean`` in ``feet_slip``. Reads
100+ site-frame foot velocity (``link_offsets`` honoured) and weights by
101+ in-contact bool. Broadcasted to ``(B,)``.
102+ """
103+ indices = list (_link_ids (asset_cfg ))
104+ data = _contact_sensor (env , sensor_name ).data
105+ in_contact = data .found .float () # (B, F)
106+ foot_vel_xy = _site_lin_vel_w (env , indices , asset_cfg .link_offsets_tensor )[..., :2 ]
107+ vel_norm = torch .norm (foot_vel_xy , dim = - 1 ) # (B, F)
108+ num_in_contact = in_contact .sum ().clamp (min = 1.0 )
109+ mean = (vel_norm * in_contact ).sum () / num_in_contact
110+ return _broadcast_scalar (mean , env )
111+
112+
113+ def landing_force_mean (env : "ManagerBasedRlEnv" , sensor_name : str ) -> torch .Tensor :
114+ """Mean ``|F|`` across (env, foot) pairs that just landed.
115+
116+ Replicates mjlab's ``Metrics/landing_force_mean`` in ``soft_landing``:
117+ contact-force magnitude on the substep ``first_contact`` fires, averaged
118+ over the landings (clamped to 1). Broadcasted to ``(B,)``.
119+ """
120+ data = _contact_sensor (env , sensor_name ).data
121+ first_contact = data .first_contact .float ()
122+ landing_impact = data .force_norm * first_contact
123+ num_landings = first_contact .sum ().clamp (min = 1.0 )
124+ mean = landing_impact .sum () / num_landings
125+ return _broadcast_scalar (mean , env )
126+
127+
128+ class peak_height_mean :
129+ """Mean swing-apex height across (env, foot) pairs at landing.
130+
131+ Stateful: tracks per-foot peak height during the air phase, samples the
132+ peak on the substep ``first_contact`` fires, then resets the slot for the
133+ next swing. Mirrors :class:`~genelab.mdp.rewards.feet_swing_height`'s
134+ bookkeeping (separate buffer so the metric works whether the reward is
135+ weight=0 or active — currently weight=−0.25 in G1).
136+
137+ Replicates mjlab's ``Metrics/peak_height_mean`` in ``feet_swing_height``
138+ (``mjlab/.../rewards.py:317``). Broadcasted to ``(B,)`` for MetricsManager.
139+ """
140+
141+ def __init__ (self , cfg : RewardTermCfg , env : "ManagerBasedRlEnv" ) -> None :
142+ asset_cfg : SceneEntityCfg = cfg .params ["asset_cfg" ]
143+ indices = list (_link_ids (asset_cfg ))
144+ self ._foot_indices : list [int ] = indices
145+ self ._offsets_tensor : torch .Tensor | None = asset_cfg .link_offsets_tensor
146+ self ._peak_heights = torch .zeros (env .num_envs , len (indices ), device = env .device )
147+
148+ def __call__ (
149+ self ,
150+ env : "ManagerBasedRlEnv" ,
151+ sensor_name : str ,
152+ asset_cfg : SceneEntityCfg ,
153+ ) -> torch .Tensor :
154+ del asset_cfg # consumed at __init__
155+ data = _contact_sensor (env , sensor_name ).data
156+ foot_z = _site_pos_w (env , self ._foot_indices , self ._offsets_tensor )[..., 2 ]
157+ # Accumulate the peak while airborne.
158+ in_air = ~ data .found
159+ self ._peak_heights = torch .where (
160+ in_air , torch .maximum (self ._peak_heights , foot_z ), self ._peak_heights
161+ )
162+ first_contact = data .first_contact .float ()
163+ peak_at_landing = self ._peak_heights * first_contact
164+ num_landings = first_contact .sum ().clamp (min = 1.0 )
165+ mean = peak_at_landing .sum () / num_landings
166+ # Zero the peak buffer for feet that just landed so the next swing
167+ # starts fresh.
168+ self ._peak_heights = torch .where (
169+ data .first_contact , torch .zeros_like (self ._peak_heights ), self ._peak_heights
170+ )
171+ return _broadcast_scalar (mean , env )
0 commit comments