mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-30 18:31:25 +00:00
a95b15ccc0
* refactor(env): introduce explicit gym ID handling in EnvConfig/factory
This commit introduces properties for the gym package/ID associated
with and environment config. They default to the current defaults
(`gym_{package_name}/{task_id}`) to avoid breaking changes, but allow
for easier use of external gym environments.
Subclasses of `EnvConfig` can override the default properties to allow
the factory to import (i.e. register) the gym env from a specific module,
and also instantiate the env from any ID string.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* more changes
* quality
* fix test
---------
Co-authored-by: Ben Sprenger <ben.sprenger@rogers.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com>
112 lines
4.0 KiB
Python
112 lines
4.0 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import importlib
|
|
|
|
import gymnasium as gym
|
|
from gymnasium.envs.registration import registry as gym_registry
|
|
|
|
from lerobot.envs.configs import AlohaEnv, EnvConfig, LiberoEnv, PushtEnv
|
|
|
|
|
|
def make_env_config(env_type: str, **kwargs) -> EnvConfig:
|
|
if env_type == "aloha":
|
|
return AlohaEnv(**kwargs)
|
|
elif env_type == "pusht":
|
|
return PushtEnv(**kwargs)
|
|
elif env_type == "libero":
|
|
return LiberoEnv(**kwargs)
|
|
else:
|
|
raise ValueError(f"Policy type '{env_type}' is not available.")
|
|
|
|
|
|
def make_env(
|
|
cfg: EnvConfig, n_envs: int = 1, use_async_envs: bool = False
|
|
) -> dict[str, dict[int, gym.vector.VectorEnv]]:
|
|
"""Makes a gym vector environment according to the config.
|
|
|
|
Args:
|
|
cfg (EnvConfig): the config of the environment to instantiate.
|
|
n_envs (int, optional): The number of parallelized env to return. Defaults to 1.
|
|
use_async_envs (bool, optional): Whether to return an AsyncVectorEnv or a SyncVectorEnv. Defaults to
|
|
False.
|
|
|
|
Raises:
|
|
ValueError: if n_envs < 1
|
|
ModuleNotFoundError: If the requested env package is not installed
|
|
|
|
Returns:
|
|
dict[str, dict[int, gym.vector.VectorEnv]]:
|
|
A mapping from suite name to indexed vectorized environments.
|
|
- For multi-task benchmarks (e.g., LIBERO): one entry per suite, and one vec env per task_id.
|
|
- For single-task environments: a single suite entry (cfg.type) with task_id=0.
|
|
|
|
"""
|
|
if n_envs < 1:
|
|
raise ValueError("`n_envs` must be at least 1")
|
|
|
|
env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
|
|
|
|
if "libero" in cfg.type:
|
|
from lerobot.envs.libero import create_libero_envs
|
|
|
|
if cfg.task is None:
|
|
raise ValueError("LiberoEnv requires a task to be specified")
|
|
|
|
return create_libero_envs(
|
|
task=cfg.task,
|
|
n_envs=n_envs,
|
|
camera_name=cfg.camera_name,
|
|
init_states=cfg.init_states,
|
|
gym_kwargs=cfg.gym_kwargs,
|
|
env_cls=env_cls,
|
|
)
|
|
elif "metaworld" in cfg.type:
|
|
from lerobot.envs.metaworld import create_metaworld_envs
|
|
|
|
if cfg.task is None:
|
|
raise ValueError("MetaWorld requires a task to be specified")
|
|
|
|
return create_metaworld_envs(
|
|
task=cfg.task,
|
|
n_envs=n_envs,
|
|
gym_kwargs=cfg.gym_kwargs,
|
|
env_cls=env_cls,
|
|
)
|
|
|
|
if cfg.gym_id not in gym_registry:
|
|
print(f"gym id '{cfg.gym_id}' not found, attempting to import '{cfg.package_name}'...")
|
|
try:
|
|
importlib.import_module(cfg.package_name)
|
|
except ModuleNotFoundError as e:
|
|
raise ModuleNotFoundError(
|
|
f"Package '{cfg.package_name}' required for env '{cfg.type}' not found. "
|
|
f"Please install it or check PYTHONPATH."
|
|
) from e
|
|
|
|
if cfg.gym_id not in gym_registry:
|
|
raise gym.error.NameNotFound(
|
|
f"Environment '{cfg.gym_id}' not registered even after importing '{cfg.package_name}'."
|
|
)
|
|
|
|
def _make_one():
|
|
return gym.make(cfg.gym_id, disable_env_checker=cfg.disable_env_checker, **(cfg.gym_kwargs or {}))
|
|
|
|
vec = env_cls([_make_one for _ in range(n_envs)], autoreset_mode=gym.vector.AutoresetMode.SAME_STEP)
|
|
|
|
# normalize to {suite: {task_id: vec_env}} for consistency
|
|
suite_name = cfg.type # e.g., "pusht", "aloha"
|
|
return {suite_name: {0: vec}}
|