54 lines
2.0 KiB
Python
54 lines
2.0 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`) are valid.
|
|
|
|
When implementing a new dataset (e.g. `AlohaDataset`), policy (e.g. `DiffusionPolicy`), or environment, follow these steps:
|
|
- Set the required class attributes: `available_datasets`.
|
|
- Set the required class attributes: `name`.
|
|
- Update variables in `lerobot/__init__.py` (e.g. `available_envs`, `available_datasets_per_envs`, `available_policies`)
|
|
- Update variables in `tests/test_available.py` by importing your new class
|
|
"""
|
|
|
|
import importlib
|
|
import pytest
|
|
import lerobot
|
|
import gymnasium as gym
|
|
|
|
from lerobot.common.datasets.xarm import XarmDataset
|
|
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():
|
|
policy_classes = [
|
|
ActionChunkingTransformerPolicy,
|
|
DiffusionPolicy,
|
|
TDMPCPolicy,
|
|
]
|
|
|
|
dataset_class_per_env = {
|
|
"aloha": AlohaDataset,
|
|
"pusht": PushtDataset,
|
|
"xarm": XarmDataset,
|
|
}
|
|
|
|
policies = [pol_cls.name for pol_cls in policy_classes]
|
|
assert set(policies) == set(lerobot.available_policies), policies
|
|
|
|
for env_name in lerobot.available_envs:
|
|
for task_name in lerobot.available_tasks_per_env[env_name]:
|
|
package_name = f"gym_{env_name}"
|
|
importlib.import_module(package_name)
|
|
gym_handle = f"{package_name}/{task_name}"
|
|
assert gym_handle in gym.envs.registry.keys(), gym_handle
|
|
|
|
dataset_class = dataset_class_per_env[env_name]
|
|
available_datasets = lerobot.available_datasets_per_env[env_name]
|
|
assert set(available_datasets) == set(dataset_class.available_datasets), f"{env_name=} {available_datasets=}"
|
|
|
|
|