注册环境

注册环境#

上一教程介绍了如何创建自定义倒立摆环境。我们根据配置中 class_type 指定的类，手动创建了环境实例。

虽然直接，但这种方法不够可扩展，因为我们有大量的环境套件。在本教程中，我们将展示如何使用 gymnasium.register() 方法将环境注册到 gymnasium 注册表中。这使我们可以通过 gymnasium.make() 函数创建环境。

代码#

本教程对应 scripts/environments 目录中的 random_agent.py 脚本。它只是一个轻量入口，调用 isaaclab_rl.entrypoints 模块中的实际实现。

simple_agents.py 代码

# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""Checkpoint-free playback workflows for Isaac Lab environments.

The zero and random agents are variations of playback that need no trained checkpoint:
the policy either infers finite zero or hold actions or samples uniform random actions.
"""

from __future__ import annotations

import argparse
import contextlib
import logging
import os
import sys
from collections.abc import Callable
from typing import Any, Literal

import gymnasium as gym
import torch

from isaaclab.app import add_launcher_args, launch_simulation
from isaaclab.envs import DirectMARLEnvCfg
from isaaclab.envs.utils.spaces import sample_space
from isaaclab.utils import math as math_utils
from isaaclab.utils import validate

import isaaclab_tasks  # noqa: F401
from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli

from .common import (
    add_video_args,
    apply_env_overrides,
    apply_video_recording,
    close_env,
    enable_cameras_for_video,
    normalize_task_name,
    pre_launch_video_config,
    video_playback_steps,
)

logger = logging.getLogger(__name__)

# PLACEHOLDER: Extension template (do not remove this comment)
with contextlib.suppress(ImportError):
    import isaaclab_tasks_experimental  # noqa: F401

PolicyName = Literal["zero", "random"]
"""Action policies supported by the checkpoint-free agents."""

_DESCRIPTIONS: dict[str, str] = {
    "zero": "Zero agent for Isaac Lab environments.",
    "random": "Random agent for Isaac Lab environments.",
}

_SEED = 42


def run(argv: list[str] | None = None, *, policy: PolicyName) -> None:
    """Run an Isaac Lab environment with a checkpoint-free policy.

    Args:
        argv: Command-line arguments excluding the executable name. Reads ``sys.argv`` when omitted.
        policy: Action policy to apply, either inferred zero actions or uniform random actions.

    Raises:
        ValueError: If the requested policy is not supported.
    """
    if policy not in _DESCRIPTIONS:
        raise ValueError(f"Unsupported policy {policy!r}. Expected one of: {sorted(_DESCRIPTIONS)}.")

    args_cli = _parse_args(argv, policy)
    torch.manual_seed(_SEED)

    env_cfg, _ = resolve_task_config(args_cli.task, "")
    apply_env_overrides(args_cli, env_cfg)
    pre_launch_video_config(env_cfg, args_cli)
    # reject unsupported configurations before launching Kit or initializing a native physics backend
    try:
        validate(env_cfg)
    except (TypeError, ValueError) as exc:
        raise SystemExit(f"Invalid environment configuration: {exc}") from None

    with launch_simulation(env_cfg, args_cli), contextlib.ExitStack() as cleanup:
        log_dir = os.path.abspath(os.path.join("logs", f"{policy}_agent", normalize_task_name(args_cli.task)))
        apply_video_recording(env_cfg, log_dir, args_cli, subdir="play")
        env = gym.make(args_cli.task, cfg=env_cfg)
        cleanup.callback(lambda: close_env(env))
        logger.info(f"Gym observation space: {env.observation_space}")
        logger.info(f"Gym action space: {env.action_space}")
        env.reset()
        if policy == "zero":
            action_policy = create_zero_action_policy(env)
        else:
            action_policy = create_random_action_policy(env)
        logger.info(f"{policy.capitalize()} agent is running, press Ctrl+C to exit...")

        budgets = [n for n in (args_cli.max_steps, video_playback_steps(args_cli, env_cfg)) if n is not None]
        max_steps = min(budgets, default=None)

        # keep running while any visualizer is open and the step budget is not exhausted
        sim = env.unwrapped.sim
        step = 0
        while sim.is_running():
            if max_steps is not None and step >= max_steps:
                break
            step += 1
            with torch.inference_mode():
                env.step(action_policy())


def create_zero_action_policy(env: gym.Env) -> Callable[[], Any]:
    """Create a policy that emits finite actions for passive environment playback.

    Manager-based environments infer hold commands for absolute task-space action terms and use literal zeros for all
    other terms. Direct-workflow environments use zero-filled samples of their declared Gymnasium spaces, including
    composite and multi-agent spaces.
    """
    unwrapped = env.unwrapped
    action_manager = getattr(unwrapped, "action_manager", None)
    if action_manager is not None:
        return _create_manager_zero_action_policy(action_manager, unwrapped)

    if isinstance(unwrapped.cfg, DirectMARLEnvCfg):
        actions = {
            agent: sample_space(space, unwrapped.device, batch_size=unwrapped.num_envs, fill_value=0)
            for agent, space in unwrapped.action_spaces.items()
        }
        return lambda: actions

    actions = sample_space(unwrapped.single_action_space, unwrapped.device, batch_size=unwrapped.num_envs, fill_value=0)
    return lambda: actions


def create_random_action_policy(env: gym.Env) -> Callable[[], torch.Tensor | dict[str, torch.Tensor]]:
    """Create a policy that samples uniform random actions in ``[-1, 1]`` for each environment and agent."""
    unwrapped = env.unwrapped
    device = unwrapped.device
    if isinstance(unwrapped.cfg, DirectMARLEnvCfg):
        return lambda: {
            agent: 2 * torch.rand((unwrapped.num_envs, *space.shape), device=device) - 1
            for agent, space in unwrapped.action_spaces.items()
        }

    return lambda: 2 * torch.rand(env.action_space.shape, device=device) - 1


def _create_manager_zero_action_policy(action_manager: Any, env: Any) -> Callable[[], torch.Tensor]:
    """Create a zero-action policy from the active action terms."""
    actions = torch.zeros_like(action_manager.action)
    term_policies = []
    index = 0
    for term_name in action_manager.active_terms:
        term = action_manager.get_term(term_name)
        term_policy = _create_action_term_zero_policy(term, env)
        if term_policy is not None:
            term_policies.append((slice(index, index + term.action_dim), term_policy))
        index += term.action_dim

    def policy() -> torch.Tensor:
        actions.zero_()
        for action_slice, term_policy in term_policies:
            actions[:, action_slice] = term_policy()
        if not torch.isfinite(actions).all():
            raise RuntimeError("Zero agent inferred non-finite actions from the current environment state.")
        return actions

    return policy


def _create_action_term_zero_policy(term: Any, env: Any) -> Callable[[], torch.Tensor] | None:
    """Create the specialized zero-action policy required by an action term."""
    term_types = {cls.__name__ for cls in type(term).__mro__}

    if "PinkInverseKinematicsAction" in term_types:
        controlled_frame_ids, controlled_frame_names = term._asset.find_bodies(
            list(term.cfg.target_eef_link_names.values()), preserve_order=True
        )
        if len(controlled_frame_ids) != len(term.cfg.target_eef_link_names):
            raise ValueError(
                "Expected one controlled body for every Pink IK target. Resolved "
                f"{controlled_frame_names} from {list(term.cfg.target_eef_link_names.values())}."
            )
        if len(controlled_frame_ids) != term._num_frame_tasks:
            raise ValueError(
                f"Pink IK has {term._num_frame_tasks} variable frame tasks but "
                f"{len(controlled_frame_ids)} controlled bodies were configured."
            )

        def pink_policy() -> torch.Tensor:
            frame_poses = term._asset.data.body_link_pose_w.torch[:, controlled_frame_ids].clone()
            frame_poses[..., :3] -= env.scene.env_origins.unsqueeze(1)
            hand_joint_positions = term._asset.data.joint_pos.torch[:, term._hand_joint_ids]
            return torch.cat((frame_poses.flatten(start_dim=1), hand_joint_positions), dim=-1)

        return pink_policy

    if "DifferentialInverseKinematicsAction" in term_types and not term.cfg.controller.use_relative_mode:

        def differential_ik_policy() -> torch.Tensor:
            ee_pos, ee_quat = term._compute_frame_pose()
            command = ee_pos if term.cfg.controller.command_type == "position" else torch.cat((ee_pos, ee_quat), dim=-1)
            return _unscale_action(command - term._offset, term._scale)

        return differential_ik_policy

    if "RMPFlowAction" in term_types and not term.cfg.use_relative_mode:

        def rmpflow_policy() -> torch.Tensor:
            ee_pos, ee_quat = term._compute_frame_pose()
            return _unscale_action(torch.cat((ee_pos, ee_quat), dim=-1), term._scale)

        return rmpflow_policy

    if "OperationalSpaceControllerAction" in term_types and term._pose_abs_idx is not None:
        term_actions = torch.zeros_like(term.raw_actions)

        def operational_space_policy() -> torch.Tensor:
            term_actions.zero_()
            term._compute_ee_pose()
            term._compute_task_frame_pose()
            if term._task_frame_pose_b is None:
                ee_pos_task = term._ee_pose_b[:, :3]
                ee_quat_task = term._ee_pose_b[:, 3:7]
            else:
                ee_pos_task, ee_quat_task = math_utils.subtract_frame_transforms(
                    term._task_frame_pose_b[:, :3],
                    term._task_frame_pose_b[:, 3:7],
                    term._ee_pose_b[:, :3],
                    term._ee_pose_b[:, 3:7],
                )
            term_actions[:, term._pose_abs_idx : term._pose_abs_idx + 3] = _unscale_action(
                ee_pos_task, term._position_scale
            )
            term_actions[:, term._pose_abs_idx + 3 : term._pose_abs_idx + 7] = _unscale_action(
                ee_quat_task, term._orientation_scale
            )
            return term_actions

        return operational_space_policy

    return None


def _unscale_action(command: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    """Map a processed command back to policy-action coordinates without division by zero."""
    return torch.where(scale != 0.0, command / scale, torch.zeros_like(command))


def _parse_args(argv: list[str] | None, policy: PolicyName) -> argparse.Namespace:
    """Parse the command line of a checkpoint-free agent and hand the remainder to Hydra."""
    parser = argparse.ArgumentParser(description=_DESCRIPTIONS[policy])
    parser.add_argument(
        "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
    )
    parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
    parser.add_argument("--task", type=str, default=None, help="Name of the task.")
    parser.add_argument(
        "--max_steps", type=int, default=None, help="Number of environment steps to run. Runs unbounded when omitted."
    )
    add_video_args(parser, action=f"the {policy} agent run")
    add_launcher_args(parser)
    # let task configs select the simulation device
    parser.set_defaults(device=None)
    args_cli, hydra_args = setup_preset_cli(parser, argv)
    enable_cameras_for_video(args_cli)
    sys.argv = [sys.argv[0]] + hydra_args
    return args_cli

代码解释#

envs.ManagerBasedRLEnv 类继承自 gymnasium.Env 类以遵循标准接口。然而，与传统的 Gym 环境不同， envs.ManagerBasedRLEnv 实现了*向量化*环境。这意味着多个环境实例同时在同一进程中运行，并且所有数据都以批处理方式返回。

同样， envs.DirectRLEnv 类也从 gymnasium.Env 类继承以用于直接工作流程。对于 envs.DirectMARLEnv ，虽然它不是从 Gymnasium 继承的，但可以以相同的方式注册和创建。

使用 gym 注册表#

要注册一个环境，我们使用 gymnasium.register() 方法。此方法接收环境名称、环境类的入口点以及环境配置类的入口点。

备注

gymnasium 注册表是一个全局注册表。因此，确保环境名称是唯一的很重要。否则，在注册环境时注册表会抛出错误。

基于管理器的环境#

对于基于管理器的环境，以下展示了 isaaclab_tasks.core.cartpole 子包中倒立摆环境的注册调用：

from . import agents

        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerCfg",
        "torchrl_cfg_entry_point": f"{agents.__name__}.torchrl_ppo_cfg:CartpolePPOCfg",
        "default_agent": "rsl_rl",
        "rsl_rl_with_symmetry_cfg_entry_point": (
            f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerWithSymmetryCfg"
        ),
        "skrl_cfg_entry_point": f"{agents.__name__}:skrl_manager_ppo_cfg.yaml",
        "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml",
    },
)

gym.register(
    id="Isaac-Cartpole-Camera",
    entry_point="isaaclab.envs:ManagerBasedRLEnv",
    disable_env_checker=True,

id 参数是环境名称。按照约定，我们为环境名称添加 Isaac- 前缀，以便在注册表中搜索。前缀后通常依次是任务名和机器人名。例如，使用 ANYmal C 在平地上进行足式运动的环境叫作 IsaacContrib-Velocity-Flat-AnymalC 。版本号 v<N> 通常用来区分同一环境的不同变体，否则环境名称可能过长、难以阅读。

entry_point 参数是环境类的入口点。入口点是一个 <module>:<class> 形式的字符串。对于 cartpole 环境，入口点是 isaaclab.envs:ManagerBasedRLEnv 。入口点用于在创建环境实例时导入环境类。

env_cfg_entry_point 参数指定了环境的默认配置。默认配置使用 isaaclab_tasks.utils.parse_env_cfg() 函数加载。然后将其传递给 gymnasium.make() 函数以创建环境实例。配置入口点可以是 YAML 文件或 Python 配置类。

直接式的环境#

对于基于直接的环境，环境注册遵循类似的模式。与将环境的入口点注册为 ManagerBasedRLEnv 类不同，我们将环境的入口点注册为环境的实现类。此外，我们将环境名称添加后缀 -Direct 以区分其与基于管理器的环境。

例如，以下展示了 isaaclab_tasks.core.cartpole 子包中倒立摆环境的注册调用：

from . import agents

        "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml",
    },
)

gym.register(
    id="Isaac-Cartpole-Camera-Direct",
    entry_point=f"{__name__}.cartpole_direct_camera_env:CartpoleCameraEnv",
    disable_env_checker=True,
    kwargs={
        "env_cfg_entry_point": f"{__name__}.cartpole_direct_camera_env_cfg:CartpoleCameraEnvCfg",
        "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml",
        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpoleCameraDirectPPORunnerCfg",

创建环境#

要将所有 isaaclab_tasks 扩展提供的环境通知 gym 注册表，我们必须在脚本开头导入模块。这将执行 __init__.py 文件，该文件遍历所有子包并注册它们各自的环境。

import isaaclab_tasks  # noqa: F401

在本教程中，任务名称从命令行中读取。任务名称用于解析默认配置以及创建环境实例。此外，其他解析的命令行参数，如环境数量、仿真设备以及是否进行渲染，用于覆盖默认配置。

    env_cfg, _ = resolve_task_config(args_cli.task, "")
    apply_env_overrides(args_cli, env_cfg)
    pre_launch_video_config(env_cfg, args_cli)
    # reject unsupported configurations before launching Kit or initializing a native physics backend
    try:
        validate(env_cfg)
    except (TypeError, ValueError) as exc:
        raise SystemExit(f"Invalid environment configuration: {exc}") from None

    with launch_simulation(env_cfg, args_cli), contextlib.ExitStack() as cleanup:
        log_dir = os.path.abspath(os.path.join("logs", f"{policy}_agent", normalize_task_name(args_cli.task)))
        apply_video_recording(env_cfg, log_dir, args_cli, subdir="play")
        env = gym.make(args_cli.task, cfg=env_cfg)

一旦创建环境，其余执行将遇到标准的重置和步进。

代码执行#

现在我们已经浏览了代码，让我们运行脚本并查看结果:

uv（推荐）

uv run python scripts/environments/random_agent.py --task Isaac-Cartpole --num_envs 32 --viz kit

这将打开一个与创建基于管理器的强化学习环境教程类似的场景。要停止仿真，您可以关闭窗口，或在终端中按 Ctrl+C 。

此外，您还可以通过显式设置 --device 的值将仿真设备从 GPU 更改为 CPU:

uv（推荐）

uv run python scripts/environments/random_agent.py --task Isaac-Cartpole --num_envs 32 --device cpu --viz kit

使用 --device cpu 标志，仿真将在 CPU 上运行。这对于调试仿真很有用。但是，与在 GPU 上运行相比，仿真速度会慢得多。