# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# 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.
# ----------------------------------------------------------------------------
"""Gym environment registration and task-package discovery utilities."""
from __future__ import annotations
import importlib
import importlib.metadata
import importlib.util
import json
import logging
import os
import sys
import torch
from copy import deepcopy
from functools import partial
from typing import TYPE_CHECKING, Dict, Type
import gymnasium as gym
from gymnasium.envs.registration import EnvSpec as GymEnvSpec
from gymnasium.envs.registration import WrapperSpec
from dexsim.utility import log_warning
if TYPE_CHECKING:
from embodichain.lab.gym.envs import BaseEnv, EmbodiedEnvCfg
from embodichain.lab.task_program.integrations import (
TaskProgramAdapterFactory,
SimulationTaskProgramRegistration,
)
_logger = logging.getLogger(__name__)
__all__ = [
"EnvSpec",
"REGISTERED_ENVS",
"TimeLimitWrapper",
"build_env",
"discover_task_packages",
"execute_init_hooks",
"get_env_spec",
"make",
"make_vec",
"register",
"register_env",
"register_env_function",
]
[docs]
class EnvSpec:
[docs]
def __init__(
self,
uid: str,
cls: Type[BaseEnv],
max_episode_steps=None,
default_kwargs: dict = None,
task_program_registration: SimulationTaskProgramRegistration | None = None,
task_program_adapter_factory: TaskProgramAdapterFactory | None = None,
supports_rl: bool = False,
):
"""A specification for a Embodied environment."""
if type(supports_rl) is not bool:
raise TypeError("supports_rl must be a bool.")
if task_program_registration is not None:
from embodichain.lab.task_program.integrations import (
SimulationTaskProgramRegistration,
)
if type(task_program_registration) is not SimulationTaskProgramRegistration:
raise TypeError(
"task_program_registration must be exactly "
"SimulationTaskProgramRegistration or None."
)
if task_program_adapter_factory is not None:
from embodichain.lab.task_program.integrations import (
TaskProgramAdapterFactory,
)
if not isinstance(
task_program_adapter_factory,
TaskProgramAdapterFactory,
):
raise TypeError(
"task_program_adapter_factory must implement "
"TaskProgramAdapterFactory or be None."
)
factory_registration = getattr(
task_program_adapter_factory,
"registration",
None,
)
if factory_registration is not None:
from embodichain.lab.task_program.integrations import (
SimulationTaskProgramRegistration,
)
if type(factory_registration) is not SimulationTaskProgramRegistration:
raise TypeError(
"task_program_adapter_factory.registration must be "
"exactly SimulationTaskProgramRegistration."
)
if (
task_program_registration is not None
and factory_registration is not None
and factory_registration is not task_program_registration
):
raise ValueError(
"task_program_adapter_factory must own the exact "
"task_program_registration."
)
if task_program_registration is None:
task_program_registration = factory_registration
self.uid = uid
self.cls = cls
self.max_episode_steps = max_episode_steps
self.default_kwargs = {} if default_kwargs is None else default_kwargs
self.supports_rl = supports_rl
self.task_program_registration = task_program_registration
self.task_program_adapter_factory = task_program_adapter_factory
def make(self, **kwargs):
_kwargs = self.default_kwargs.copy()
_kwargs.update(kwargs)
if self.task_program_adapter_factory is not None:
supplied_factory = _kwargs.get("task_program_adapter_factory")
if (
supplied_factory is not None
and supplied_factory is not self.task_program_adapter_factory
):
raise ValueError(
"A registered Task Program adapter factory cannot be "
"overridden at environment construction."
)
_kwargs["task_program_adapter_factory"] = self.task_program_adapter_factory
return self.cls(**_kwargs)
@property
def gym_spec(self):
"""Return a gym EnvSpec for this env"""
entry_point = self.cls.__module__ + ":" + self.cls.__name__
return GymEnvSpec(
self.uid,
entry_point,
max_episode_steps=self.max_episode_steps,
kwargs=self.default_kwargs,
)
REGISTERED_ENVS: Dict[str, EnvSpec] = {}
[docs]
def register(
name: str,
cls: Type[BaseEnv],
max_episode_steps=None,
default_kwargs: dict = None,
task_program_registration: SimulationTaskProgramRegistration | None = None,
task_program_adapter_factory: TaskProgramAdapterFactory | None = None,
supports_rl: bool = False,
):
"""Register a Embodied environment."""
# hacky way to avoid circular import errors when users inherit a task in DexSim and try to register it themselves
from embodichain.lab.gym.envs import BaseEnv, BaseEnv
if name in REGISTERED_ENVS:
log_warning(f"Env {name} already registered")
if not (issubclass(cls, BaseEnv) or issubclass(cls, BaseEnv)):
raise TypeError(f"Env {name} must inherit from BaseEnv or BaseEnv")
REGISTERED_ENVS[name] = EnvSpec(
name,
cls,
max_episode_steps=max_episode_steps,
default_kwargs=default_kwargs,
supports_rl=supports_rl,
task_program_registration=task_program_registration,
task_program_adapter_factory=task_program_adapter_factory,
)
[docs]
class TimeLimitWrapper(gym.Wrapper):
"""like the standard gymnasium timelimit wrapper but fixes truncated variable to be a batched array"""
[docs]
def __init__(self, env: gym.Env, max_episode_steps: int):
super().__init__(env)
prev_frame_locals = sys._getframe(1).f_locals
frame = sys._getframe(1)
# check for user supplied max_episode_steps during gym.make calls
if frame.f_code.co_name == "make" and "max_episode_steps" in prev_frame_locals:
if prev_frame_locals["max_episode_steps"] is not None:
max_episode_steps = prev_frame_locals["max_episode_steps"]
# do some wrapper surgery to remove the previous timelimit wrapper
# with gymnasium 0.29.1, this will remove the timelimit wrapper and nothing else.
curr_env = env
while curr_env is not None:
if isinstance(curr_env, gym.wrappers.TimeLimit):
self.env = curr_env.env
break
self._max_episode_steps = self.base_env.max_episode_steps
@property
def base_env(self) -> BaseEnv:
return self.env.unwrapped
@property
def device(self) -> torch.device:
return self.base_env.device
@property
def num_envs(self) -> int:
return self.base_env.num_envs
[docs]
def step(self, action):
observation, reward, terminated, truncated, info = self.env.step(action)
truncated = truncated | (self.base_env.elapsed_steps >= self._max_episode_steps)
return observation, reward, terminated, truncated, info
[docs]
def make(env_id, **kwargs):
"""Instantiate a Embodied environment.
Args:
env_id (str): Environment ID.
as_gym (bool, optional): Add TimeLimit wrapper as gym.
**kwargs: Keyword arguments to pass to the environment.
"""
if env_id not in REGISTERED_ENVS:
raise KeyError("Env {} not found in registry".format(env_id))
env_spec = REGISTERED_ENVS[env_id]
env = env_spec.make(**kwargs)
return env
[docs]
def get_env_spec(env_id: str) -> EnvSpec:
"""Return one registered environment specification or fail closed."""
if type(env_id) is not str or not env_id or env_id != env_id.strip():
raise ValueError("env_id must be a non-empty string without outer whitespace.")
try:
return REGISTERED_ENVS[env_id]
except KeyError as exc:
raise KeyError(f"Env {env_id!r} not found in registry.") from exc
[docs]
def build_env(env_id: str, base_env_cfg: EmbodiedEnvCfg):
"""Create an environment from a registered env id.
A thin convenience wrapper around :func:`make` that deep-copies the base
config so callers can safely mutate the resulting environment's cfg
without affecting shared defaults. This helper used to live in the task
package; it now lives with the registry so that core code paths such as RL
training do not need to depend on an official task package.
Args:
env_id: Registered environment id (see :func:`register_env`).
base_env_cfg: Base environment configuration to instantiate with.
Returns:
The instantiated environment.
"""
return make(env_id, cfg=deepcopy(base_env_cfg))
[docs]
def make_vec(env_id, **kwargs):
env = gym.make(env_id, **kwargs)
return env
[docs]
def register_env(
uid: str,
max_episode_steps=None,
override=False,
*,
supports_rl: bool = False,
task_program_registration: SimulationTaskProgramRegistration | None = None,
task_program_adapter_factory: TaskProgramAdapterFactory | None = None,
**kwargs,
):
"""A decorator to register Embodied environments.
Args:
uid (str): unique id of the environment.
max_episode_steps (int): maximum number of steps in an episode.
override (bool): whether to override the environment if it is already registered.
supports_rl: Whether the environment has a supported RL training path.
Notes:
- `max_episode_steps` is processed differently from other keyword arguments in gym.
`gym.make` wraps the env with `gym.wrappers.TimeLimit` to limit the maximum number of steps.
- `gym.EnvSpec` uses kwargs instead of **kwargs!
"""
try:
json.dumps(kwargs)
except TypeError:
raise RuntimeError(
f"You cannot register_env with non json dumpable kwargs, e.g. classes or types. If you really need to do this, it is recommended to create a mapping of string to the unjsonable data and to pass the string in the kwarg and during env creation find the data you need"
)
def _register_env(cls):
cls = register_env_function(
cls,
uid,
override,
max_episode_steps,
supports_rl=supports_rl,
task_program_registration=task_program_registration,
task_program_adapter_factory=task_program_adapter_factory,
**kwargs,
)
return cls
return _register_env
[docs]
def register_env_function(
cls,
uid,
override=False,
max_episode_steps=None,
*,
supports_rl: bool = False,
task_program_registration: SimulationTaskProgramRegistration | None = None,
task_program_adapter_factory: TaskProgramAdapterFactory | None = None,
**kwargs,
):
if uid in REGISTERED_ENVS:
if override:
from gymnasium.envs.registration import registry
log_warning(f"Override registered env {uid}")
REGISTERED_ENVS.pop(uid)
registry.pop(uid)
else:
log_warning(f"Env {uid} is already registered. Skip registration.")
return cls
register(
uid,
cls,
max_episode_steps=max_episode_steps,
default_kwargs=deepcopy(kwargs),
supports_rl=supports_rl,
task_program_registration=task_program_registration,
task_program_adapter_factory=task_program_adapter_factory,
)
# Register for gym
gym.register(
uid,
entry_point=partial(make, env_id=uid),
vector_entry_point=partial(make_vec, env_id=uid),
max_episode_steps=max_episode_steps,
disable_env_checker=True, # Temporary solution as we allow empty observation spaces
kwargs=deepcopy(kwargs),
)
return cls
def _import_task_package(ep: importlib.metadata.EntryPoint):
"""Import a task package and ensure its auto-registration runs.
Legacy editable installs from before official tasks were bundled in the
main distribution can expose two editable projects from the same checkout.
The repository-level ``embodichain_tasks/`` container may then be picked up
as a namespace package that shadows the real
``embodichain_tasks/embodichain_tasks/`` package.
A namespace package never executes ``__init__.py``, so the package's
``import_packages()`` call -- which triggers every ``@register_env`` --
is skipped and no environments are registered.
When such shadowing is detected (the imported module has no ``__file__``),
this helper locates the real ``__init__.py`` beneath one of the namespace's
search locations (``<location>/<top_level>/__init__.py``) and loads it
directly so registration runs. This is a no-op for regular, non-shadowed
installs.
Args:
ep: The ``embodichain.tasks`` entry point to import.
Returns:
The imported module (the real package when possible, otherwise the
namespace module).
"""
module_name = ep.value
top_level = module_name.partition(".")[0]
mod = importlib.import_module(module_name)
if getattr(mod, "__file__", None) is not None:
return mod
# Only top-level packages can be reliably force-loaded here; dotted entry
# point values fall back to the plain import above.
if module_name != top_level:
return mod
# Namespace shadowing: search each namespace location for the real
# package's __init__.py and load it in place of the namespace module.
for location in list(getattr(mod, "__path__", [])):
init_path = os.path.join(location, top_level, "__init__.py")
if not os.path.isfile(init_path):
continue
pkg_dir = os.path.dirname(init_path)
spec = importlib.util.spec_from_file_location(
top_level, init_path, submodule_search_locations=[pkg_dir]
)
if spec is None or spec.loader is None:
continue
real_mod = importlib.util.module_from_spec(spec)
sys.modules[top_level] = real_mod
spec.loader.exec_module(real_mod)
return real_mod
return mod
[docs]
def discover_task_packages() -> list[str]:
"""Import all registered task packages via ``embodichain.tasks`` entry_points.
Each task package recursively imports its task modules, which triggers
``@register_env`` → ``gym.register()``. After this call, all tasks from all
installed packages are available in gymnasium's global registry.
Returns:
List of entry point names that were successfully imported.
"""
imported: list[str] = []
try:
eps = importlib.metadata.entry_points(group="embodichain.tasks")
except TypeError:
# Python < 3.12: entry_points() requires a keyword argument
eps = importlib.metadata.entry_points().get("embodichain.tasks", [])
for ep in eps:
try:
_import_task_package(ep)
imported.append(ep.name)
except Exception:
_logger.warning(
f"Failed to import task package '{ep.name}' ({ep.value})",
exc_info=True,
)
return imported
[docs]
def execute_init_hooks() -> list[str]:
"""Execute all registered init hooks via ``embodichain.init`` entry_points.
Hooks are called in entry_points declaration order. An exception from one
hook does not prevent others from executing.
Each entry point value must be in the format ``"module.path:function_name"``.
The function must accept no arguments and return ``None``.
Returns:
List of hook names that were executed successfully.
"""
executed: list[str] = []
try:
eps = importlib.metadata.entry_points(group="embodichain.init")
except TypeError:
# Python < 3.12
eps = importlib.metadata.entry_points().get("embodichain.init", [])
for ep in eps:
try:
module_name, func_name = ep.value.split(":", 1)
module = importlib.import_module(module_name)
func = getattr(module, func_name)
func()
executed.append(ep.name)
except Exception:
_logger.warning(
f"Init hook '{ep.name}' ({ep.value}) failed",
exc_info=True,
)
return executed