forked from tangger/lerobot
1a1308d62f
add fixes for reproducibility only try to start env if it is closed revision fix normalization and data type Improve README Improve README Tests are passing, Eval pretrained model works, Add gif Update gif Update gif Update gif Update gif Update README Update README update minor Update README.md Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com> Update README.md Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com> Address suggestions Update thumbnail + stats Update thumbnail + stats Update README.md Co-authored-by: Alexander Soare <alexander.soare159@gmail.com> Add more comments Add test_examples.py
65 lines
2.3 KiB
Python
65 lines
2.3 KiB
Python
"""
|
|
This test verifies that all environments, datasets, policies listed in `lerobot/__init__.py` can be sucessfully
|
|
imported and that their class attributes (eg. `available_datasets`, `name`, `available_tasks`) corresponds.
|
|
|
|
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
|
|
"""
|
|
|
|
import pytest
|
|
import lerobot
|
|
|
|
from lerobot.common.envs.aloha.env import AlohaEnv
|
|
from lerobot.common.envs.pusht.env import PushtEnv
|
|
from lerobot.common.envs.simxarm import SimxarmEnv
|
|
|
|
from lerobot.common.datasets.simxarm import SimxarmDataset
|
|
from lerobot.common.datasets.aloha import AlohaDataset
|
|
from lerobot.common.datasets.pusht import PushtDataset
|
|
|
|
from lerobot.common.policies.act.policy import ActionChunkingTransformerPolicy
|
|
from lerobot.common.policies.diffusion.policy import DiffusionPolicy
|
|
from lerobot.common.policies.tdmpc.policy import TDMPCPolicy
|
|
|
|
|
|
def test_available():
|
|
pol_classes = [
|
|
ActionChunkingTransformerPolicy,
|
|
DiffusionPolicy,
|
|
TDMPCPolicy,
|
|
]
|
|
|
|
env_classes = [
|
|
AlohaEnv,
|
|
PushtEnv,
|
|
SimxarmEnv,
|
|
]
|
|
|
|
dat_classes = [
|
|
AlohaDataset,
|
|
PushtDataset,
|
|
SimxarmDataset,
|
|
]
|
|
|
|
policies = [pol_cls.name for pol_cls in pol_classes]
|
|
assert set(policies) == set(lerobot.available_policies)
|
|
|
|
envs = [env_cls.name for env_cls in env_classes]
|
|
assert set(envs) == set(lerobot.available_envs)
|
|
|
|
tasks_per_env = {env_cls.name: env_cls.available_tasks for env_cls in env_classes}
|
|
for env in envs:
|
|
assert set(tasks_per_env[env]) == set(lerobot.available_tasks_per_env[env])
|
|
|
|
datasets_per_env = {env_cls.name: dat_cls.available_datasets for env_cls, dat_cls in zip(env_classes, dat_classes)}
|
|
for env in envs:
|
|
assert set(datasets_per_env[env]) == set(lerobot.available_datasets_per_env[env])
|
|
|
|
|