refactor(env): make unilab a self-contained env package

.github/workflows/ci.yml

@@ -73,7 +73,7 @@ jobs:
           enable-cache: true
       - name: Install dependencies
         run: |
-          uv sync \
+          uv sync --extra unirl \
             --no-install-package torch \
             --no-install-package nvidia-cublas-cu12 \
             --no-install-package nvidia-cuda-cupti-cu12 \
@@ -109,7 +109,7 @@ jobs:
           enable-cache: true
       - name: Install dependencies
         run: |
-          uv sync \
+          uv sync --extra unirl \
             --no-install-package torch \
             --no-install-package nvidia-cublas-cu12 \
             --no-install-package nvidia-cuda-cupti-cu12 \
@@ -228,7 +228,7 @@ jobs:
         run: sudo apt-get update && sudo apt-get install -y libwayland-client0
       - name: Install dependencies
         run: |
-          uv sync --extra mujoco --extra mjwarp \
+          uv sync --extra mujoco --extra mjwarp --extra unirl \
             --no-install-package torch \
             --no-install-package nvidia-cublas-cu12 \
             --no-install-package nvidia-cuda-cupti-cu12 \

.github/workflows/tensor-runtime-ci.yml

@@ -43,7 +43,7 @@ jobs:
           python-version: "3.11"
       - name: Install CPU runtime
         run: |
-          uv sync --extra mujoco \
+          uv sync --extra mujoco --extra unirl \
             --no-install-package torch \
             --no-install-package nvidia-cublas-cu12 \
             --no-install-package nvidia-cuda-cupti-cu12 \
@@ -89,7 +89,7 @@ jobs:
           enable-cache: false
           python-version: "3.11"
       - name: Install tensor runtime
-        run: uv sync --extra mujoco --extra mjwarp --extra newton --extra motrix
+        run: uv sync --extra mujoco --extra mjwarp --extra newton --extra motrix --extra unirl
       - name: CUDA task contract tests
         run: >
           uv run --no-sync pytest -q

AGENTS.md

@@ -54,14 +54,15 @@ architecture changes.
 
 - UniLab provides environments and adapters, and now also owns the RL
   algorithm and training-execution stack in `unilab.rl` (merged from the
-  unilab_rl repo @ 3654e4b). Its layers: `contracts` (env injection and
-  observation semantics), `adapters/rsl_rl`, `algos` (APPO/SAC/FlashSAC),
-  `runtime` (distributed/ipc/offpolicy), and `logging`. Internal DAG:
-  `contracts` depends only on type/tensor libraries; `runtime` must not
-  depend on `algos`/`training`/`base`; `algos` must not construct concrete
-  envs; the RL core must not import `base`/`envs`/`tasks`/`managers`/
-  `training`/`scripts`/`unisim`. These boundaries are guarded by
-  `tests/test_library_import_boundary.py`. The PPO path is a direct rsl-rl
+  unilab_rl repo @ 3654e4b). The env-injection contract and observation
+  semantics live in `unilab.base.contracts`; the `unilab.rl` layers are
+  `adapters/rsl_rl`, `algos` (APPO/SAC/FlashSAC), `runtime`
+  (distributed/ipc/offpolicy), and `logging`. Internal DAG: `base.contracts`
+  depends only on type/tensor libraries; `runtime` must not depend on
+  `algos`/`training`/`base`; `algos` must not construct concrete envs; the RL
+  core must not import `base`/`envs`/`tasks`/`managers`/`training`/`scripts`/
+  `unisim` except the `unilab.base.contracts` leaf. These boundaries are
+  guarded by `tests/test_library_import_boundary.py`. The PPO path is a direct rsl-rl
   integration (`unilab.rl.adapters.rsl_rl`: VecEnv adapter and distributed
   helpers; runner and algorithm are upstream rsl-rl). The env factory must
   be a pickleable `EnvFactory` for spawn collectors. The `unilab.rl`,

Makefile

@@ -1,6 +1,6 @@
 .PHONY: sync
 sync:
-	uv sync --extra mujoco
+	uv sync --extra mujoco --extra unirl
 
 .PHONY: setup
 setup: sync
@@ -14,12 +14,12 @@ install-completion:
 sync-rocm:
 	@cp pyproject.rocm.toml pyproject.toml
 	@if [ -f uv.rocm.lock ]; then cp uv.rocm.lock uv.lock; fi
-	uv sync --extra mujoco
+	uv sync --extra mujoco --extra unirl
 	cp uv.lock uv.rocm.lock
 
 .PHONY: sync-xpu
 sync-xpu:
-	uv sync --extra mujoco --no-install-package torch
+	uv sync --extra mujoco --extra unirl --no-install-package torch
 	uv pip install torch==2.7.0 --torch-backend xpu
 
 .PHONY: format

docs/sphinx/AGENTS.md

@@ -119,7 +119,7 @@ language-independent absolute path.
    - algorithms and tasks: `src/unilab/conf/`, `src/unilab/scripts/train_*.py`,
      `src/unilab/rl/algos/`
    - env contract: `src/unilab/base/torch_env.py`,
-     `src/unilab/rl/contracts/env.py`
+     `src/unilab/base/contracts/env.py`
    - backend contract: `unisim.backend.base`
    - registry: `src/unilab/base/registry.py`
    - runner/IPC: `src/unilab/rl/runtime/async_runner.py`,

docs/sphinx/source/api_reference/algos/index.md

@@ -16,7 +16,7 @@ The PPO (RSL-RL) integration lives in `unilab.rl.adapters.rsl_rl`:
 
 UniLab keeps the training *entrypoints* (`src/unilab/scripts/train_*.py`),
 which inject environments into the RL runners through
-`unilab.rl.contracts.env.EnvFactory`; see `src/unilab/base/env_factory.py` for the
+`unilab.base.contracts.env.EnvFactory`; see `src/unilab/base/env_factory.py` for the
 registry-backed adapter.
 
 All trainers conform to a single runner contract — see

docs/sphinx/source/en/1-getting_started/6-migration_1_4_to_1_5.md

@@ -28,7 +28,7 @@ the restored form — for example
 | `EventTermCfg.min_step_count_between_reset` | ignored for reset events | per-row trigger filtering restored (1.3.x semantics) |
 | Reset owners | `reset_transaction(reset_mask)` / `reset_committed(reset_mask)` | `reset_transaction(env_ids)` / `reset_committed(env_ids)` |
 | Recorder hooks | `record_pre_reset(reset_mask)` / `record_post_reset(reset_mask)` | `record_pre_reset(env_ids)` / `record_post_reset(env_ids)` |
-| RL env contract | `EnvProtocol.reset(env_indices)` | `EnvProtocol.reset(env_ids)` (see `src/unilab/rl/contracts/env.py`) |
+| RL env contract | `EnvProtocol.reset(env_indices)` | `EnvProtocol.reset(env_ids)` (see `src/unilab/base/contracts/env.py`) |
 
 `env_ids` is a one-dimensional integer tensor of row indices; `None` (and, at
 manager boundaries, a `slice`) selects every row. The restored form matches

docs/sphinx/source/en/3-deployment/3-framework_migration/7-from_unilab_rl.md

@@ -26,10 +26,10 @@ runtime, and the RL logging stack ship in the base install.
 
 | Old import (`uni_rl.*`) | New import |
 |---|---|
-| `uni_rl.env_contract` | `unilab.rl.contracts.env` |
+| `uni_rl.env_contract` | `unilab.base.contracts.env` |
 | `uni_rl.training_state` | no replacement (`unilab.training.experiment` owns checkpoint save/load) |
-| `uni_rl.utils.observations` | `unilab.rl.contracts.observations` |
-| `uni_rl.algos.common.device` (probe) | `unilab.rl.contracts.probe` |
+| `uni_rl.utils.observations` | `unilab.base.contracts.observations` |
+| `uni_rl.algos.common.device` (probe) | `unilab.base.contracts.probe` |
 | `uni_rl.algos.appo` / `.sac` / `.flash_sac` / `.common` | `unilab.rl.algos.*` (same subpackage names) |
 | `uni_rl.algos.common.collector_timing` | `unilab.rl.runtime.collector_timing` |
 | `uni_rl.ipc.async_runner` | `unilab.rl.runtime.async_runner` |

docs/sphinx/source/zh_CN/1-getting_started/6-migration_1_4_to_1_5.md

@@ -25,7 +25,7 @@ UniLab 1.5 将公共 reset 契约恢复为 IsaacLab/mjlab 风格的 `env_ids` 
 | `EventTermCfg.min_step_count_between_reset` | reset 事件不生效 | 恢复逐行触发过滤（1.3.x 语义） |
 | reset owner | `reset_transaction(reset_mask)` / `reset_committed(reset_mask)` | `reset_transaction(env_ids)` / `reset_committed(env_ids)` |
 | recorder 钩子 | `record_pre_reset(reset_mask)` / `record_post_reset(reset_mask)` | `record_pre_reset(env_ids)` / `record_post_reset(env_ids)` |
-| RL 环境契约 | `EnvProtocol.reset(env_indices)` | `EnvProtocol.reset(env_ids)`（见 `src/unilab/rl/contracts/env.py`） |
+| RL 环境契约 | `EnvProtocol.reset(env_indices)` | `EnvProtocol.reset(env_ids)`（见 `src/unilab/base/contracts/env.py`） |
 
 `env_ids` 是一维整数行索引张量；`None`（manager 边界还接受 `slice`）表示全部行。
 恢复后的形态与 v1.3.x 公共契约一致，针对 1.3.x 编写的代码无需修改。

docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/7-from_unilab_rl.md

@@ -24,10 +24,10 @@
 
 | 旧 import（`uni_rl.*`） | 新 import |
 |---|---|
-| `uni_rl.env_contract` | `unilab.rl.contracts.env` |
+| `uni_rl.env_contract` | `unilab.base.contracts.env` |
 | `uni_rl.training_state` | 无替代（checkpoint 保存/加载由 `unilab.training.experiment` 负责） |
-| `uni_rl.utils.observations` | `unilab.rl.contracts.observations` |
-| `uni_rl.algos.common.device`（probe） | `unilab.rl.contracts.probe` |
+| `uni_rl.utils.observations` | `unilab.base.contracts.observations` |
+| `uni_rl.algos.common.device`（probe） | `unilab.base.contracts.probe` |
 | `uni_rl.algos.appo` / `.sac` / `.flash_sac` / `.common` | `unilab.rl.algos.*`（子包名不变） |
 | `uni_rl.algos.common.collector_timing` | `unilab.rl.runtime.collector_timing` |
 | `uni_rl.ipc.async_runner` | `unilab.rl.runtime.async_runner` |

pyproject.rocm.toml

@@ -32,17 +32,13 @@ dependencies = [
     "etils",
     "hydra-core>=1.3",
     "omegaconf",
-    "tensordict",
-    "rsl-rl-lib>=5.0.0",
     "packaging",
     "mediapy",
-    "tensorboard",
     "setuptools",
     "rich",
     "tqdm",
     "typing-extensions",
     "lark>=1.3.1",
-    "wandb>=0.22.3",
     "onnxruntime<1.20 ; python_version < '3.11'",
     "onnxruntime>=1.20 ; python_version >= '3.11'",
     "huggingface_hub>=0.25",
@@ -76,6 +72,13 @@ motrix = ["unisim-core[motrix]==2.2.0"]
 genesis = ["genesis-world==1.3.3"]
 # Keep the optional tensor-ray extra aligned with pyproject.toml.
 ray = ["uniray[warp]==0.1.0"]
+# Keep the training-stack extra aligned with pyproject.toml.
+unirl = [
+    "tensordict",
+    "rsl-rl-lib>=5.0.0",
+    "tensorboard",
+    "wandb>=0.22.3",
+]
 
 [dependency-groups]
 # The RL stack (APPO/off-policy/multi-GPU) is merged into unilab itself

pyproject.toml

@@ -43,17 +43,13 @@ dependencies = [
     "etils",
     "hydra-core>=1.3",
     "omegaconf",
-    "tensordict",
-    "rsl-rl-lib>=5.0.0",
     "packaging",
     "mediapy",
-    "tensorboard>=2.21.0",
     "setuptools",
     "rich",
     "tqdm",
     "typing-extensions",
     "lark>=1.3.1",
-    "wandb>=0.22.3",
     "onnxruntime<1.20 ; python_version < '3.11'",
     "onnxruntime>=1.20 ; python_version >= '3.11'",
     "huggingface_hub>=0.25",
@@ -117,6 +113,16 @@ superdex = [
 # caster factory. The PyPI distribution is `uniray`; the import namespace is
 # `uni_ray`.
 ray = ["uniray[warp]==0.1.0"]
+# The merged RL/training stack (unilab.rl) needs these optional packages;
+# `pip install unilab` stays a lightweight env without them. Single extra on
+# purpose: APPO/SAC/FlashSAC/off-policy share these; only PPO (rsl-rl) is
+# lighter and does not justify its own extra.
+unirl = [
+    "tensordict",
+    "rsl-rl-lib>=5.0.0",
+    "tensorboard>=2.21.0",
+    "wandb>=0.22.3",
+]
 
 [dependency-groups]
 # The RL stack (APPO/off-policy/multi-GPU) is merged into unilab itself

scripts/benchmark/rl/benchmark_offpolicy_collector_active.py

@@ -771,7 +771,7 @@ def _make_env(
     env_cfg_override: dict[str, Any] | None,
 ):
     from unilab.base.config_adapter import create_env
-    from unilab.rl.contracts.observations import get_obs_dims
+    from unilab.base.contracts.observations import get_obs_dims
 
     env = create_env(cfg, num_envs=int(cfg.algo.num_envs), env_cfg_override=env_cfg_override)
     if env.state is None:
@@ -860,7 +860,7 @@ def _run_active_window_case(
     policy_inference_context: TimingStats | None = None,
     policy_inference_context_evidence: dict[str, Any] | None = None,
 ) -> CollectorResult:
-    from unilab.rl.contracts.observations import split_obs_dict
+    from unilab.base.contracts.observations import split_obs_dict
     from unilab.rl.runtime.ipc.replay_buffer import ReplayBuffer
     from unilab.rl.runtime.offpolicy.worker import timeout_bootstrap_mask
 

scripts/benchmark/rl/benchmark_replay_buffer_placement.py

@@ -312,7 +312,7 @@ def _resolve_targets(
 
 def _resolve_env_shape(cfg: DictConfig, algo: str) -> ReplayShape:
     from unilab.base.config_adapter import EnvConfigAdapter, create_env
-    from unilab.rl.contracts.observations import get_obs_dims
+    from unilab.base.contracts.observations import get_obs_dims
     from unilab.training import ensure_registries
 
     ensure_registries()

scripts/benchmark/rl/benchmark_sac_replay_buffer_sampling.py

@@ -255,7 +255,7 @@ def _compose_offpolicy_cfg(task: str = DEFAULT_TASK, sim: str = DEFAULT_SIM) ->
 
 def _resolve_env_shape(cfg: DictConfig) -> ReplayShape:
     from unilab.base.config_adapter import EnvConfigAdapter, create_env
-    from unilab.rl.contracts.observations import get_obs_dims
+    from unilab.base.contracts.observations import get_obs_dims
     from unilab.training import ensure_registries
 
     ensure_registries()

src/unilab/base/__init__.py

@@ -1,7 +1,21 @@
-"""Environment registry and base classes."""
+"""Environment registry and base classes.
 
-from unilab.base.registry import ensure_registries
+``ensure_registries`` is re-exported lazily so importing ``unilab.base`` (and
+therefore ``unilab.base.contracts``) does not pull in the registry, the torch
+env, or any backend package. Most callers import the concrete submodule
+directly (``from unilab.base import registry``).
+"""
 
-__all__ = [
-    "ensure_registries",
-]
+from __future__ import annotations
+
+from typing import Any
+
+__all__ = ["ensure_registries"]
+
+
+def __getattr__(name: str) -> Any:
+    if name == "ensure_registries":
+        from unilab.base.registry import ensure_registries
+
+        return ensure_registries
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/unilab/base/contracts/__init__.py

@@ -0,0 +1,41 @@
+"""Env injection contracts and observation semantics (single owner).
+
+Cold-path contract definitions only; this subpackage depends on stdlib,
+typing, and Torch types, never on concrete envs, registries, or algorithms.
+"""
+
+from unilab.base.contracts.env import (
+    EnvAlgoCapabilities,
+    EnvAlgoCapabilitiesProtocol,
+    EnvFactory,
+    EnvPlayCapabilitiesProtocol,
+    EnvProtocol,
+    EnvStateProtocol,
+    SupportsAlgoCapabilitiesProtocol,
+    SupportsBackendDiagnosticsProtocol,
+    SupportsCompletionReceiptsProtocol,
+    get_algo_capabilities,
+    get_backend_runtime_diagnostics,
+    get_produces_completion_receipts,
+)
+from unilab.base.contracts.observations import get_obs_dims, split_obs_dict
+from unilab.base.contracts.probe import EnvContractProbe, probe_env_contract
+
+__all__ = [
+    "EnvAlgoCapabilities",
+    "EnvAlgoCapabilitiesProtocol",
+    "EnvContractProbe",
+    "EnvFactory",
+    "EnvPlayCapabilitiesProtocol",
+    "EnvProtocol",
+    "EnvStateProtocol",
+    "SupportsAlgoCapabilitiesProtocol",
+    "SupportsBackendDiagnosticsProtocol",
+    "SupportsCompletionReceiptsProtocol",
+    "get_algo_capabilities",
+    "get_backend_runtime_diagnostics",
+    "get_obs_dims",
+    "get_produces_completion_receipts",
+    "probe_env_contract",
+    "split_obs_dict",
+]

src/unilab/rl/contracts/env.py → src/unilab/base/contracts/env.py

@@ -1,4 +1,4 @@
-"""THE tensor-native env contract consumed by ``unilab.rl`` runners.
+"""THE tensor-native env contract consumed by RL runners and collectors.
 
 This layer never constructs environments itself and never imports an env
 registry. Environments are injected by the caller (e.g. UniLab's training

src/unilab/rl/contracts/observations.py → src/unilab/base/contracts/observations.py

@@ -1,7 +1,7 @@
-"""Observation-group helpers for the ``unilab.rl`` env contract.
+"""Observation-group helpers for the env contract.
 
 Generic helpers over the ``obs_groups_spec`` / observation-dict contract
-defined in ``unilab.rl.contracts.env``. Observation values are Torch tensors at
+defined in ``unilab.base.contracts.env``. Observation values are Torch tensors at
 the env boundary (see ``env_contract``); the helpers stay value-agnostic so
 host-side diagnostics with NumPy carriers can reuse them.
 """

src/unilab/rl/contracts/probe.py → src/unilab/base/contracts/probe.py

@@ -6,8 +6,8 @@
 from dataclasses import dataclass
 from typing import Any
 
-from unilab.rl.contracts.env import EnvFactory, get_produces_completion_receipts
-from unilab.rl.contracts.observations import get_obs_dims as _get_obs_dims
+from unilab.base.contracts.env import EnvFactory, get_produces_completion_receipts
+from unilab.base.contracts.observations import get_obs_dims as _get_obs_dims
 
 
 @dataclass(frozen=True)

src/unilab/base/env_factory.py

@@ -1,12 +1,11 @@
 """Picklable ``EnvFactory`` adapters bridging the RL env contract to the
 UniLab env registry.
 
-RL runners and collectors (``unilab.rl``) never construct environments
-themselves; the caller injects an ``EnvFactory`` (see
-``unilab.rl.contracts.env``). Because collectors run in ``multiprocessing``
-spawn subprocesses, the factory must be picklable by reference — the adapters
-below are module-level functions bound with ``functools.partial`` (never
-closures or lambdas).
+RL runners and collectors never construct environments themselves; the caller
+injects an ``EnvFactory`` (see ``unilab.base.contracts.env``). Because
+collectors run in ``multiprocessing`` spawn subprocesses, the factory must be
+picklable by reference — the adapters below are module-level functions bound
+with ``functools.partial`` (never closures or lambdas).
 """
 
 from __future__ import annotations
@@ -20,7 +19,7 @@
 if TYPE_CHECKING:
     # ``EnvFactory``/``EnvProtocol`` are typing-only Protocols, so the import
     # is annotation-only and must not run at module import time.
-    from unilab.rl.contracts.env import EnvFactory, EnvProtocol
+    from unilab.base.contracts.env import EnvFactory, EnvProtocol
 
 
 def make_registry_env(

src/unilab/cli.py

@@ -61,6 +61,12 @@ class Route:
 
 
 def package_root() -> Path:
+    """Return the source-checkout package root (``src/unilab``).
+
+    Internal: this path only exists in a source checkout, not an installed
+    wheel. Downstream packages must not rely on it; register tasks through the
+    ``unilab.tasks`` entry-point contract instead.
+    """
     return Path(__file__).resolve().parent