Skip to content

fix(offpolicy): default nan guard passes the wrong guard type to TorchEnv #1793

Closed

Description

@TATP-233

Fixed reproduction

Current develop/tensor-runtime@b6a7c54f, default SAC nan-guard setting (training.nan_guard.enabled=true):

CUDA_VISIBLE_DEVICES=0 uv run --no-sync train --algo sac --task g1_walk_flat --sim mujoco \
  training.devices=[0] +training.inference_transport=cuda training.no_play=true \
  training.export_onnx=false algo.max_iterations=50 training.log_dir=/tmp/mjbatch-owner-smoke

The collector crashes during startup:

File ".../uni_rl/offpolicy/worker.py", line 358, in _run_collector
    env.set_nan_guard(NanGuard(...))
File ".../unilab/base/torch_env.py", line 776, in set_nan_guard
TypeError: TorchEnv nan guard must be TensorNanGuard, got NanGuard

The same 50-iteration command completes when training.nan_guard.enabled=false, so this blocks the default-enabled diagnostic path rather than training itself.

Problem

uni_rl.offpolicy.worker constructs uni_rl.utils.nan_guard.NanGuard, while TorchEnv's contract requires UniLab's TensorNanGuard. The packages intentionally do not import each other, so the worker cannot construct UniLab's type directly. This needs a formal cross-package guard contract (for example a structural protocol/adapter or a guard factory passed from the UniLab runner) rather than an import dependency.

Acceptance criteria

  • Default training.nan_guard.enabled=true off-policy run starts and completes.
  • uni_rl still does not import UniLab.
  • Contract tests cover the collector-side guard injection boundary.
  • CUDA tensor-runtime runs retain finite ctrl/state/observation checks without extra hot-path transfers.

Activity

  1. TATP-233 commented on Sep 30, 2026

    @TATP-233
    CollaboratorAuthor

    Fixed by #1799 plus unilabsim/unilab_rl#71. UniLab injects a pickleable TensorNanGuard factory across the collector boundary; uni_rl accepts the environment-owned factory without importing UniLab. Merged to develop/tensor-runtime.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions