unilab.training

Shared training helpers for entrypoint scripts.

Deprecated compatibility layer: the re-exports below are kept for existing callers only. New code should import from the canonical locations instead (unilab.utils.* for layer-0 helpers, unilab.base.config_adapter / unisim.backend.base for base-layer helpers).

class unilab.training.BackendAdapter[source]

Bases: object

Build env/play overrides from the final composed config.

Parameters:
__init__(cfg, *, root_dir, algo_name=None, scene_materializer=<function materialize_scene_visual_override>)[source]
Parameters:
build_task_env_cfg_override()[source]

Build env_cfg_override from the resolved reward + env sections.

Return type:

dict[str, Any]

build_play_env_cfg_override()[source]

Build play-mode overrides from an optional backend-agnostic play profile.

Return type:

dict[str, Any]

class unilab.training.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, collector_device=None, seed_info=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

class unilab.training.HardwareMonitor[source]

Bases: object

Monitor CPU, GPU, memory usage.

__init__()[source]
get_metrics()[source]

Get current hardware metrics.

Return type:

Dict[str, float]

unilab.training.algo_config_dict(cfg)[source]

Resolve the composed cfg.algo subtree into a plain mutable dict.

Parameters:

cfg (DictConfig)

Return type:

dict[str, Any]

unilab.training.apply_env_nan_guard(env, training_cfg)[source]

Attach a NanGuard to env when training.nan_guard is enabled.

Parameters:

  • env (Any)

  • training_cfg (Any)

Return type:

None

unilab.training.assert_offpolicy_task_choice_matches_algo(cfg, *, algo_name=None)[source]

Reject off-policy configs whose cfg.algo.algo does not match the requested algo.

Parameters:

  • cfg (DictConfig)

  • algo_name (str | None)

Return type:

None

unilab.training.build_run_dir_name(timestamp, sim_backend, *, world_size=1)[source]

Return the canonical run directory name shared by all training entries.

Parameters:

  • timestamp (str)

  • sim_backend (str)

  • world_size (int)

Return type:

str

unilab.training.create_env(cfg, *, num_envs, env_cfg_override=None, sim_backend=None, task_name=None)[source]

Construct an environment via the registry using the current Hydra config.

Parameters:
unilab.training.ensure_registries()[source]

Import env modules so registry-based entrypoints can instantiate tasks.

Return type:

None

unilab.training.format_play_checkpoint_error(cfg, *, task_log_root, load_path, load_path_dir)[source]

Build the user-facing diagnostic for an unresolvable play checkpoint.

Parameters:
Return type:

str

unilab.training.get_entrypoint_log_root(root_dir, *, algo_log_name, log_root=None)[source]

Resolve the log root for non-Hydra entrypoints using training helper semantics.

Parameters:
Return type:

Path

unilab.training.get_hydra_runtime_choice(cfg, group)[source]

Return a selected Hydra config-group choice when runtime metadata is available.

Parameters:

  • cfg (DictConfig)

  • group (str)

Return type:

str | None

unilab.training.get_latest_checkpoint(run_dir, *, suffix='.pt')[source]

Return the latest model checkpoint inside a run directory.

Parameters:
Return type:

Path | None

unilab.training.get_latest_run(log_dir)[source]

Return the lexicographically latest run directory under a task log root.

Parameters:

log_dir (str | Path)

Return type:

Path | None

unilab.training.get_log_root(root_dir, cfg)[source]

Resolve the algorithm log root, honoring optional training.log_root overrides.

Parameters:

  • root_dir (str | Path)

  • cfg (DictConfig)

Return type:

Path

unilab.training.log_playback_plan(plan, *, prefix='')[source]

Print user-facing playback status for a resolved backend plan.

Parameters:

  • plan (BackendPlayRenderPlan)

  • prefix (str)

Return type:

None

unilab.training.parse_checkpoint_path(cfg, *, root_dir, load_run=None, task_name=None, checkpoint=None, suffix='.pt')[source]

Resolve a checkpoint path from Hydra config and repository root.

Parameters:
Return type:

tuple[Path | None, Path | None]

unilab.training.resolve_checkpoint_path(base_log_dir, load_run, *, suffix='.pt')[source]

Resolve a latest or explicit checkpoint path from a task log root.

Parameters:
Return type:

tuple[Path | None, Path | None]

unilab.training.resolve_nan_guard_cfg(training_cfg)[source]

Build the shared NanGuardCfg from training.nan_guard, or None.

Parameters:

training_cfg (Any)

Return type:

NanGuardCfg | None

unilab.training.resolve_task_checkpoint_path(root_dir, *, task_name, load_run, algo_log_name, checkpoint=None, suffix='.pt', log_root=None)[source]

Resolve checkpoint paths for auxiliary entrypoints through shared training semantics.

Parameters:
Return type:

tuple[Path | None, Path | None]

unilab.training.should_run_playback(*, play_only, no_play, play_render_mode)[source]

Return whether train/eval should enter playback for the configured mode.

Parameters:
Return type:

bool

class unilab.training.TrainingSeedInfo[source]

Bases: object

Configured and effective seed metadata for a training run.

Parameters:
configured_seed: int | None
configured_seed_source: str | None
effective_seed: int | None
to_dict()[source]
Return type:

dict[str, int | str | None]

__init__(configured_seed, configured_seed_source, effective_seed)
Parameters:
unilab.training.apply_configured_training_seed(cfg, *, torch_runtime=True, cuda=True)[source]

Resolve and apply the configured training seed before runtime construction.

Parameters:
Return type:

TrainingSeedInfo

unilab.training.apply_training_seed(seed, *, torch_runtime=True, cuda=True)[source]

Apply a seed to the runtimes used by training entrypoints.

Parameters:
Return type:

int | None

unilab.training.derive_worker_seed(base_seed, worker_index=0)[source]

Derive deterministic subprocess seeds from the effective run seed.

Parameters:
Return type:

int | None

unilab.training.resolve_appo_checkpoint_path(base_log_dir, load_run)[source]

Resolve an APPO checkpoint under a task log root, returning string paths.

Parameters:
Return type:

tuple[str | None, str | None]

unilab.training.resolve_offpolicy_checkpoint_path(root_dir, algo_log_name, task, load_run)[source]

Resolve an off-policy checkpoint from the repo-rooted log tree.

Parameters:
Return type:

tuple[str | None, str | None]

unilab.training.resolve_training_seed(cfg)[source]

Resolve the configured seed, preferring the algorithm-level contract.

Parameters:

cfg (Any)

Return type:

TrainingSeedInfo

unilab.training.setup_logger(log_dir, algo_name, *, echo=True, filename='train.log')[source]

Create a simple file-backed logger for script-local progress messages.

Parameters:
Return type:

Logger

Modules

common

Shared helpers for training entrypoints.

experiment

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

onnx_export

ONNX export and onnxruntime numeric verification helpers for play entrypoints.

run

Run directory and checkpoint resolution helpers.