Language Language English 简体中文

unilab.training

Shared training helpers for entrypoint scripts.

The re-exports below keep the owner-facing unilab.training import path clear. New code should import from canonical owner modules instead (unilab.training.*, unilab.base.config_adapter, and unisim.backend.base).

class unilab.training.EnvConfigAdapter

Bases: object

Build env/play overrides from the final composed config.

Parameters:

• cfg (DictConfig)

• root_dir (str | Path)

• algo_name (str | None)

• scene_materializer (Callable[..., str])

__init__(cfg, *, root_dir, algo_name=None, scene_materializer=<function materialize_scene_visual_override>)

Parameters:

• cfg (DictConfig)

• root_dir (str | Path)

• algo_name (str | None)

• scene_materializer (Callable[..., str])

build_task_env_cfg_override()

Build env_cfg_override from the resolved reward + env sections.

Return type:

dict[str, Any]

build_play_env_cfg_override()

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

Return type:

dict[str, Any]

class unilab.training.ExperimentTracker

Bases: object

Tracks experiment metadata locally and optionally in Weights & Biases.

Parameters:

• root_dir (str | Path)

• log_dir (str | Path)

• algo_name (str)

• task_name (str)

• sim_backend (str)

• training_cfg (Any)

• full_cfg (Any)

• device (str | None)

• configured_device (str | None)

• collector_device (str | None)

• device_resolution (Mapping[str, Mapping[str, str]] | None)

• seed_info (Any | None)

• tf32_matmul_policy (Any | None)

__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)

Parameters:

• root_dir (str | Path)

• log_dir (str | Path)

• algo_name (str)

• task_name (str)

• sim_backend (str)

• training_cfg (Any)

• full_cfg (Any)

• device (str | None)

• configured_device (str | None)

• collector_device (str | None)

• device_resolution (Mapping[str, Mapping[str, str]] | None)

• seed_info (Any | None)

• tf32_matmul_policy (Any | None)

property run: Any | None

property run_url: str | None

property wandb_settings: dict[str, Any]

start()

Return type:

None

update_summary(summary=None)

Parameters:

summary (dict[str, Any] | None)

Return type:

None

log_video(video_path, key='media/play_video')

Parameters:

• video_path (str | Path | None)

• key (str)

Return type:

None

finish()

Return type:

None

unilab.training.algo_config_dict(cfg)

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)

Attach a NanGuard to env when diagnostics are enabled.

Parameters:

• env (Any)

• training_cfg (Any)

Return type:

None

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

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)

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)

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

Parameters:

• cfg (DictConfig)

• num_envs (int)

• env_cfg_override (dict[str, Any] | None)

• sim_backend (str | None)

• task_name (str | None)

unilab.training.ensure_registries()

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)

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

Parameters:

• cfg (DictConfig)

• task_log_root (Path)

• load_path (Path | None)

• load_path_dir (Path | None)

Return type:

str

unilab.training.get_hydra_runtime_choice(cfg, group)

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')

Return the latest model checkpoint inside a run directory.

Parameters:

• run_dir (str | Path)

• suffix (str)

Return type:

Path | None

unilab.training.get_latest_run(log_dir)

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=None, *, algo_log_name=None, log_root=None)

Resolve the algorithm log root.

The configured training.log_root (or the explicit log_root argument) wins; otherwise the UNILAB_TEST_LOG_ROOT override applies; otherwise the default is <root_dir>/logs/<algo_log_name>. Pass either a composed Hydra cfg or an explicit algo_log_name.

Parameters:

• root_dir (str | Path)

• cfg (DictConfig | None)

• algo_log_name (str | None)

• log_root (str | Path | None)

Return type:

Path

unilab.training.graceful_keyboard_interrupt(entrypoint)

Let entrypoint cleanup finish, then exit Ctrl+C without a traceback.

Runners own their finally blocks. This outer boundary converts only the resulting KeyboardInterrupt into the conventional shell status after all collector, logger, and distributed-process cleanup has completed.

Parameters:

entrypoint (TypeVar(_F, bound= Callable[..., Any]))

Return type:

TypeVar(_F, bound= Callable[..., Any])

unilab.training.is_viser_play_render_mode(play_render_mode)

Return whether the configured mode selects the browser-based viser viewer.

viser is a UniLab-level playback mode layered on top of the unisim PLAY_RENDER_MODES contract; it renders through the shared MuJoCo playback shell instead of a backend-native renderer, so it never reaches SimBackend.resolve_play_render_plan.

Parameters:

play_render_mode (str | None)

Return type:

bool

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

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

Parameters:

• plan (PlayRenderPlan)

• prefix (str)

Return type:

None

unilab.training.nonfatal_play_step(step)

Keep independent post-training play artifacts (ONNX export, video render) from blocking each other: log the failure and continue with the rest.

Parameters:

step (str)

Return type:

Iterator[None]

unilab.training.resolve_checkpoint_path(task_log_root, load_run, *, checkpoint=None, suffix='.pt')

Resolve a checkpoint from a task log root.

load_run selects the run: "-1" picks the latest run under task_log_root; any other value is first tried as a direct path and then relative to task_log_root, and may name either a run directory or a checkpoint file. checkpoint optionally selects one checkpoint inside the resolved run (an iteration number or a filename).

Returns (checkpoint_path, run_dir). A resolved run that lacks the requested (or any) checkpoint yields (None, run_dir); an unresolvable run yields (None, None).

Parameters:

• task_log_root (str | Path)

• load_run (str | int | PathLike[str])

• checkpoint (str | int | None)

• suffix (str)

Return type:

tuple[Path | None, Path | None]

unilab.training.resolve_configured_checkpoint_path(cfg, *, root_dir, load_run=None, task_name=None, checkpoint=None, suffix='.pt')

Resolve a checkpoint path from Hydra config and repository root.

Parameters:

• cfg (DictConfig)

• root_dir (str | Path)

• load_run (str | int | PathLike[str] | None)

• task_name (str | None)

• checkpoint (str | int | None)

• suffix (str)

Return type:

tuple[Path | None, Path | None]

unilab.training.resolve_nan_guard_cfg(training_cfg)

Build the shared tensor diagnostics config, or None.

Parameters:

training_cfg (Any)

Return type:

NanGuardCfg | None

unilab.training.resolve_run_log_dir(cfg, *, world_size=1, timestamp=None, root_dir=None)

Resolve one canonical run directory shared by all distributed ranks.

Parameters:

• cfg (DictConfig)

• world_size (int)

• timestamp (str | None)

• root_dir (str | Path | None)

Return type:

str

unilab.training.should_run_playback(*, play_only, no_play, play_render_mode)

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

Parameters:

• play_only (bool)

• no_play (bool)

• play_render_mode (str | None)

Return type:

bool

class unilab.training.TrainingSeedInfo

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)

configured_seed: int | None

configured_seed_source: str | None

effective_seed: int | None

to_dict()

Return type:

dict[str, int | str | None]

__init__(configured_seed, configured_seed_source, effective_seed)

Parameters:

• configured_seed (int | None)

• configured_seed_source (str | None)

• effective_seed (int | None)

unilab.training.apply_configured_training_seed(cfg, *, torch_runtime=True, cuda=True)

Resolve and apply the configured training seed before runtime construction.

Parameters:

• cfg (Any)

• torch_runtime (bool)

• cuda (bool)

Return type:

TrainingSeedInfo

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

Apply a seed to the runtimes used by training entrypoints.

Parameters:

• seed (int | None)

• torch_runtime (bool)

• cuda (bool)

Return type:

int | None

unilab.training.resolve_training_seed(cfg)

Resolve the configured seed from the algorithm-level algo.seed contract.

Parameters:

cfg (Any)

Return type:

TrainingSeedInfo

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

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

Parameters:

• log_dir (str | Path)

• algo_name (str)

• echo (bool)

• filename (str)

Return type:

Logger

Modules

common	Shared helpers for training entrypoints.
cuda_mps_cli	User-owned CUDA MPS daemon management for the UniLab CLI.
cuda_process_sharing	Fail-closed validation for CUDA process-sharing execution modes.
dp_config	Hydra-side data-parallel configuration assembly.
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.
seed	Shared training seed contract helpers.
soak	Resource and progress monitoring for long-duration training soaks.
tensor_diagnostics	Finite-state diagnostics for the tensor-native environment runtime.
tensor_runtime	Hydra config resolution for the off-policy tensor-runtime contract.
tensor_runtime_hooks	Pickleable runtime hooks injected across the unilab.rl spawn boundary.