Metadata-Version: 2.4
Name: unilab-rl
Version: 1.4.11
Summary: Reinforcement learning algorithms and async runtimes extracted from UniLab (APPO/SAC/FlashSAC)
Keywords: reinforcement-learning,robotics,gpu-simulation
Author: UniLab contributors
License-Expression: Apache-2.0
License-File: LICENSE
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Requires-Dist: numpy
Requires-Dist: torch>=2.7.0
Requires-Dist: tensordict
Requires-Dist: omegaconf
Requires-Dist: hydra-core
Requires-Dist: rsl-rl-lib>=5.0.0
Requires-Dist: tensorboard
Requires-Dist: wandb>=0.22.3
Requires-Dist: packaging
Requires-Dist: rich
Requires-Python: >=3.10, <3.14
Project-URL: Repository, https://github.com/unilabsim/unilab_rl
Description-Content-Type: text/markdown

# unilab-rl

[![PyPI](https://img.shields.io/pypi/v/unilab-rl)](https://pypi.org/project/unilab-rl/)
[![CI](https://github.com/unilabsim/unilab_rl/actions/workflows/ci.yml/badge.svg)](https://github.com/unilabsim/unilab_rl/actions/workflows/ci.yml)
[![License](https://img.shields.io/badge/license-Apache--2.0-blue.svg)](LICENSE)

English | [简体中文](README_zh.md)

Reinforcement learning algorithms and asynchronous runtimes extracted from
[UniLab](https://github.com/unilabsim/UniLab), packaged as a standalone,
simulator-agnostic library.

- Distribution name: `unilab-rl`
- Import namespace: `uni_rl`
- Repository: [unilabsim/unilab_rl](https://github.com/unilabsim/unilab_rl)

## Relationship with UniLab

`uni_rl` is the RL algorithm and async-runtime layer of the UniLab project,
split out into its own package. [UniLab](https://github.com/unilabsim/UniLab)
remains the consumer side: it owns the physics backends, task suites, and
training entrypoints, and injects environments into `uni_rl` through
`uni_rl.env_contract.EnvFactory`. `uni_rl` never imports `unilab` / `unisim`
and never constructs environments itself, so any vectorized environment
satisfying the contract — including simulators outside UniLab — can drive the
algorithms in this package.

UniLab consumes `uni_rl` as an optional extra (`unilab[uni_rl]`) for APPO,
off-policy algorithms, and multi-GPU data-parallel PPO launches; its
single-process PPO path drives upstream rsl_rl directly. Install `unilab-rl`
directly when you want to reuse its algorithms and async runtime with your own
environment stack.

> Naming note: the originally intended distribution name `uni-rl` is
> unregistrable on PyPI because it ultranormalizes to the existing `unirl`
> project. The distribution is therefore published as `unilab-rl`; the import
> namespace remains `uni_rl` as designed.

## Contents

- **Async PPO (APPO)**: native collector/learner multiprocess implementation
  (actor/critic networks built on
  [rsl_rl](https://github.com/leggedrobotics/rsl_rl) model classes)
- **Off-policy**: SAC and FlashSAC with double-buffer async runners;
  FlashSAC supports optional age-biased replay sampling
- **Runtime infrastructure**: shared-memory rollout/replay buffers, replay
  pipelines, data-parallel gradient sync, memory budgeting, tensorboard/wandb
  training loggers, and a trace recorder

## Layout

- `uni_rl.algos.*` — the algorithm layer: async on-policy (`appo`),
  off-policy learners (`sac`, `flash_sac`), and shared
  algorithm helpers (`common`)
- `uni_rl.ipc` — runtime infrastructure: async runner, shared-memory
  rollout/replay buffers, replay pipelines, DP gradient sync, memory budget
- `uni_rl.offpolicy` — the generic off-policy double-buffer runner scaffolding
- `uni_rl.logging` — tensorboard/wandb training loggers, trace recorder
- `uni_rl.utils` — device, seed, nan-guard, observation helpers
- `uni_rl.env_contract` — the injected env factory/protocol contract

## Installation

```bash
pip install unilab-rl
# or, with uv:
uv add unilab-rl
```

Requires Python 3.10–3.13 and PyTorch ≥ 2.7.

## Usage

`uni_rl` does not construct environments. Inject a picklable env factory
(`EnvFactory = Callable[[int, Mapping | None], EnvProtocol]`) into the runner
of your chosen algorithm:

```python
from collections.abc import Mapping

from uni_rl.env_contract import EnvProtocol


def make_env(num_envs: int, cfg: Mapping | None) -> EnvProtocol:
    """Top-level factory (picklable by reference; no closures/lambdas)."""
    ...
```

The env contract is tensor-native and autoresetting: `step()` accepts a
contiguous float32 Torch action tensor on the env's authoritative `device`,
observations/rewards/termination flags are Torch tensors keyed by observation
group (`obs_groups_spec`), and `reset()` accepts a one-dimensional Torch
integer tensor (or `None`) and returns `(obs, info)`. NumPy remains a
serialization/diagnostic carrier only, not the trainer boundary. See the module
docstring in
[`src/uni_rl/env_contract.py`](src/uni_rl/env_contract.py) for the full
contract, and the *new algorithm recipe* section in
[`AGENTS.md`](AGENTS.md) for how to plug in a custom algorithm via
`runtime_resolver` without forking.

The breaking canonical TensorBoard/W&B field contract and the historical
old-to-new migration table are documented in
[`docs/metrics.md`](docs/metrics.md).

## Off-policy cold-path preparation

The double-buffer runner performs learner-owned warmup after DP initialization
and before starting the collector. Implement
`prepare_for_collection(context: OffPolicyWarmupContext)` on a custom learner
for compilation, graph capture, and other cold paths. Preparation must leave
weights, optimizers, schedulers, RNG state, and counters unchanged; compiler and
graph caches are the only sanctioned retained effects. A custom SAC runtime
may instead provide `OffPolicyRuntime.learner_prepare_hook`, and actor adapters
may provide `warmup_actions`.

Coordination failures use learner phase/progress and process liveness rather
than a wall-clock performance SLA. `training.inference_request_timeout_sec` is
deprecated and ignored; remove it from owner YAML during migration.

## FlashSAC learner acceleration controls

FlashSAC exposes three separate learner controls:

- `use_compile` enables or disables CUDA Inductor compilation. Omitting it
  preserves the auto-enabled NVIDIA CUDA default; an explicit `false` selects
  the eager per-update path.
- `compile_full_objectives` chooses whether Inductor compiles the complete
  critic/actor objectives or the smaller loss-tensor regions. Omitting it
  preserves the full-objective NVIDIA default; it has no effect when
  compilation is disabled.
- `use_whole_cycle_cuda_graph` enables or disables learner-owned whole-cycle
  CUDA Graph capture. Omitting it preserves the single-NVIDIA default. This
  layer requires compilation, so `use_compile=false` also disables it; setting
  it to `false` leaves the selected Inductor objective scope enabled.

## Design contract

`uni_rl` does **not** depend on any simulator or environment library.
Algorithm behavior is owned by the algo modules under `uni_rl.algos.*`;
runtime infrastructure (`ipc`, `logging`, `offpolicy`, `utils`,
`env_contract`) lives at the top level and never depends on the algorithm
layer. See UniLab's training entrypoints for reference env integrations.

## Development

```bash
make sync      # install dependencies (uv)
make test      # pytest
make format    # ruff check --fix + ruff format
uv run mypy src/uni_rl && uv run pyright   # type gates
```

## Citation

If you use `unilab-rl` in your research, please cite the UniLab paper:

```bibtex
@article{jia2026unilab,
  title   = {UniLab: A Heterogeneous Architecture for Robot RL Beyond GPU-Dominant Paradigms},
  author  = {Yufei Jia and Zhanxiang Cao and Mingrui Yu and Heng Zhang and Shenyu Chen and Dixuan Jiang and Meng Li and Xiaofan Li and Yiyang Liu and Junzhe Wu and Zheng Li and XiLin Fang and Tingyu Cui and Shengcheng Fu and Haoyang Li and Anqi Wang and Zifan Wang and Dongjie Zhu and Chenyu Cao and Zhenbiao Huang and Ziang Zheng and Jie Lu and Xin Ma and Zhengyang Wei and Xiang Zhao and Tianyue Zhan and Ye He and Yuxiang Chen and Yizhou Jiang and Yue Li and Haizhou Ge and Yuhang Dong and Fan Jia and Ziheng Zhang and Meng Zhang and Xiwa Deng and Zhixing Chen and Hanyang Shao and Chenxin Dong and Yixuan Li and Yizhi Chen and Bokui Chen and Kaifeng Zhang and Hanqing Cui and Yusen Qin and Ruqi Huang and Lei Han and Tiancai Wang and Xiang Li and Yue Gao and Guyue Zhou},
  journal = {arXiv preprint arXiv:2605.30313},
  year    = {2026},
  url     = {https://arxiv.org/abs/2605.30313}
}
```

## License

Apache-2.0, same as UniLab.
