# ----------------------------------------------------------------------------
# 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.
# ----------------------------------------------------------------------------
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
_logger = logging.getLogger(__name__)
[docs]
class EnvSpec:
[docs]
def __init__(
self,
uid: str,
cls: Type[BaseEnv],
max_episode_steps=None,
default_kwargs: dict = None,
):
"""A specification for a Embodied environment."""
self.uid = uid
self.cls = cls
self.max_episode_steps = max_episode_steps
self.default_kwargs = {} if default_kwargs is None else default_kwargs
def make(self, **kwargs):
_kwargs = self.default_kwargs.copy()
_kwargs.update(kwargs)
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
):
"""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
)
[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
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 (``embodichain_tasks.rl``); it now lives with the registry so
that core code paths such as RL training do not need to depend on a task
package. ``embodichain_tasks.rl`` re-exports it for backward
compatibility.
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))
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, **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.
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, **kwargs)
return cls
return _register_env
def register_env_function(cls, uid, override=False, max_episode_steps=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),
)
# 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.
In editable ("development") installs, a task package that lives in a
subdirectory of another editable project can be shadowed by a same-named
namespace package. For example, when both ``embodichain`` and
``embodichain_tasks`` are installed editable from the same checkout, the
``embodichain`` editable finder exposes the repository root on the import
path; the ``embodichain_tasks/`` project container (which has no
``__init__.py`` at its root) is then 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
def discover_task_packages() -> list[str]:
"""Import all registered task packages via ``embodichain.tasks`` entry_points.
Each task package's ``__init__.py`` recursively imports its sub-packages,
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
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