feat: use rich logging (colors) with hydra (Metta-AI#208)

## Changes - Add support for rich colored logging with millisecond precision - Create utility to override Hydra's default logger with our rich logger - Add time remaining estimation to epoch progress logs (shows in sec/min/hours/days) ## Why Hydra sets up its own logger by default, but our rich logger provides better visibility with colors and precise timestamps. The time remaining estimation helps developers monitor long-running training jobs more effectively. fixes Metta-AI#125

Author: rjwalters

metta/rl/pufferlib/trainer.py

@@ -109,10 +109,11 @@ def __init__(
  if self._master:
  print(policy_record.policy())
 
- if policy_record.metadata["action_names"] != self.vecenv.driver_env.action_names():
+ action_names = self.vecenv.driver_env.action_names()
+ if policy_record.metadata["action_names"] != action_names:
  raise ValueError(
  "Action names do not match between policy and environment: "
- f"{policy_record.metadata['action_names']} != {self.vecenv.driver_env.action_names()}"
+ f"{policy_record.metadata['action_names']} != {action_names}"
  )
 
  self._initial_pr = policy_record
@@ -175,13 +176,15 @@ def __init__(
 
  def train(self):
  self.train_start = time.time()
+ self.steps_start = self.agent_step
+
  logger.info("Starting training")
 
+ # it doesn't make sense to evaluate more often than checkpointing since we need a saved policy to evaluate
  if (
  self.trainer_cfg.evaluate_interval != 0
  and self.trainer_cfg.evaluate_interval < self.trainer_cfg.checkpoint_interval
  ):
- # it doesn't make sense to evaluate more often than checkpointing since we need a saved policy to evaluate
  raise ValueError("evaluate_interval must be at least as large as checkpoint_interval")
 
  logger.info(f"Training on {self.device}")
@@ -195,9 +198,25 @@ def train(self):
  # Processing stats
  self._process_stats()
 
+ # log progress
+ steps_per_second = (self.agent_step - self.steps_start) / (time.time() - self.train_start)
+ remaining_steps = self.trainer_cfg.total_timesteps - self.agent_step
+ remaining_time_sec = remaining_steps / steps_per_second
+
+ # Format remaining time in appropriate units
+ if remaining_time_sec < 60:
+ time_str = f"{remaining_time_sec:.0f} sec"
+ elif remaining_time_sec < 3600:
+ time_str = f"{remaining_time_sec / 60:.1f} min"
+ elif remaining_time_sec < 86400: # Less than a day
+ time_str = f"{remaining_time_sec / 3600:.1f} hours"
+ else:
+ time_str = f"{remaining_time_sec / 86400:.1f} days"
+
  logger.info(
- f"Epoch {self.epoch} - {self.agent_step} "
- f"({100.00 * self.agent_step / self.trainer_cfg.total_timesteps:.2f}%)"
+ f"Epoch {self.epoch} - {self.agent_step} [{steps_per_second:.0f}/sec]"
+ f" ({100.00 * self.agent_step / self.trainer_cfg.total_timesteps:.2f}%)"
+ f" - {time_str} remaining"
  )
 
  # Checkpointing trainer

metta/sim/simulation_config.py

@@ -9,6 +9,7 @@
 
 class SimulationConfig(Config):
  """Configuration for a single simulation run."""
+
  env: str
  device: str
  num_envs: int
@@ -24,26 +25,23 @@ class SimulationConfig(Config):
 
 class SimulationSuiteConfig(SimulationConfig):
  """A suite of named simulations, with suite-level defaults injected."""
+
  run_dir: str
  simulations: Dict[str, SimulationConfig]
 
- # —— don't need env bc all the simulations will specify —— 
+ # —— don't need env bc all the simulations will specify ——
  env: Optional[str] = None
 
  @model_validator(mode="before")
- def _propagate_defaults(cls, values: dict) -> dict:
+ def propagate_defaults(cls, values: dict) -> dict:
  # collect any suite-level overrides that are present & non-None
  suite_defaults = {
- k: v for k, v in values.items()
- if k in ("env", "device", "num_envs", "num_episodes") and v is not None
+ k: v for k, v in values.items() if k in ("env", "device", "num_envs", "num_episodes") and v is not None
  }
-
  raw_sims = values.get("simulations", {}) or {}
  merged: Dict[str, dict] = {}
  for name, sim_cfg in raw_sims.items():
  # sim_cfg is a dict; override only where sim_cfg provides a key
  merged[name] = {**suite_defaults, **sim_cfg}
  values["simulations"] = merged
  return values
-
-

metta/util/logging.py

@@ -1,7 +1,10 @@
+import logging
 import os
 import sys
+from datetime import datetime
 
 from loguru import logger
+from rich.logging import RichHandler
 
 
 def remap_io(logs_path: str):
@@ -12,14 +15,72 @@ def remap_io(logs_path: str):
  stderr = open(stderr_log_path, "a")
  sys.stderr = stderr
  sys.stdout = stdout
- logger.remove() # Remove default handler
- logger.remove() # Remove default handler
- # logger.add(
- # sys.stdout, colorize=True,
- # format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level: <8}</level> | "
- # "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - "
- # "<level>{message}</level>")
+ logger.remove()
 
 
 def restore_io():
  sys.stderr = sys.__stderr__
+ sys.stdout = sys.__stdout__
+
+
+# Create a custom formatter that supports milliseconds
+class MillisecondFormatter(logging.Formatter):
+ def formatTime(self, record, datefmt=None):
+ created = datetime.fromtimestamp(record.created)
+ # Convert microseconds to milliseconds (keep only 3 digits)
+ msec = created.microsecond // 1000
+ if datefmt:
+ # Replace %f with just 3 digits for milliseconds
+ datefmt = datefmt.replace("%f", f"{msec:03d}")
+ else:
+ datefmt = "[%H:%M:%S.%03d]"
+ return created.strftime(datefmt) % msec
+
+
+# Create a custom handler that always shows the timestamp
+class AlwaysShowTimeRichHandler(RichHandler):
+ def emit(self, record):
+ # Force a unique timestamp for each record
+ record.created = record.created + (record.relativeCreated % 1000) / 1000000
+ super().emit(record)
+
+
+def get_log_level(provided_level=None):
+ """
+ Determine log level based on priority:
+ 1. Environment variable LOG_LEVEL
+ 2. Provided level parameter
+ 3. Default to INFO
+ """
+ # Check environment variable first
+ env_level = os.environ.get("LOG_LEVEL")
+ if env_level:
+ return env_level.upper()
+
+ # Check provided level next
+ if provided_level:
+ return provided_level.upper()
+
+ # Default to INFO
+ return "INFO"
+
+
+def setup_mettagrid_logger(name: str, level=None) -> logging.Logger:
+ # Get the appropriate log level based on priority
+ log_level = get_log_level(level)
+
+ # Remove all handlers from the root logger
+ root_logger = logging.getLogger()
+ for handler in root_logger.handlers[:]:
+ root_logger.removeHandler(handler)
+
+ # Add back our custom Rich handler
+ rich_handler = AlwaysShowTimeRichHandler(rich_tracebacks=True)
+ formatter = MillisecondFormatter("%(message)s", datefmt="[%H:%M:%S.%f]")
+ rich_handler.setFormatter(formatter)
+ root_logger.addHandler(rich_handler)
+
+ # Set the level
+ root_logger.setLevel(getattr(logging, log_level))
+
+ return logging.getLogger(name)

tests/sim/test_simulation_config.py

@@ -1,11 +1,10 @@
 """
-Unit‑tests for SimulationSuiteConfig ⇄ SimulationConfig behaviour.
-
+Unit‑tests for SimulationSuiteConfig ⇄ SimulationConfig behavior.
 Covered
 -------
 * suite‑level defaults propagate into children
 * child‑level overrides win
-* missing required keys always raise  (allow_missing removed)
+* missing required keys always raise (allow_missing removed)
 """
 
 from typing import Dict
@@ -19,21 +18,29 @@
 # ---------------------------------------------------------------------------
 # constants
 # ---------------------------------------------------------------------------
-
 ROOT_ENV, CHILD_A, CHILD_B = "env/root", "env/a", "env/b"
 DEVICE, RUN_DIR = "cpu", "./runs/test"
 
 
-def _build(cfg: Dict):
- return SimulationSuiteConfig(OmegaConf.create(cfg))
+@pytest.fixture
+def build_simulation_suite_config():
+ def _build(cfg: Dict):
+ # First create the OmegaConf object
+ dict_config = OmegaConf.create(cfg)
+
+ # Convert to a Python dictionary
+ regular_dict = OmegaConf.to_container(dict_config, resolve=True)
+
+ # Now create the SimulationSuiteConfig using the model_validate method
+ return SimulationSuiteConfig.model_validate(regular_dict)
+
+ return _build
 
 
 # ---------------------------------------------------------------------------
 # propagation & overrides
 # ---------------------------------------------------------------------------
-
-
-def test_propogate_defaults_and_overrides():
+def test_propagate_defaults_and_overrides(build_simulation_suite_config):
  cfg = {
  "env": ROOT_ENV,
  "num_envs": 4,
@@ -45,9 +52,8 @@ def test_propogate_defaults_and_overrides():
  "b": {"env": CHILD_B, "num_envs": 8}, # overrides num_envs
  },
  }
- suite = _build(cfg)
+ suite = build_simulation_suite_config(cfg)
  a, b = suite.simulations["a"], suite.simulations["b"]
-
  # device and num_envs both propagated, even though num_envs has a default
  assert (a.device, a.num_envs) == (DEVICE, 4)
  assert (b.device, b.num_envs) == (DEVICE, 8)
@@ -56,20 +62,17 @@ def test_propogate_defaults_and_overrides():
 # ---------------------------------------------------------------------------
 # allow_extra – child nodes
 # ---------------------------------------------------------------------------
-
-
 @pytest.mark.parametrize(
  "has_extra, should_pass",
  [
  (False, True),
  (True, False),
  ],
 )
-def test_allow_extra_child_keys(has_extra, should_pass):
+def test_allow_extra_child_keys(build_simulation_suite_config, has_extra, should_pass):
  child_node = {"env": CHILD_A}
  if has_extra:
  child_node["foo"] = "bar" # <- unknown key
-
  cfg = {
  "env": ROOT_ENV,
  "num_envs": 4,
@@ -78,21 +81,18 @@ def test_allow_extra_child_keys(has_extra, should_pass):
  "run_dir": RUN_DIR,
  "simulations": {"sim": child_node},
  }
-
  if should_pass:
- suite = _build(cfg)
+ suite = build_simulation_suite_config(cfg)
  assert suite.simulations["sim"].device == DEVICE
  else:
  with pytest.raises(ValueError):
- _build(cfg)
+ build_simulation_suite_config(cfg)
 
 
 # ---------------------------------------------------------------------------
 # missing required keys should always error
 # ---------------------------------------------------------------------------
-
-
-def test_missing_device_always_errors():
+def test_missing_device_always_errors(build_simulation_suite_config):
  cfg = {
  "env": ROOT_ENV,
  "num_envs": 4,
@@ -101,10 +101,10 @@ def test_missing_device_always_errors():
  "simulations": {"sim": {}}, # required 'device' omitted
  }
  with pytest.raises(ValidationError):
- _build(cfg)
+ build_simulation_suite_config(cfg)
 
 
-def test_missing_suite_env_is_allowed():
+def test_missing_suite_env_is_allowed(build_simulation_suite_config):
  cfg = {
  "run_dir": RUN_DIR,
  "device": DEVICE,
@@ -116,5 +116,5 @@ def test_missing_suite_env_is_allowed():
  }
  },
  }
- suite = _build(cfg)
+ suite = build_simulation_suite_config(cfg)
  assert suite.simulations["sim"].env == CHILD_A

tools/analyze.py

@@ -1,18 +1,18 @@
 """Analysis tool for MettaGrid evaluation results."""
 
-import logging
-
 import hydra
 from omegaconf import DictConfig
 
 from metta.eval.report import dump_stats, generate_report
+from metta.util.logging import setup_mettagrid_logger
 from metta.util.runtime_configuration import setup_mettagrid_environment
 
 
 @hydra.main(version_base=None, config_path="../configs", config_name="analyze_job")
 def main(cfg: DictConfig) -> None:
  setup_mettagrid_environment(cfg)
- logger = logging.getLogger(__name__)
+ logger = setup_mettagrid_logger("analyze")
+
  view_type = "latest"
  logger.info(f"Generating {view_type} report")
  dump_stats(cfg)

tools/play.py

@@ -1,4 +1,3 @@
-import logging
 import os
 import signal # Aggressively exit on ctrl+c
 import sys
@@ -9,31 +8,37 @@
 from metta.agent.policy_store import PolicyStore
 from metta.sim.simulation_config import SimulationConfig
 from metta.util.config import Config
+from metta.util.logging import setup_mettagrid_logger
 from metta.util.runtime_configuration import setup_mettagrid_environment
 from metta.util.wandb.wandb_context import WandbContext
 
 signal.signal(signal.SIGINT, lambda sig, frame: os._exit(0))
 
 
-logging.basicConfig(level="INFO")
-logger = logging.getLogger(__name__)
-
-
 class PlayJob(Config):
  sim: SimulationConfig
  policy_uri: str
 
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
 
 @hydra.main(version_base=None, config_path="../configs", config_name="play_job")
-def play(cfg):
+def main(cfg) -> int:
  setup_mettagrid_environment(cfg)
 
+ logger = setup_mettagrid_logger("play")
+ logger.info(f"Playing {cfg.run}")
+
  with WandbContext(cfg) as wandb_run:
  policy_store = PolicyStore(cfg, wandb_run)
+
  play_job = PlayJob(cfg.play_job)
  policy_record = policy_store.policy(play_job.policy_uri)
  metta.sim.simulator.play(play_job.sim, policy_record)
 
+ return 0
+
 
 if __name__ == "__main__":
- sys.exit(play())
+ sys.exit(main())