seohongpark / ocbench

OCBench 
 Paper   Project page   Blog post 
# Overview
OCBench is a controllable robotic manipulation benchmark for studying behavioral cloning (BC) and reinforcement learning (RL).
OCBench provides GPU-accelerated tasks and scripted policies that mimic the properties of human demonstrations.
This enables controlled studies of BC and RL algorithms with virtually unlimited data.
# Quick start
### Installation
OCBench can be easily installed via PyPI:
 [code] 
The additional `mjwarp` and `train` dependencies are required for the GPU-accelerated (MJWarp) environments and reference implementations, respectively.
If you only want to use the CPU environments, you can install OCBench without the optional dependencies:
 [code] 
### Usage
OCBench environments follow the Gymnasium interface.
Here's an example of instantiating and running a CPU environment:
 [code] 
OCBench defaults to GPU-accelerated (MJWarp) environments.
Here's an example of running 1024 parallel environments:
 [code] 
# Tasks and datasets
### Tasks
OCBench provides 28 manipulation tasks across 5 types of environments.
See the paper for more details.
| Environment | Tasks |
| --- | --- |
| `block-single` | `task1` |
| `block-double` | `task1`, `task2` |
| `block-triple` | `task1`, `task2` |
| `block-quadruple` | `task1`, `task2`, `task3` |
| `chamber-easy` | `task1`, `task2`, `task3` |
| `chamber-medium` | `task1`, `task2` |
| `chamber-hard` | `task1`, `task2` |
| `switch-3x3` | `task1`, `task2` |
| `switch-4x4` | `task1`, `task2` |
| `switch-5x5` | `task1`, `task2` |
| `hanoi-single` | `task1`, `task2` |
| `hanoi-double` | `task1`, `task2` |
| `hanoi-triple` | `task1`, `task2` |
| `bowling` | `task1` |
OCBench also provides several variants of the environments:
* `lite`: The `lite` variant has much shorter horizons and simplified scripted policies. This can be useful for faster prototyping.
* `cpu`: The `cpu` variant uses CPU-based MuJoCo instead of GPU-based MJWarp.
* `visual`: The `visual` variant provides pixel-based observations instead of state-based observations.
| Environment name | `lite` | `cpu` | `visual` |
| --- | --- | --- | --- |
| `block-single-task1-v0` | No | No | No |
| `block-lite-single-task1-v0` | Yes | No | No |
| `block-cpu-single-task1-v0` | No | Yes | No |
| `block-cpu-lite-single-task1-v0` | Yes | Yes | No |
| `visual-block-single-task1-v0` | No | No | Yes |
| `visual-block-lite-single-task1-v0` | Yes | No | Yes |
| `visual-block-cpu-single-task1-v0` | No | Yes | Yes |
| `visual-block-cpu-lite-single-task1-v0` | Yes | Yes | Yes |
### Datasets
OCBench datasets are hosted on Hugging Face.
The datasets will be automatically downloaded to `~/.cache/ocbench/datasets` by default.
You can also specify a dataset root directory or path by setting `dataset_path` or `dataset_root`.
For `lite`, `cpu`, and `visual` variants, we only provide datasets for selected tasks.
In particular, we do not provide datasets for full-horizon `cpu` environments.
However, you can still gener