Policy-Lightning is a PyTorch Lightning-based implementation of popular policy learning algorithms. It is specifically designed for embodied multi-agent manipulation tasks, offering clean abstractions for training, evaluation, and integration with simulators.
conda create -n policy-lt python=3.12
conda activate policy-lt
# Replace 'cu***' with your CUDA version, e.g., cu124
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu***
# Install project dependencies
pip install -r requirements.txtCurrently, we support data collection through RoboFactory, an automated multi-agent simulation and recording framework built on top of ManiSkill.
We plan to support additional simulators and datasets in the future — contributions are welcome! To convert your existing dataset into our supported format, please refer to the Data Conversion Guide.
We provide small-scale demo datasets for quick testing and validation. You can download them from hugging face.
2D Diffusion Policy:
python workspace --config-name=dp2 task=2a_lift_barrier3D Diffusion Policy:
python workspace --config-name=dp3 task=2a_lift_barrier_3dCustom Policy:
To integrate your own policy architecture:
-
Implement your custom policy in
./policy/. -
Create a corresponding configuration in
./config/.
Optional (for custom tasks):
-
Place your dataset in
./data/. -
Add a task configuration under
./config/task/.
python workspace --config-name=[custom_polic] task=[custom_task]