Language

unilab.base.torch_env.TorchEnv¶

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)

Methods

`__init__`(cfg, backend, num_envs, *[, device])
`apply_action`(actions, state)	Transform an action tensor into a backend control tensor.
`capture_play_video_frame`()	Capture one detached RGB video frame through the env contract.
`close`()	Clean up environment resources
`export_training_state`()
`get_backend_runtime_diagnostics`()	Serialize the backend's tensor runtime diagnostics on the cold path.
`get_physics_state_layout`()	Return the backend physics-state playback layout.
`get_physics_state_snapshot`()	Return a detached physics snapshot for playback/video export.
`get_playback_mocap_state`([env_index])	Return detached `(mocap_pos, mocap_quat)` playback state.
`get_playback_model`([env_index])	Return the backend playback model for one vectorized environment.
`get_scene_visual_model_file`()	Return the backend visual model file on the playback cold path.
`import_training_state`(state)
`init_play_renderer`([render_spacing, ...])	Initialize backend-native playback rendering when available.
`init_state`()	Initialize environment and return initial state
`render`([mode])	Render the current state to an RGB array through the play renderer.
`reset`([env_ids])	Reset the selected environment rows; `None` resets all rows.
`resolve_play_render_plan`(*, ...)	Resolve high-level playback mode through the concrete backend.
`run_playback`(*, initialize, step, num_steps)	Execute playback through the concrete backend.
`run_playback_mode`(*, play_render_mode, ...)	Resolve configured playback mode and execute it through the backend contract.
`set_autoreset`(enabled)
`set_nan_guard`(guard)	Attach the owner-provided diagnostic guard.
`step`(actions)	Step the environment with owner-declared action carrier.
`update_state`(state)	Compute tensor observations, reward, and termination state.

Attributes

`action_space`	Action space
`backend`	Public backend owner used by Manager terms and tensor contracts.
`cfg`	The configuration of the environment
`device`
`num_envs`	return the size of the env if it is vectorized
`obs_groups_spec`	Map from observation group name to its dimension.
`observation_space`	Observation space
`play_capabilities`	Return env-facing play/render capabilities.
`produces_completion_receipts`	RL env-contract declaration of the backend's step receipts.
`state`	Current environment state (None before first reset)

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

Return type:: dict[str, Any]

export_training_state()[source]¶

Return type:: dict[str, Any]

import_training_state(state)[source]¶

Parameters:: state (Mapping[str, Any])
Return type:: None

close()[source]¶

Clean up environment resources

Return type:: None

abstract apply_action(actions, state)[source]¶

Transform an action tensor into a backend control tensor.

Parameters:

actions (Tensor)
state (TorchEnvState)

Return type:

Tensor

abstract update_state(state)[source]¶

Compute tensor observations, reward, and termination state.

Parameters:: state (TorchEnvState)
Return type:: TorchEnvState

abstract property action_space: Space¶: Action space

run_playback_mode(*, play_render_mode, play_steps, output_video, initialize, step, render_spacing=None, render_offset_mode=None, frame_state_getter=None, camera_kwargs=None, debug_overlay_getter=None, on_frame=None, on_plan=None)¶

Resolve configured playback mode and execute it through the backend contract.

Overlay gating follows the resolved plan mode: record playback requires play_capabilities.supports_debug_overlay (backends fail closed otherwise); interactive playback requires supports_interactive_debug_overlay — when the backend lacks it the getter is dropped with a warning so interactive playback still runs without overlays.

camera_kwargs is a record/offline camera profile. Some native interactive viewers (for example the Isaac family) fail closed when they receive one because the user owns the camera in the window; in that case the profile is dropped with a warning and playback retried without it, so interactive playback still runs.

Parameters:

play_render_mode (str | None)
play_steps (int | None)
output_video (str | PathLike[str] | None)
initialize (Callable[[], Any])
step (Callable[[Any], Any])
render_spacing (float | None)
render_offset_mode (str | None)
frame_state_getter (Callable[[], ndarray] | None)
camera_kwargs (CameraCfg | Mapping[str, Any] | None)
debug_overlay_getter (Callable[[], Sequence[Sequence[DebugPrimitive] | None] | None] | None)
on_frame (Callable[[int, ndarray], ndarray | None] | None)
on_plan (Callable[[PlayRenderPlan], None] | None)

Return type:

str | None