opentau.utils.logging_utils

Utilities for tracking and logging training metrics.

This module provides classes for tracking metrics during training, including AverageMeter for computing running averages and MetricsTracker for managing multiple metrics with step tracking.

Classes

AverageMeter(name[, fmt])	Computes and stores the average and current value.
MetricsTracker(batch_size, metrics[, ...])	A helper class to track and log metrics over time.

class opentau.utils.logging_utils.AverageMeter(name: str, fmt: str = ':f')

Bases: object

Computes and stores the average and current value.

Adapted from https://github.com/pytorch/examples/blob/main/imagenet/main.py

Parameters:

• name – Name of the metric being tracked.

• fmt – Format string for displaying the average value. Defaults to “:f”.

__init__(name: str, fmt: str = ':f')

reset() → None

Reset all accumulated statistics to zero.

update(val: float, n: int = 1) → None

Update the meter with a new value.

Parameters:

• val – New value to add.

• n – Number of samples this value represents. Defaults to 1.

class opentau.utils.logging_utils.MetricsTracker(batch_size: int, metrics: dict[str, AverageMeter], initial_step: int = 0)

Bases: object

A helper class to track and log metrics over time.

Usage pattern:

# initialize, potentially with non-zero initial step (e.g. if resuming run)
metrics = {"loss": AverageMeter("loss", ":.3f")}
train_metrics = MetricsTracker(cfg, dataset, metrics, initial_step=step)

# update metrics derived from step (samples, episodes, epochs) at each training step
train_metrics.step()

# update various metrics
loss = policy.forward(batch)
train_metrics.loss = loss

# display current metrics
logging.info(train_metrics)

# export for wandb
wandb.log(train_metrics.to_dict())

# reset averages after logging
train_metrics.reset_averages()

__init__(batch_size: int, metrics: dict[str, AverageMeter], initial_step: int = 0)

Initialize the metrics tracker.

Parameters:

• batch_size – Number of samples per gradient update.

• metrics – Dictionary of metric names to AverageMeter instances.

• initial_step – Starting step number (useful when resuming training). Defaults to 0.

reset_averages() → None

Resets average meters.

step() → None

Updates metrics that depend on ‘step’ for one step.

to_dict(use_avg: bool = True) → dict[str, int | float]

Convert current metrics to a dictionary.

Parameters:

use_avg – If True, use average values; otherwise use current values. Defaults to True.

Returns:

Dictionary containing steps, samples, and all metric values.