purge some uninteresting stats from wandb logging (#4546)

# Summary
Reduce training/environment stats volume and simplify StatsReporter
output. This trims low‑value metrics, removes weight/dormant‑neuron
reporting, and narrows rolling averages to a small set of required env
metrics.

# Why
We’re intentionally cutting noisy or redundant metrics to lower logging
overhead and make dashboards easier to read, while keeping core signals
intact.

# What changed
## Environment stats production (mettagrid)
- Removed label completion tracking/emission from `StatsTracker`.
- Reward estimate diffs and env timing stats are still produced in the
env but are now filtered out before logging.

## Stats filtering (metta/rl/stats.py)
- `filter_movement_metrics` now drops:
  - `env_attributes/*`
  - `env_reward_estimates/*`
  - `env_timing_per_epoch/*`
  - `env_timing_cumulative/*`
- Removed stale `env_label_completions/*` filter entry (metric no longer
produced in this branch).
- Core movement metrics kept: only the four direction counters.

## StatsReporter (metta/rl/training/stats_reporter.py)
- Removed weight metrics + dormant‑neuron analysis.
- Removed `parameters/*` payload (learning rate, epoch steps,
minibatches, schedulefree extras).
- Hyperparameters now come only from optimizer param groups (lr,
schedulefree scheduled_lr, lr_max).
- Rolling averages computed only for `default_zero_metrics` instead of
all env metrics.

## W&B logger (metta/rl/training/wandb_logger.py)
- Removed direct logging of `latest_losses_stats`; losses now only flow
through StatsReporter.
  - If StatsReporter is disabled, loss metrics will not appear.

## System monitor (mettagrid)
- Dropped static counters: cpu_count, cpu_count_logical/physical,
memory_total_mb, gpu_count.

# Behavior / compatibility notes
- W&B/env metric output is substantially reduced. Any dashboards/scripts
that reference removed metrics will need updates.
- Loss metrics are no longer emitted by WandbLogger when StatsReporter
is disabled.

# Testing
Not run (stats/logging changes only).

# Files touched
- metta/rl/stats.py
- metta/rl/training/stats_reporter.py
- metta/rl/training/wandb_logger.py
- packages/mettagrid/python/src/mettagrid/envs/stats_tracker.py
- packages/mettagrid/python/src/mettagrid/profiling/system_monitor.py
- tests/rl/test_stats_reporter_defaults.py


[Asana
Task](https://app.asana.com/1/1209016784099267/project/1210348820405981/task/1212600739220124)

Author: relh

metta/rl/stats.py

@@ -50,7 +50,7 @@ def accumulate_rollout_stats(
 
 
 def filter_movement_metrics(stats: dict[str, Any]) -> dict[str, Any]:
-    """Filter movement metrics to only keep core values, removing derived stats."""
+    """Filter verbose environment metrics while keeping core values."""
     filtered = {}
 
     # Core movement metrics we want to keep (without any suffix)
@@ -60,14 +60,18 @@ def filter_movement_metrics(stats: dict[str, Any]) -> dict[str, Any]:
         "env_agent/movement.direction.down",
         "env_agent/movement.direction.left",
         "env_agent/movement.direction.right",
-        "env_agent/movement.sequential_rotations",
-        "env_agent/movement.rotation.to_up",
-        "env_agent/movement.rotation.to_down",
-        "env_agent/movement.rotation.to_left",
-        "env_agent/movement.rotation.to_right",
     }
+    noisy_prefixes = (
+        "env_reward_estimates/",
+        "env_timing_per_epoch/",
+        "env_timing_cumulative/",
+    )
 
     for key, value in stats.items():
+        if key.startswith("env_attributes/"):
+            continue
+        if key.startswith(noisy_prefixes):
+            continue
         # Check if this is a core metric (exact match)
         if key in core_metrics:
             filtered[key] = value

metta/rl/training/stats_reporter.py

@@ -7,14 +7,11 @@
 
 import numpy as np
 import torch
-import torch.nn as nn
 from pydantic import Field
 
 from metta.common.wandb.context import WandbRun
-from metta.rl.model_analysis import compute_dormant_neuron_stats
 from metta.rl.stats import accumulate_rollout_stats, compute_timing_stats, process_training_stats
 from metta.rl.training.component import TrainerComponent
-from metta.rl.utils import should_run
 from mettagrid.base_config import Config
 
 logger = logging.getLogger(__name__)
@@ -43,11 +40,9 @@ def _to_scalar(value: Any) -> Optional[float]:
 def build_wandb_payload(
     processed_stats: dict[str, Any],
     timing_info: dict[str, Any],
-    weight_stats: dict[str, Any],
     grad_stats: dict[str, float],
     system_stats: dict[str, Any],
     memory_stats: dict[str, Any],
-    parameters: dict[str, Any],
     hyperparameters: dict[str, Any],
     *,
     agent_step: int,
@@ -87,12 +82,9 @@ def _update(items: dict[str, Any], *, prefix: str = "") -> None:
     _update(experience_stats, prefix="experience/")
 
     _update(processed_stats.get("environment_stats", {}))
-    _update(parameters, prefix="parameters/")
     _update(hyperparameters, prefix="hyperparameters/")
-
     _update(system_stats)
     _update({f"trainer_memory/{k}": v for k, v in memory_stats.items()})
-    _update(weight_stats)
     _update(grad_stats)
     _update(timing_info.get("timing_stats", {}))
 
@@ -107,10 +99,6 @@ class StatsReporterConfig(Config):
     grad_mean_variance_interval: int = 50
     interval: int = 1
     """How often to report stats (in epochs)"""
-    analyze_weights_interval: int = 0
-    """How often to compute weight metrics (0 disables)."""
-    dormant_neuron_threshold: float = 1e-6
-    """Threshold for considering a neuron dormant based on mean absolute weight magnitude."""
     rolling_window: int = Field(default=5, ge=1, description="Number of epochs for metric rolling averages")
     default_zero_metrics: tuple[str, ...] = Field(
         default_factory=lambda: ("env_game/assembler.heart.created",),
@@ -355,27 +343,16 @@ def _build_wandb_payload(
         timing_info = compute_timing_stats(timer=timer, agent_step=agent_step)
         self._normalize_steps_per_second(timing_info, agent_step)
 
-        weight_stats = self._collect_weight_stats(policy=policy, epoch=epoch)
-        dormant_stats = self._compute_dormant_neuron_stats(policy=policy)
-        if dormant_stats:
-            weight_stats.update(dormant_stats)
         system_stats = self._collect_system_stats()
         memory_stats = self._collect_memory_stats()
-        parameters = self._collect_parameters(
-            experience=experience,
-            optimizer=optimizer,
-            timing_info=timing_info,
-        )
-        hyperparameters = self._collect_hyperparameters(trainer_cfg=trainer_cfg, parameters=parameters)
+        hyperparameters = self._collect_hyperparameters(optimizer=optimizer)
 
         return build_wandb_payload(
             processed_stats=processed,
             timing_info=timing_info,
-            weight_stats=weight_stats,
             grad_stats=self._state.grad_stats,
             system_stats=system_stats,
             memory_stats=memory_stats,
-            parameters=parameters,
             hyperparameters=hyperparameters,
             agent_step=agent_step,
             epoch=epoch,
@@ -386,7 +363,10 @@ def _augment_with_rolling_averages(self, processed: dict[str, Any]) -> None:
         if not isinstance(env_stats, dict):
             return
 
-        tracked_keys = set(env_stats.keys()) | set(self._state.rolling_stats.keys())
+        tracked_keys = set(self._config.default_zero_metrics)
+        for key in list(self._state.rolling_stats.keys()):
+            if key not in tracked_keys:
+                del self._state.rolling_stats[key]
         window = self._config.rolling_window
 
         for key in tracked_keys:
@@ -402,9 +382,7 @@ def _augment_with_rolling_averages(self, processed: dict[str, Any]) -> None:
                 continue
             history.append(scalar)
             env_stats.setdefault(key, scalar)
-            # Skip creating .avg versions for env_per_label metrics
-            if not (key.startswith("env_per_label_rewards/") or key.startswith("env_per_label_chest_deposits/")):
-                env_stats[f"{key}.avg"] = sum(history) / len(history)
+            env_stats[f"{key}.avg"] = sum(history) / len(history)
 
     def _normalize_steps_per_second(self, timing_info: dict[str, Any], agent_step: int) -> None:
         """Adjust SPS to account for agent steps accumulated before a resume."""
@@ -433,43 +411,6 @@ def _normalize_steps_per_second(self, timing_info: dict[str, Any], agent_step: i
         if isinstance(timing_stats, dict):
             timing_stats["timing_cumulative/sps"] = sps
 
-    def _collect_weight_stats(self, *, policy: Any, epoch: int) -> dict[str, float]:
-        interval = self._config.analyze_weights_interval
-        if not interval:
-            policy_config = getattr(policy, "config", None)
-            interval = getattr(policy_config, "analyze_weights_interval", 0) if policy_config else 0
-
-        if not interval or not should_run(epoch, interval):
-            return {}
-
-        if not hasattr(policy, "compute_weight_metrics"):
-            return {}
-
-        weight_stats: dict[str, float] = {}
-        try:
-            for metrics in policy.compute_weight_metrics():
-                name = metrics.get("name", "unknown")
-                for key, value in metrics.items():
-                    if key == "name":
-                        continue
-                    scalar = _to_scalar(value)
-                    if scalar is None:
-                        continue
-                    weight_stats[f"weights/{key}/{name}"] = scalar
-        except Exception as exc:  # pragma: no cover - safeguard against model-specific failures
-            logger.warning("Failed to compute weight metrics: %s", exc, exc_info=True)
-        return weight_stats
-
-    def _compute_dormant_neuron_stats(self, *, policy: Any) -> dict[str, float]:
-        if not isinstance(policy, nn.Module):
-            return {}
-        threshold = getattr(self._config, "dormant_neuron_threshold", 1e-6)
-        try:
-            return compute_dormant_neuron_stats(policy, threshold=threshold)
-        except Exception as exc:  # pragma: no cover - safeguard against model-specific failures
-            logger.debug("Failed to compute dormant neuron stats: %s", exc, exc_info=True)
-            return {}
-
     def _collect_system_stats(self) -> dict[str, Any]:
         system_monitor = getattr(self.context, "system_monitor", None)
         if system_monitor is None:
@@ -490,79 +431,19 @@ def _collect_memory_stats(self) -> dict[str, Any]:
             logger.debug("Memory monitor stats failed: %s", exc, exc_info=True)
             return {}
 
-    def _collect_parameters(
-        self,
-        *,
-        experience: Any,
-        optimizer: torch.optim.Optimizer,
-        timing_info: dict[str, Any],
-    ) -> dict[str, Any]:
-        learning_rate = getattr(self.context.config.optimizer, "learning_rate", 0)
-        if optimizer and optimizer.param_groups:
-            learning_rate = optimizer.param_groups[0].get("lr", learning_rate)
-
-        parameters: dict[str, Any] = {
-            "learning_rate": learning_rate,
-            "epoch_steps": timing_info.get("epoch_steps", 0),
-            "num_minibatches": getattr(experience, "num_minibatches", 0),
-        }
-
-        # Add ScheduleFree optimizer information
-        if optimizer and optimizer.param_groups:
-            param_group = optimizer.param_groups[0]
-            is_schedulefree = "train_mode" in param_group
-
-            if is_schedulefree:
-                scheduled_lr = param_group.get("scheduled_lr")
-                if scheduled_lr is not None:
-                    parameters["schedulefree_scheduled_lr"] = scheduled_lr
-                lr_max = param_group.get("lr_max")
-                if lr_max is not None:
-                    parameters["schedulefree_lr_max"] = lr_max
-
-        return parameters
-
-    def _collect_hyperparameters(
-        self,
-        *,
-        trainer_cfg: Any,
-        parameters: dict[str, Any],
-    ) -> dict[str, Any]:
+    def _collect_hyperparameters(self, *, optimizer: torch.optim.Optimizer) -> dict[str, Any]:
         hyperparameters: dict[str, Any] = {}
-        if "learning_rate" in parameters:
-            hyperparameters["learning_rate"] = parameters["learning_rate"]
-
-        optimizer_cfg = getattr(trainer_cfg, "optimizer", None)
-        if optimizer_cfg:
-            hyperparameters["optimizer_type"] = optimizer_cfg.type
-            if "schedulefree" in optimizer_cfg.type:
-                warmup_steps = getattr(optimizer_cfg, "warmup_steps", None)
-                if warmup_steps is not None:
-                    hyperparameters["schedulefree_warmup_steps"] = warmup_steps
-
-        losses = getattr(trainer_cfg, "losses", None)
-        loss_configs = getattr(losses, "loss_configs", {}) if losses else {}
-        if isinstance(loss_configs, dict):
-            ppo_actor_cfg = loss_configs.get("ppo_actor")
-            if ppo_actor_cfg is not None:
-                for attr in ("clip_coef", "ent_coef", "norm_adv", "target_kl"):
-                    value = getattr(ppo_actor_cfg, attr, None)
-                    if value is None:
-                        continue
-                    hyperparameters[f"ppo_actor_{attr}"] = value
-
-            ppo_critic_cfg = loss_configs.get("ppo_critic")
-            if ppo_critic_cfg is not None:
-                for attr in (
-                    "vf_coef",
-                    "vf_clip_coef",
-                    "clip_vloss",
-                    "critic_update",
-                    "aux_coef",
-                    "beta",
-                ):
-                    value = getattr(ppo_critic_cfg, attr, None)
-                    if value is None:
-                        continue
-                    hyperparameters[f"ppo_critic_{attr}"] = value
+        param_groups = optimizer.param_groups
+        if not param_groups:
+            return hyperparameters
+        param_group = param_groups[0]
+        learning_rate = param_group.get("lr")
+        if learning_rate is not None:
+            hyperparameters["learning_rate"] = learning_rate
+        scheduled_lr = param_group.get("scheduled_lr")
+        if scheduled_lr is not None:
+            hyperparameters["schedulefree_scheduled_lr"] = scheduled_lr
+        lr_max = param_group.get("lr_max")
+        if lr_max is not None:
+            hyperparameters["schedulefree_lr_max"] = lr_max
         return hyperparameters

metta/rl/training/wandb_logger.py

@@ -54,10 +54,6 @@ def on_epoch_end(self, epoch: int) -> None:  # noqa: D401 - documented in base c
         for k, v in elapsed.items():
             self._prev_elapsed[k] = float(v)
 
-        for key, value in context.latest_losses_stats.items():
-            metric_key = key if "/" in key else f"loss/{key}"
-            payload[metric_key] = float(value)
-
         self._wandb_run.log(payload)
 
     def on_training_complete(self) -> None:  # noqa: D401

packages/mettagrid/python/src/mettagrid/envs/stats_tracker.py

@@ -15,7 +15,6 @@ def __init__(self, stats_writer: StatsWriter):
         self._stats_writer = stats_writer
         self._episode_start_ts = datetime.datetime.now()
         self._episode_end_ts = None
-        self._label_completions = {"completed_tasks": [], "completion_rates": {}}
         self._per_label_rewards = {}
         self._per_label_chest_deposits = {}  # Track chest deposits per label
 
@@ -55,11 +54,6 @@ def on_episode_end(self) -> None:
                 config.game.reward_estimates["worst_case_optimal_reward"] - mean_reward
             )
 
-        self._update_label_completions()
-
-        # only plot label completions once we have a full moving average window, to prevent initial bias
-        if len(self._label_completions["completed_tasks"]) >= 50:
-            infos["label_completions"] = self._label_completions["completion_rates"]
         self._per_label_rewards[config.label] = mean_reward
         infos["per_label_rewards"] = self._per_label_rewards
 
@@ -147,21 +141,3 @@ def _add_timing_info(self) -> None:
 
     def on_close(self) -> None:
         self._stats_writer.close()
-
-    def _update_label_completions(self, moving_avg_window: int = 500) -> None:
-        """Update label completions."""
-        label = self._sim.config.label
-
-        # keep track of a list of the last 500 labels
-        if len(self._label_completions["completed_tasks"]) >= moving_avg_window:
-            self._label_completions["completed_tasks"].pop(0)
-        self._label_completions["completed_tasks"].append(label)
-
-        # moving average of the completion rates
-        self._label_completions["completion_rates"] = {t: 0 for t in set(self._label_completions["completed_tasks"])}
-        for t in self._label_completions["completed_tasks"]:
-            self._label_completions["completion_rates"][t] += 1
-        self._label_completions["completion_rates"] = {
-            t: self._label_completions["completion_rates"][t] / len(self._label_completions["completed_tasks"])
-            for t in self._label_completions["completion_rates"]
-        }

packages/mettagrid/python/src/mettagrid/profiling/system_monitor.py

@@ -103,14 +103,10 @@ def _initialize_default_metrics(self):
         self._metric_collectors = {
             # CPU metrics
             "cpu_percent": lambda: psutil.cpu_percent(interval=0),
-            "cpu_count": lambda: psutil.cpu_count(),
-            "cpu_count_logical": lambda: psutil.cpu_count(logical=True),
-            "cpu_count_physical": lambda: psutil.cpu_count(logical=False) or psutil.cpu_count(logical=True),
             # Memory metrics
             "memory_percent": lambda: psutil.virtual_memory().percent,
             "memory_available_mb": lambda: psutil.virtual_memory().available / (1024 * 1024),
             "memory_used_mb": lambda: psutil.virtual_memory().used / (1024 * 1024),
-            "memory_total_mb": lambda: psutil.virtual_memory().total / (1024 * 1024),
             # Process-specific metrics
             "process_memory_mb": lambda: psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024),
             "process_cpu_percent": lambda: self._process.cpu_percent(),
@@ -154,7 +150,6 @@ def _initialize_default_metrics(self):
             # Add aggregate metrics (rename to make it clear they're aggregates)
             self._metric_collectors.update(
                 {
-                    "gpu_count": lambda: gpu_count,
                     "gpu_utilization_avg": self._get_gpu_utilization_cuda,
                     "gpu_memory_percent_avg": self._get_gpu_memory_percent_cuda,
                     "gpu_memory_used_mb_total": self._get_gpu_memory_used_mb_cuda,

tests/rl/test_stats_reporter_defaults.py

@@ -34,7 +34,7 @@ def _reporter(existing_heart: float | None = None) -> StatsReporter:
         stopwatch=timer,
         experience=SimpleNamespace(stats=lambda: {}),
         policy=None,
-        optimizer=SimpleNamespace(param_groups=[]),
+        optimizer=None,
         epoch=0,
         agent_step=0,
         run_name=None,
@@ -61,7 +61,7 @@ def test_heart_metric_zero_fill_and_preserve(existing: float | None, expected: f
         agent_step=0,
         epoch=0,
         timer=reporter.context.stopwatch,
-        optimizer=reporter.context.optimizer,
+        optimizer=SimpleNamespace(param_groups=[]),
     )
 
     assert payload["env_game/assembler.heart.created"] == expected