# ----------------------------------------------------------------------------
# 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.
# ----------------------------------------------------------------------------
"""Shared environment and Policy construction for RL training and evaluation."""
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from gymnasium import spaces
import torch
from embodichain.lab.gym.envs.managers.cfg import EventCfg
from embodichain.lab.gym.utils.gym_utils import config_to_cfg, get_manager_modules
from embodichain.lab.gym.utils.profiler import EnvProfilerCfg
from embodichain.lab.gym.utils.registration import build_env
from embodichain.lab.sim import SimulationManagerCfg
from embodichain.lab.sim.cfg import RenderCfg
from embodichain.learning.rl.env import build_learning_env
from embodichain.learning.models import build_mlp_from_cfg
from embodichain.learning.rl.policies import build_policy, resolve_policy_obs_groups
from embodichain.learning.rl.utils import (
dict_to_tensordict,
flatten_observation_groups,
)
from embodichain.utils.utility import load_config
__all__ = [
"PolicyRuntime",
"build_gym_policy_runtime",
"build_learning_policy_runtime",
]
@dataclass(frozen=True)
class _GymEnvironmentRuntime:
"""A simulator task reconstructed from one training configuration."""
env: Any
env_id: str
env_cfg: Any
gym_config: dict[str, Any]
gym_config_path: Path
[docs]
@dataclass(frozen=True)
class PolicyRuntime:
"""An Environment and Policy reconstructed from one training configuration."""
env: Any
policy: torch.nn.Module
device: torch.device
env_id: str
env_cfg: Any | None = None
gym_config: dict[str, Any] | None = None
gym_config_path: Path | None = None
[docs]
def close(self) -> None:
"""Close the Environment without terminating the current process."""
_close_environment(self.env)
def _resolve_config_reference(
value: str | Path,
*,
base_dir: str | Path | None = None,
) -> Path:
"""Resolve a referenced config relative to its containing config file."""
path = Path(value).expanduser()
if path.is_absolute():
return path
if base_dir is not None:
candidate = Path(base_dir).expanduser().resolve() / path
if candidate.exists():
return candidate
return path
def _build_learning_environment(
config: dict[str, Any],
*,
device: torch.device,
num_envs: int,
) -> tuple[str, Any]:
"""Build the lightweight Environment declared by a training config."""
env_block = config["trainer"]["learning_env"]
if isinstance(env_block, str):
env_name = env_block
env_config: dict[str, Any] = {}
else:
env_name = env_block["name"]
env_config = dict(env_block.get("cfg", {}))
return str(env_name), build_learning_env(
str(env_name),
num_envs=num_envs,
device=device,
**env_config,
)
def _observation_space_dim(
observation_space: spaces.Space,
groups: tuple[str, ...] | None,
) -> int:
"""Return the flattened size of selected observation groups."""
if groups is None:
return int(spaces.utils.flatdim(observation_space))
total = 0
for group in groups:
selected_space = observation_space
for key in group.split("."):
if not isinstance(selected_space, spaces.Dict):
raise TypeError(
f"Observation group '{group}' requires a Dict space at '{key}'."
)
if key not in selected_space.spaces:
raise KeyError(
f"Observation group '{group}' was not found in the environment space."
)
selected_space = selected_space.spaces[key]
total += int(spaces.utils.flatdim(selected_space))
return total
[docs]
def build_learning_policy_runtime(
config: dict[str, Any],
*,
device: torch.device,
num_envs: int,
) -> PolicyRuntime:
"""Build a lightweight Environment and its configured Policy."""
env_name, env = _build_learning_environment(
config,
device=device,
num_envs=num_envs,
)
try:
policy = _build_learning_policy(config["policy"], env, device)
except Exception:
env.close()
raise
return PolicyRuntime(env, policy, device, env_name)
def _trainer_events_need_renderer(events: Sequence[EventCfg]) -> bool:
"""Inspect active, resolved event classes without constructing functors."""
from embodichain.lab.gym.envs.managers.record import (
record_camera_data,
validation_cameras,
)
for event in events:
func = event.func
if not isinstance(func, type):
continue
if issubclass(func, record_camera_data):
return True
if issubclass(func, validation_cameras) and event.params.get("cameras", []):
return True
return False
def _build_gym_environment(
config: dict[str, Any],
*,
simulation_device: torch.device,
num_envs: int | None,
headless: bool,
renderer: str | None,
gpu_id: int,
seed: int | None = None,
config_dir: str | Path | None = None,
profiler: EnvProfilerCfg | None = None,
trainer_events: Sequence[EventCfg] = (),
) -> _GymEnvironmentRuntime:
"""Build the simulator Environment declared by a training config."""
trainer_cfg = config["trainer"]
gym_config_path = _resolve_config_reference(
trainer_cfg["gym_config"],
base_dir=config_dir,
)
gym_config = load_config(gym_config_path)
env_cfg = config_to_cfg(gym_config, manager_modules=get_manager_modules())
if num_envs is not None:
env_cfg.num_envs = int(num_envs)
if seed is not None:
env_cfg.seed = int(seed)
if env_cfg.sim_cfg is None:
env_cfg.sim_cfg = SimulationManagerCfg()
env_cfg.sim_cfg.sim_device = simulation_device
env_cfg.sim_cfg.headless = headless
if renderer is not None:
env_cfg.sim_cfg.render_cfg = RenderCfg(renderer=renderer)
env_cfg.sim_cfg.gpu_id = (
simulation_device.index
if simulation_device.type == "cuda" and simulation_device.index is not None
else gpu_id
)
if env_cfg.sim_cfg.render_cfg.renderer == "auto" and _trainer_events_need_renderer(
trainer_events
):
from embodichain.lab.sim import SimulationManager, cfg as sim_cfg
from embodichain.lab.sim.utility.render_utils import select_default_renderer
process_renderer = SimulationManager._get_process_renderer()
env_cfg.sim_cfg.render_cfg.renderer = (
process_renderer
if sim_cfg.DEFAULT_RENDERER == "auto"
and process_renderer not in (None, "no-render")
else select_default_renderer(env_cfg.sim_cfg.gpu_id)
)
env_cfg.profiler = profiler
env = build_env(gym_config["id"], base_env_cfg=env_cfg)
return _GymEnvironmentRuntime(
env=env,
env_id=str(gym_config["id"]),
env_cfg=env_cfg,
gym_config=gym_config,
gym_config_path=gym_config_path.resolve(),
)
[docs]
def build_gym_policy_runtime(
config: dict[str, Any],
*,
device: torch.device,
num_envs: int | None,
headless: bool,
renderer: str | None,
gpu_id: int,
seed: int | None = None,
config_dir: str | Path | None = None,
profiler: EnvProfilerCfg | None = None,
simulation_device: torch.device | None = None,
trainer_events: Sequence[EventCfg] = (),
) -> PolicyRuntime:
"""Build a simulator task and the Policy declared by its training config.
Args:
config: Training configuration containing trainer and policy blocks.
device: Policy tensor device.
num_envs: Optional override for the number of environments.
headless: Whether to suppress the native window.
renderer: Explicit renderer override; None preserves the task's render
configuration, including demand-based auto selection.
gpu_id: Renderer GPU index when the simulation device has no CUDA index.
seed: Optional environment seed override.
config_dir: Directory used to resolve the referenced task configuration.
profiler: Optional environment profiler configuration.
simulation_device: Optional simulation device distinct from the policy.
trainer_events: Active training and evaluation events, resolved before
the first World so their cameras can reserve the process renderer.
Returns:
The constructed environment, policy and resolved task configuration.
"""
task = _build_gym_environment(
config,
simulation_device=simulation_device or device,
num_envs=num_envs,
headless=headless,
renderer=renderer,
gpu_id=gpu_id,
seed=seed,
config_dir=config_dir,
profiler=profiler,
trainer_events=trainer_events,
)
env = task.env
try:
sample_observation, _ = env.reset()
sample_observation_td = dict_to_tensordict(sample_observation, device)
actor_obs_groups, critic_obs_groups = resolve_policy_obs_groups(
config["policy"]
)
actor_observation_dim = int(
flatten_observation_groups(
sample_observation_td,
actor_obs_groups,
).shape[-1]
)
uses_separate_critic_obs = actor_obs_groups != critic_obs_groups
critic_observation_dim = (
int(
flatten_observation_groups(
sample_observation_td,
critic_obs_groups,
).shape[-1]
)
if uses_separate_critic_obs
else actor_observation_dim
)
environment_action_dim = int(env.action_space.shape[-1])
policy = _build_gym_policy(
config["policy"],
env=env,
device=device,
actor_observation_dim=actor_observation_dim,
critic_observation_dim=critic_observation_dim,
uses_separate_critic_obs=uses_separate_critic_obs,
action_dim=environment_action_dim,
)
except Exception:
_close_environment(env)
raise
return PolicyRuntime(
env=env,
policy=policy,
device=device,
env_id=task.env_id,
env_cfg=task.env_cfg,
gym_config=task.gym_config,
gym_config_path=task.gym_config_path,
)
def _build_gym_policy(
policy_block: dict[str, Any],
*,
env: Any,
device: torch.device,
actor_observation_dim: int,
critic_observation_dim: int,
uses_separate_critic_obs: bool,
action_dim: int,
) -> torch.nn.Module:
configured_action_dim = int(policy_block.get("action_dim", action_dim))
if configured_action_dim != action_dim:
raise ValueError(
f"Configured policy.action_dim={configured_action_dim} does not match "
f"env action dim {action_dim}."
)
policy_name = str(policy_block["name"]).lower()
if policy_name == "actor_critic":
actor_cfg = policy_block.get("actor")
critic_cfg = policy_block.get("critic")
if actor_cfg is None or critic_cfg is None:
raise ValueError(
"ActorCritic requires policy.actor and policy.critic definitions."
)
return build_policy(
policy_block,
actor_observation_dim,
env.action_space,
device,
actor=build_mlp_from_cfg(actor_cfg, actor_observation_dim, action_dim),
critic=build_mlp_from_cfg(critic_cfg, critic_observation_dim, 1),
critic_obs_space=(
critic_observation_dim if uses_separate_critic_obs else None
),
)
if policy_name == "actor_only":
actor_cfg = policy_block.get("actor")
if actor_cfg is None:
raise ValueError("ActorOnly requires a policy.actor definition.")
return build_policy(
policy_block,
actor_observation_dim,
env.action_space,
device,
actor=build_mlp_from_cfg(actor_cfg, actor_observation_dim, action_dim),
)
return build_policy(
policy_block,
env.observation_space,
env.action_space,
device,
)
def _build_learning_policy(
policy_block: dict[str, Any],
env: Any,
device: torch.device,
) -> torch.nn.Module:
actor_obs_groups, critic_obs_groups = resolve_policy_obs_groups(policy_block)
actor_observation_dim = _observation_space_dim(
env.single_observation_space,
actor_obs_groups,
)
uses_separate_critic_obs = actor_obs_groups != critic_obs_groups
critic_observation_dim = (
_observation_space_dim(env.single_observation_space, critic_obs_groups)
if uses_separate_critic_obs
else actor_observation_dim
)
action_dim = int(env.single_action_space.shape[-1])
actor_cfg = policy_block.get("actor")
critic_cfg = policy_block.get("critic")
policy = build_policy(
policy_block,
(
actor_observation_dim
if policy_block["name"].lower() in {"actor_critic", "actor_only"}
else env.single_observation_space
),
env.single_action_space,
device,
actor=(
build_mlp_from_cfg(actor_cfg, actor_observation_dim, action_dim)
if actor_cfg is not None
else None
),
critic=(
build_mlp_from_cfg(critic_cfg, critic_observation_dim, 1)
if critic_cfg is not None
else None
),
critic_obs_space=(critic_observation_dim if uses_separate_critic_obs else None),
)
if "initial_log_std" in policy_block and hasattr(policy, "log_std"):
with torch.no_grad():
policy.log_std.fill_(float(policy_block["initial_log_std"]))
return policy
def _close_environment(env: Any) -> None:
target = getattr(env, "unwrapped", env)
if getattr(target, "sim", None) is not None:
target.close(exit_process=False)
else:
env.close()