UniLab

repository·main·Indexed 21 days ago

https://github.com/unilabsim/unilab

Universal Lab for Robot Learning (version 0.1.0) is a heterogeneous architecture for Robot Reinforcement Learning (RL) designed to decouple CPU-based physics simulation from GPU-based policy training. It enables high-throughput training via CPU parallel simulation streams and shared memory. The project includes a Numba-fused kernel implementation for efficient reward and termination calculations, as well as comprehensive bilingual documentation covering user and developer guides, API references, and Architecture Decision Records (ADRs).

Tokens
137.6K
Snippets
324
Records
715
Agent score
74%

What's inside unilab

  1. Overview of unilab.utils utilities

    main

    The unilab.utils module provides a collection of shared utility functions used across different environments and scripts. It includes tools for:

    • Device Probing: Checking hardware availability and device properties.
    • Tensor Helpers: Manipulating and managing tensors.
    • Support-Matrix Bookkeeping: Managing compatibility and support matrices.
    • NaN Guards: Protecting against or detecting NaN values in computations.
    • Geometry & Rotation: Pure-NumPy implementations of geometric and rotational operations.
  2. Overview of UniLab Deployment Journeys

    main

    UniLab provides deployment playbooks for moving trained policies across different hardware, simulation backends, and source frameworks. Every deployment tutorial follows a structured format:

    1. What you start with: The trained artifact and its associated configuration.
    2. What changes: The minimal set of edits required in code, YAML files, and assets.
    3. How to validate: Concrete commands and checkpoints to ensure the deployment is successful.

    There are three primary deployment paths:

    • Sim → Real: Preparing policies for physical hardware (e.g., G1, Go2, or Allegro hands) using ONNX exports and contract checks.
    • Sim → Sim (MuJoCo ↔ Motrix): Switching a task between different simulation backends (MuJoCo and Motrix) without requiring a full retraining process.
    • Framework Migration: Porting tasks and policies from other reinforcement learning frameworks like Isaac Lab, Legged Gym, rsl_rl, or skrl into the UniLab ecosystem.
  3. Overview of available training algorithms

    main

    UniLab provides several reinforcement learning algorithms, each with a specific training style, entrypoint script, and configuration file. Use the following table to identify which script and configuration to use for your desired algorithm:

    AlgorithmStyleEntrypointConfig Evidence
    PPOsynchronous on-policyscripts/train_rsl_rl.pyconf/ppo/config.yaml
    APPOasync on-policyscripts/train_appo.pyconf/appo/config.yaml
    SACoff-policyscripts/train_offpolicy.pyconf/offpolicy/algo/sac.yaml
    TD3off-policyscripts/train_offpolicy.pyconf/offpolicy/algo/td3.yaml
    FlashSACoff-policyscripts/train_offpolicy.pyconf/offpolicy/algo/flashsac.yaml
    HIM-PPOheight-estimator PPO pathscripts/train_him_ppo.pyconf/ppo_him/config.yaml
    HORAteacher/student distillation pathscripts/train_hora_distill.pyconf/hora_distill/config.yaml
    MLX PPOsynchronous on-policy (Apple Silicon)scripts/train_mlx_ppo.pyconf/ppo/config_mlx.yaml

    For general CLI flags applicable to all algorithms, refer to the CLI reference documentation.

  4. Overview of the unilab package structure

    main

    The unilab package is organized into several functional layers: core foundations, a learning stack, and various subsystems. Understanding this hierarchy helps in locating specific components for simulation, training, or utility tasks.

    Core Foundations

    These modules define the fundamental contracts and interfaces used by the entire library:

    • unilab.base: Contains essential interfaces like NpEnv, SimBackend, Registry, and Scene.
    • unilab.envs: Provides concrete task implementations such as locomotion, manipulation, and motion tracking.

    Learning Stack

    These modules facilitate the training and execution of reinforcement learning agents:

    • unilab.algos: Implementations of RL algorithms (PPO, APPO, SAC, TD3) in both PyTorch and MLX.
    • unilab.training: Orchestration tools, runtime helpers, monitoring, and reward bookkeeping.
    • unilab.ipc: Shared-memory primitives for connecting CPU workers to GPU learners.
    • unilab.backend: Adapters for simulation engines like MuJoCo and Motrix that implement the SimBackend interface.

    Subsystems

    Specialized modules for environment augmentation and tooling:

    • unilab.dr: Declarative domain randomization.
    • unilab.terrains: Procedural and heightfield terrain generation.
    • unilab.visualization: Scene rendering and viser bridges.
    • unilab.tools: Scene export, NaN visualization, and ONNX export.
    • unilab.utils: Math, IO, and numerical helpers.
    • unilab.logging: Structured logging and bridges for W&B and TensorBoard.
  5. Overview of UniLab operational tools

    main

    UniLab provides a suite of operational tools designed for policy export, training failure inspection, metadata transmission, and scene instantiation. The available tools include:

    • ONNX Export: Used to export policies from replay paths and validate runtime inputs.
    • W&B & TensorBoard: Used to configure run logs and experiment metadata.
    • NaN Visualizer: Used to inspect NaN guard dumps during PPO (Proximal Policy Optimization) runs.
    • Scene Export: Used to export MuJoCo scenes and copied assets for inspection.
    • Robot Import: Used to integrate robot assets, model descriptions, control interfaces, and keyframes.
  6. Overview of unilab.algos learning algorithms

    main

    The unilab.algos package provides various reinforcement learning algorithms split into two primary backends based on the hardware and framework requirements:

    • unilab.algos.torch: Implements algorithms using PyTorch, including PPO (via RSL-RL), APPO, FastSAC, FastTD3, FlashSAC, HIM-PPO, HORA + distillation, and a generic off-policy runner.
    • unilab.algos.mlx: Implements Apple-silicon native PPO using the MLX framework.

    All trainers in this package follow a unified runner contract, ensuring consistent lifecycle management across different algorithm implementations.

  7. Overview of UniLab Contracts

    main

    Contracts in UniLab define the strict boundaries and requirements that different system components must preserve to ensure interoperability. These contracts govern the interaction between task code, backend adapters, runners, and algorithm entrypoints.

    Key contract areas include:

    • Env contract: Defines NpEnvState, reset/step shapes, observation groups, and wrappers.
    • Backend contract: Defines SimBackend ownership and explicit capability support.
    • Task owner contract: Governs Hydra owner YAML identity and backend selection.
    • Domain randomization contract: Defines boundaries for initialization, reset, intervals, and backend capabilities.
    • Runner lifecycle contract: Defines ownership for runner start, stop, and checkpointing operations.
  8. Explore PyTorch reinforcement learning algorithms in `unilab.algos.torch`

    main

    The unilab.algos.torch module provides a collection of PyTorch-based reinforcement learning algorithm implementations. The available submodules include:

    • unilab.algos.torch.common: Common utilities and base classes.
    • unilab.algos.torch.appo: APPO (Asynchronous Proximal Policy Optimization).
    • unilab.algos.torch.fast_sac: Fast Soft Actor-Critic.
    • unilab.algos.torch.fast_td3: Fast Twin Delayed DDPG.
    • unilab.algos.torch.flash_sac: Flash Soft Actor-Critic.
    • unilab.algos.torch.him_ppo: Hindsight Imitation Learning PPO.
    • unilab.algos.torch.hora: HORA algorithm.
    • unilab.algos.torch.offpolicy: Various off-policy algorithm implementations.
  9. Access API reference and project documentation

    main

    The UniLab documentation provides several reference entry points for developers:

    • API Reference: Contains autodoc outputs for public Python interfaces, including type information.
    • Glossary: A collection of standardized terminology used throughout the documentation and codebase.
    • Changelog: Documentation of released versions and release notes.
    • Backend Support Matrix: An automatically generated matrix showing the compatibility between different backends and tasks, including selection guidance.
    • ADR Index: An index of Architecture Decision Records (ADRs) and their current status.
  10. Architecture Overview

    main

    The UniLab architecture is organized around several core pillars that define how the system operates, how ownership is distributed, and how components are registered. To understand the system, you should explore the following conceptual areas:

    • Runtime Model: Understand the Runner lifecycle, the separation between worker and learner roles, and how data flows through the system.
    • Layer Boundaries: Learn the ownership rules for different layers, specifically scripts, envs, backends, and algorithms.
    • Scene Composition: Understand how scene fragments, resources, and cold-path materialization work together to build simulation environments.
    • Registry: Learn how the Bootstrap import mechanism works and how to register env and backend components.
  11. Navigate the UniLab User Guide

    main

    The UniLab User Guide provides reference material for daily use after installation. If you are configuring UniLab for the first time, you should first read the Getting Started guide.

    The guide is organized into the following functional areas:

    • Training: Covers CLI routing, Hydra owner YAML configuration, logging, checkpoints, Docker usage, and multi-GPU setups.
    • Algorithms: Provides comparisons and details for PPO, APPO, SAC, TD3, FlashSAC, MLX PPO, HIM-PPO, and HORA.
    • Backends: Guidance on selecting between MuJoCo or Motrix based on owner YAML and backend capability evidence.
    • Tasks: Information on finding tasks related to motion control, action tracking, manipulation, and mobile manipulation.
    • Domain Randomization: Instructions for setting up reset, init, and interval randomization via task owner configurations.
    • Tooling: Details on exporting ONNX, checking for NaNs, sending W&B logs, and exporting scenes.
    • Manipulation: Specific instructions for tasks involving Allegro, Sharpa, and Go2+Airbot.
    • Terrain: (Included in the table of contents) for terrain-related configurations.
  12. Explore manipulation environments in unilab.envs.manipulation

    main

    The unilab.envs.manipulation module provides specialized environments for robotic manipulation tasks. It includes specific implementations for in-hand manipulation and Stewart platform-based tasks.

    Available environments include:

    • unilab.envs.manipulation.allegro_inhand
    • unilab.envs.manipulation.sharpa_inhand
    • unilab.envs.manipulation.stewart