60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
"""
|
|
This file contains lists of available environments, dataset and policies to reflect the current state of LeRobot library.
|
|
We do not want to import all the dependencies, but instead we keep it lightweight to ensure fast access to these variables.
|
|
|
|
Example:
|
|
```python
|
|
import lerobot
|
|
print(lerobot.available_envs)
|
|
print(lerobot.available_tasks_per_env)
|
|
print(lerobot.available_datasets_per_env)
|
|
print(lerobot.available_datasets)
|
|
print(lerobot.available_policies)
|
|
```
|
|
|
|
Note:
|
|
When implementing a concrete class (e.g. `AlohaDataset`, `PushtEnv`, `DiffusionPolicy`), you need to:
|
|
1. set the required class attributes:
|
|
- for classes inheriting from `AbstractDataset`: `available_datasets`
|
|
- for classes inheriting from `AbstractEnv`: `name`, `available_tasks`
|
|
- for classes inheriting from `AbstractPolicy`: `name`
|
|
2. update variables in `lerobot/__init__.py` (e.g. `available_envs`, `available_datasets_per_envs`, `available_policies`)
|
|
3. update variables in `tests/test_available.py` by importing your new class
|
|
"""
|
|
|
|
from lerobot.__version__ import __version__ # noqa: F401
|
|
|
|
available_envs = [
|
|
"aloha",
|
|
"pusht",
|
|
"simxarm",
|
|
]
|
|
|
|
available_tasks_per_env = {
|
|
"aloha": [
|
|
"sim_insertion",
|
|
"sim_transfer_cube",
|
|
],
|
|
"pusht": ["pusht"],
|
|
"simxarm": ["lift"],
|
|
}
|
|
|
|
available_datasets_per_env = {
|
|
"aloha": [
|
|
"aloha_sim_insertion_human",
|
|
"aloha_sim_insertion_scripted",
|
|
"aloha_sim_transfer_cube_human",
|
|
"aloha_sim_transfer_cube_scripted",
|
|
],
|
|
"pusht": ["pusht"],
|
|
"simxarm": ["xarm_lift_medium"],
|
|
}
|
|
|
|
available_datasets = [dataset for env in available_envs for dataset in available_datasets_per_env[env]]
|
|
|
|
available_policies = [
|
|
"act",
|
|
"diffusion",
|
|
"tdmpc",
|
|
]
|