unilab.training.experiment

Shared experiment tracking utilities for local files and W&B.

Functions

build_run_config_payload(*, run_metadata, ...)

Assemble the canonical {"run", "config", ...} run-config payload.

build_wandb_run_name(algo_name, task_name, ...)

build_wandb_settings(training_cfg, *, ...)

get_git_info(root_dir)

patch_rsl_rl_action_std_logging(runner)

Log the current distribution's public standard deviation, including state dependence.

patch_rsl_rl_resume_state()

Persist and restore UniLab progress through rsl-rl checkpoints.

patch_rsl_rl_tensorboard_logging(runner[, ...])

Batch and optionally throttle rsl-rl's TensorBoard scalar writes.

patch_rsl_rl_wandb_writer()

Patch rsl-rl W&B writer so it can reuse an already-open run.

write_run_config_snapshot(log_dir, *, ...[, ...])

Write the run-config snapshot into log_dir and return the payload.

Classes

ExperimentTracker

Tracks experiment metadata locally and optionally in Weights & Biases.
unilab.training.experiment.get_git_info(root_dir)[source]
Parameters:

root_dir (str | Path)

Return type:

dict[str, Any]

unilab.training.experiment.build_run_config_payload(*, run_metadata, full_cfg, contract_snapshot=None)[source]

Assemble the canonical {"run", "config", ...} run-config payload.

Parameters:
Return type:

dict[str, Any]

unilab.training.experiment.write_run_config_snapshot(log_dir, *, run_metadata, full_cfg, filename='run_config.json', contract_snapshot=None, trailing_newline=False)[source]

Write the run-config snapshot into log_dir and return the payload.

Parameters:
Return type:

dict[str, Any]

unilab.training.experiment.build_wandb_run_name(algo_name, task_name, log_dir)[source]
Parameters:
Return type:

str

unilab.training.experiment.build_wandb_settings(training_cfg, *, algo_name, task_name, sim_backend, log_dir)[source]
Parameters:
Return type:

dict[str, Any]

class unilab.training.experiment.ExperimentTracker[source]

Bases: object

Tracks experiment metadata locally and optionally in Weights & Biases.

Parameters:
__init__(*, root_dir, log_dir, algo_name, task_name, sim_backend, training_cfg, full_cfg, device=None, configured_device=None, collector_device=None, device_resolution=None, seed_info=None, tf32_matmul_policy=None)[source]
Parameters:
property run: Any | None
property run_url: str | None
property wandb_settings: dict[str, Any]
start()[source]
Return type:

None

update_summary(summary=None)[source]
Parameters:

summary (dict[str, Any] | None)

Return type:

None

log_video(video_path, key='media/play_video')[source]
Parameters:
Return type:

None

finish()[source]
Return type:

None

unilab.training.experiment.patch_rsl_rl_action_std_logging(runner)[source]

Log the current distribution’s public standard deviation, including state dependence.

Parameters:

runner (Any)

Return type:

None

unilab.training.experiment.patch_rsl_rl_tensorboard_logging(runner, log_interval=1)[source]

Batch and optionally throttle rsl-rl’s TensorBoard scalar writes.

Wraps runner.logger.writer so each step’s ~15-30 add_scalar calls become a single event record, and gates backend writes to every log_interval iterations (console output and episode bookkeeping are unaffected; the final iteration is always logged). No-op for non-TB backends and when no writer exists.

Parameters:

  • runner (Any)

  • log_interval (int)

Return type:

None

unilab.training.experiment.patch_rsl_rl_wandb_writer()[source]

Patch rsl-rl W&B writer so it can reuse an already-open run.

Return type:

None

unilab.training.experiment.patch_rsl_rl_resume_state()[source]

Persist and restore UniLab progress through rsl-rl checkpoints.

Without this patch, rsl-rl’s Logger.__init__ writes tot_time = 0 and tot_timesteps = 0 and OnPolicyRunner.load never refreshes them, so the Train/mean_reward/time and Train/mean_episode_length/time TensorBoard scalars (which use int(self.tot_time) as their step) restart from 0 on every resumed run and visually overlap the original segment. See issue #441.

The patch wraps OnPolicyRunner.save / OnPolicyRunner.load to round-trip plain progress in the saved dict: :rtype: None

  • unilab_logger_state: Logger.tot_time / tot_timesteps (required);

  • unilab_env_training_state: the wrapped TorchEnv’s versioned training state, written only when the environment publishes one.

Loading fails closed (issue #2416). A checkpoint must carry a valid unilab_logger_state; pre-patch rsl-rl checkpoints without it are rejected with a message describing the minimum supported format instead of silently restarting progress from 0. When the active environment publishes training state, the checkpoint must also carry its versioned unilab_env_training_state; environments without exportable state neither write nor require the key, and no placeholder value is fabricated.

Checkpoints are immutable load inputs. The loader never repairs, rewrites, upgrades, or migrates a checkpoint file. Environment state is plain versioned JSON data; no Torch tensor, backend, environment, or manager object is serialized through this path.