unilab.training.tensor_diagnostics

Finite-state diagnostics for the tensor-native environment runtime.

The guard’s detection path stays in Torch: it never converts observations, rewards, controls, or physics snapshots to host arrays. Host conversion and artifact writing happen only in NanGuard.dump(), which is an explicit abnormal-diagnostic boundary.

Classes

NanGuard

Detect bad tensor rows on-device and export dumps at an explicit host boundary.

NanGuardCfg

Pickleable configuration for the tensor finite-state guard.
class unilab.training.tensor_diagnostics.NanGuardCfg[source]

Bases: object

Pickleable configuration for the tensor finite-state guard.

Parameters:
enabled: bool = False
buffer_size: int = 100
max_envs_to_dump: int = 5
output_dir: str | None = None
__init__(enabled=False, buffer_size=100, max_envs_to_dump=5, output_dir=None)
Parameters:
class unilab.training.tensor_diagnostics.NanGuard[source]

Bases: object

Detect bad tensor rows on-device and export dumps at an explicit host boundary.

Parameters:
__init__(cfg, num_envs, supports_state_playback)[source]
Parameters:
capture(physics_state)[source]

Retain a detached device tensor snapshot for a later abnormal dump.

Parameters:

physics_state (Tensor | None)

Return type:

None

check(obs, reward, step=0)[source]

Return device-resident bad-row indices, or None when every row is finite.

Parameters:
Return type:

Tensor | None

check_ctrl(ctrl, step=0)[source]

Return device-resident bad control-row indices, or None when finite.

Parameters:
Return type:

Tensor | None

dump(nan_env_ids, model_file, step)[source]

Export one abnormal-state artifact at an explicit host boundary.

Parameters:
Return type:

str | None