Language

unilab.base.torch_env¶

Tensor-native environment lifecycle base.

TorchEnv is the sole Manager-Based runtime carrier. It owns the state, step, selected-row autoreset, timeout, finite, training-state, and backend playback contracts.

Classes

`TorchEnv`	Backend-agnostic Torch environment lifecycle.
`TorchEnvState`	Tensor-native vectorized environment transition.

class unilab.base.torch_env.TorchEnv[source]¶

Bases: ABEnv

Backend-agnostic Torch environment lifecycle.

Concrete owners implement action transformation, tensor state update, and selected-row (env_ids) reset. The base deliberately rejects NumPy input and never sanitizes non-finite tensors: invalid actions, controls, observations, and rewards fail closed.

Parameters:

cfg (EnvCfg)
backend (SimBackend)
num_envs (int)
device (device | str)

__init__(cfg, backend, num_envs, *, device='cpu')[source]¶

Parameters:

cfg (EnvCfg)
backend (SimBackend)
num_envs (int)
device (device | str)

property device: device¶

property cfg: EnvCfg¶: The configuration of the environment

property backend: SimBackend¶: Public backend owner used by Manager terms and tensor contracts.

property produces_completion_receipts: bool¶

RL env-contract declaration of the backend’s step receipts.

True when the bound backend publishes an event-backed async completion receipt on every step (replay CUDA ingress then fails closed without one). False for synchronous producers whose step outputs are already ordered at return, where no receipt exists and none is awaited. This is a cold-path declaration; step keeps publishing whatever receipt the backend returns.

property num_envs: int¶: return the size of the env if it is vectorized

property state: TorchEnvState | None¶: Current environment state (None before first reset)

property obs_groups_spec: dict[str, int]¶: Map from observation group name to its dimension.

property observation_space: Space¶: Observation space

init_state()[source]¶

Initialize environment and return initial state

Return type:: TorchEnvState

reset(env_ids=None)[source]¶

Reset the selected environment rows; None resets all rows.

The public selector follows the IsaacLab/mjlab env_ids row-index form; the v1.4 fixed-width bool mask form is removed.

Parameters:: env_ids (Tensor | None)
Return type:: tuple[dict[str, Tensor], dict[str, Any]]

step(actions)[source]¶

Step the environment with owner-declared action carrier.

The abstract base stays runtime-neutral. Concrete owners validate the declared carrier; the Torch lifecycle accepts Torch tensors.

Parameters:: actions (Tensor)
Return type:: TorchEnvState

init_play_renderer(render_spacing=None, render_offset_mode=None, *, headless=False, capture=False, width=1280, height=720, camera_kwargs=None)[source]¶

Initialize backend-native playback rendering when available.

Parameters:

render_spacing (float | None)
render_offset_mode (str | None)
headless (bool)
capture (bool)
width (int)
height (int)
camera_kwargs (CameraCfg | Mapping[str, Any] | None)

Return type:

None

resolve_play_render_plan(*, play_render_mode, play_steps, output_video)[source]¶

Resolve high-level playback mode through the concrete backend.

Parameters:

play_render_mode (str | None)
play_steps (int | None)
output_video (str | PathLike[str] | None)

Return type:

PlayRenderPlan

run_playback(*, initialize, step, num_steps, output_video=None, render_spacing=None, render_offset_mode=None, headless=None, record_video=None, frame_state_getter=None, camera_kwargs=None, debug_overlay_getter=None, on_frame=None)[source]¶

Execute playback through the concrete backend.

Parameters:

initialize (Callable[[], Any])
step (Callable[[Any], Any])
num_steps (int | None)
output_video (str | PathLike[str] | None)
render_spacing (float | None)
render_offset_mode (str | None)
headless (bool | None)
record_video (bool | None)
frame_state_getter (Callable[[], np.ndarray] | None)
camera_kwargs (CameraCfg | Mapping[str, Any] | None)
debug_overlay_getter (DebugOverlayGetter | None)
on_frame (Callable[[int, np.ndarray], np.ndarray | None] | None)

Return type:

str | None

property play_capabilities: EnvPlayCapabilities¶: Return env-facing play/render capabilities.

render(mode='rgb_array')[source]¶

Render the current state to an RGB array through the play renderer.

Parameters:: mode (str)
Return type:: ndarray

capture_play_video_frame()[source]¶

Capture one detached RGB video frame through the env contract.

Return type:: ndarray

get_physics_state_snapshot()[source]¶

Return a detached physics snapshot for playback/video export.

Return type:: ndarray

get_playback_model(env_index=None)[source]¶

Return the backend playback model for one vectorized environment.

Parameters:: env_index (int | None)
Return type:: Any

get_physics_state_layout()[source]¶

Return the backend physics-state playback layout.

Return type:: Any

get_playback_mocap_state(env_index=0)[source]¶

Return detached (mocap_pos, mocap_quat) playback state.

Parameters:: env_index (int)
Return type:: tuple[ndarray, ndarray]

get_scene_visual_model_file()[source]¶

Return the backend visual model file on the playback cold path.

Return type:: str | None

set_autoreset(enabled)[source]¶

Parameters:: enabled (bool)
Return type:: None

set_nan_guard(guard)[source]¶

Attach the owner-provided diagnostic guard.

Detection stays device-resident. Artifact export occurs only at the guard’s explicit abnormal-diagnostic host boundary.

Parameters:: guard (NanGuard)
Return type:: None

get_backend_runtime_diagnostics()[source]¶

Serialize the backend’s tensor runtime diagnostics on the cold path.

This is the collection boundary consumed by training runtimes (the unilab.rl off-policy collector publishes it into the run manifest). The payload is plain JSON: available and backend_type plus, when available, one diagnostics entry per backend optimization with requested/enabled/disable_reason. Backends that declare no optional tensor runtime optimization, or whose diagnostic probe fails, record an explicit unavailable marker instead of a null field.