"""Metrics computation and logging utilities for PPO training."""
from typing import Any, Optional, Union
from collections.abc import Mapping
import warnings
import jax
import jax.numpy as jp
from flax import nnx
from jaxtyping import Array, Float, PyTree
from nnx_ppo.networks.types import StatefulModule
from nnx_ppo.algorithms.types import LoggingLevel
from nnx_ppo.algorithms.rollout import Transition
[docs]
def compute_metrics(
loss_metrics: dict[str, PyTree],
rollout_data: Transition,
logging_level: LoggingLevel,
percentile_levels: Optional[tuple[int, ...]] = None,
) -> dict[str, Any]:
"""Compute training metrics from loss metrics and rollout data.
Args:
loss_metrics: Dictionary of loss values from ppo_loss.
rollout_data: Transition data from the rollout.
logging_level: Which metrics to include.
percentile_levels: Percentiles to compute (e.g., (0, 25, 50, 75, 100)).
If None, uses mean/std instead.
Returns:
Dictionary of computed metrics.
"""
metrics = {}
for k, v in loss_metrics.items():
_log_metric(metrics, k, v, percentile_levels)
if LoggingLevel.ENV_METRICS in logging_level:
_log_metric(metrics, "env", rollout_data.metrics["env"], percentile_levels)
if LoggingLevel.NETWORK_METRICS in logging_level:
_log_metric(metrics, "net", rollout_data.metrics["net"], percentile_levels)
if LoggingLevel.ROLLOUT_STATS in logging_level:
_log_metric(
metrics, "rollout_batch/reward", rollout_data.rewards, percentile_levels
)
_log_metric(
metrics,
"rollout_batch/action",
rollout_data.network_output.actions,
percentile_levels,
)
metrics["rollout_batch/done_rate"] = rollout_data.done.mean()
metrics["rollout_batch/truncation_rate"] = rollout_data.truncated.mean()
if LoggingLevel.ROLLOUT_OBS in logging_level:
_log_metric(
metrics, "rollout_batch/obs", rollout_data.obs, percentile_levels
)
if LoggingLevel.ACTOR_EXTRA in logging_level:
_log_metric(
metrics,
"loglikelihood",
rollout_data.network_output.loglikelihoods,
percentile_levels,
)
if LoggingLevel.CRITIC_EXTRA in logging_level:
_log_metric(
metrics,
"losses/predicted_value",
rollout_data.network_output.value_estimates,
percentile_levels,
)
return metrics
def _log_metric(
metrics: dict[str, Any],
name: str,
x: Union[Mapping[Union[str, int], Any], Float[Array, "..."]],
percentile_levels: Optional[tuple[int, ...]] = None,
) -> None:
"""Log a metric with either percentiles or mean/std.
Args:
metrics: Dictionary to add metrics to (mutated in place).
name: Base name for the metric.
x: Value to log (can be array or nested mapping).
percentile_levels: Percentiles to compute. If None, uses mean/std.
"""
if isinstance(x, Mapping):
for k, v in x.items():
_log_metric(metrics, f"{name}/{k}", v, percentile_levels)
return
# Boolean arrays (e.g. termination flags): log fraction-true, not mean/std/percentiles
if hasattr(x, "dtype") and jp.issubdtype(x.dtype, jp.bool_):
metrics[name] = jp.mean(x)
elif percentile_levels is None or len(percentile_levels) == 0:
metrics[f"{name}/mean"] = jp.mean(x)
metrics[f"{name}/std"] = jp.std(x)
else:
percentiles = jp.percentile(x, jp.array(percentile_levels))
for pl, p in zip(percentile_levels, percentiles):
metrics[f"{name}/p{int(pl)}"] = p
[docs]
def log_weight_stats(
metrics: dict[str, Any],
networks: StatefulModule,
percentile_levels: Optional[tuple[int, ...]] = None,
) -> None:
"""Log weight statistics over all of the network's parameters.
Args:
metrics: Dictionary to add metrics to (mutated in place).
networks: The PPO network to extract weights from.
percentile_levels: Percentiles to compute. If None, uses mean/std.
"""
params = nnx.state(networks, nnx.Param)
leaves = jax.tree.leaves(params)
if not leaves:
warnings.warn("Network has no nnx.Param leaves; skipping weight logging.")
return None
weights = jp.concatenate([p.flatten() for p in leaves])
_log_metric(metrics, "weights", weights, percentile_levels)