From 2e0d86edbe38e831241a8c9672dbe84095dc9192 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 5 May 2026 09:37:34 +0000 Subject: [PATCH 01/58] Sync warp env's viewport check with stable env's API PR #5297 (Decouple Renderer from Camera) replaced `sim.get_setting('/isaaclab/visualizer')` (returned a comma-separated string) with `sim.has_active_visualizers()` plus per-type queries; the warp env path was missed, so it crashes with "AttributeError: 'dict' object has no attribute 'split'" during env construction. Mirror the stable env's pattern: - `has_active_visualizers()` for the gating predicate - `has_kit()` so kitless Newton-only runs (`--viz rerun`) skip ViewportCameraController --- .../envs/manager_based_env_warp.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py index 8c558b925947..03f8cc9dfba4 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py @@ -33,6 +33,7 @@ from isaaclab.ui.widgets import ManagerLiveVisualizer from isaaclab.utils.seed import configure_seed from isaaclab.utils.timer import Timer +from isaaclab.utils.version import has_kit from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp as InteractiveScene from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode, ManagerCallSwitch @@ -154,13 +155,11 @@ def __init__(self, cfg: ManagerBasedEnvCfg): # Persistent scalar buffer for global env step count (stable pointer for capture). self._global_env_step_count_wp = wp.zeros((1,), dtype=wp.int32, device=self.device) - # set up camera viewport controller - # viewport is not available in other rendering modes so the function will throw a warning - # FIXME: This needs to be fixed in the future when we unify the UI functionalities even for - # non-rendering modes. - viz_str = self.sim.get_setting("/isaaclab/visualizer") or "" - available_visualizers = [v.strip() for v in viz_str.split(",") if v.strip()] - if "kit" in available_visualizers and bool(viz_str): + # set up camera viewport controller — mirror stable env (PR #5297 changed + # the visualizer API; this branch was missed). ViewportCameraController + # uses omni.kit; skip it in kitless Newton-only runs. + has_visualizers = self.sim.has_active_visualizers() + if (self.sim.has_gui or has_visualizers) and has_kit(): self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) else: self.viewport_camera_controller = None From 5fb7327351fba10816bbf23c901a4b720c10dc65 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 5 May 2026 09:37:37 +0000 Subject: [PATCH 02/58] Alias warp manager term cfgs to stable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Warp's ObservationTermCfg / RewardTermCfg / TerminationTermCfg had no behavioural difference from their stable counterparts — same fields, same base class. The override only carried Warp-first docstrings while keeping the classes as siblings of stable's, which broke `isinstance(stable_term, warp.TermCfg)` checks inside the experimental managers when a stable cfg is fed through the warp runtime. Re-export stable's term cfg classes directly. Term funcs still follow the warp-first `func(env, out, **params)` signature at runtime; the type annotation just lives on stable's class instead of warp's. --- .../managers/manager_term_cfg.py | 95 +++---------------- 1 file changed, 13 insertions(+), 82 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py index 3fab3bfc5ffc..604b138003d4 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py @@ -5,90 +5,21 @@ """Configuration terms for different managers (experimental, Warp-first). -This module is a passthrough to :mod:`isaaclab.managers.manager_term_cfg` except for -the following term configs which are overridden for Warp-first execution: - -- :class:`ObservationTermCfg` -- :class:`RewardTermCfg` -- :class:`TerminationTermCfg` +The warp manager classes accept the same term cfg shapes as the stable +managers, so this module simply re-exports the stable term cfg classes +to keep ``isinstance(stable_term, warp.TermCfg)`` true. This is what +lets the :class:`isaaclab_experimental.envs.warp_frontend.WarpFrontend` +adapter run a stable cfg through the warp runtime without rewrapping +every term. + +The ``func`` callable on each term is still expected to follow the +warp-first ``func(env, out, **params) -> None`` signature when run on +the warp runtime; only the *type* is shared with stable. """ from __future__ import annotations -from collections.abc import Callable -from dataclasses import MISSING - +# Re-export stable manager term cfg classes verbatim. +# `from … import *` carries `ObservationTermCfg`, `RewardTermCfg`, +# `TerminationTermCfg`, `ManagerTermBaseCfg`, etc. from isaaclab.managers.manager_term_cfg import * # noqa: F401,F403 -from isaaclab.managers.manager_term_cfg import ManagerTermBaseCfg as _ManagerTermBaseCfg -from isaaclab.utils import configclass - - -@configclass -class RewardTermCfg(_ManagerTermBaseCfg): - """Configuration for a reward term. - - The function is expected to write the (unweighted) reward values into a - pre-allocated Warp buffer provided by the manager. - - Expected signature: - - - ``func(env, out, **params) -> None`` - - where ``out`` is a Warp array of shape ``(num_envs,)`` with float32 dtype. - """ - - func: Callable[..., None] = MISSING - """The function to be called to fill the pre-allocated reward buffer.""" - - weight: float = MISSING - """The weight of the reward term.""" - - -@configclass -class TerminationTermCfg(_ManagerTermBaseCfg): - """Configuration for a termination term (experimental, Warp-first). - - The function is expected to write termination flags into a pre-allocated Warp buffer. - - Expected signature: - - - ``func(env, out, **params) -> None`` - - where ``out`` is a Warp array of shape ``(num_envs,)`` with boolean dtype. - """ - - func: Callable[..., None] = MISSING - """The function to be called to fill the pre-allocated termination buffer.""" - - time_out: bool = False - """Whether the termination term contributes towards episodic timeouts. Defaults to False.""" - - -@configclass -class ObservationTermCfg(_ManagerTermBaseCfg): - """Configuration for an observation term (experimental, Warp-first). - - The function is expected to write observation values into a pre-allocated Warp buffer provided - by the observation manager. - - Expected signature: - - - ``func(env, out, **params) -> None`` - - where ``out`` is a Warp array of shape ``(num_envs, obs_term_dim)`` with float32 dtype. - - Notes: - - The stable fields (noise/modifiers/history) are kept for config compatibility, but the - experimental Warp-first observation manager may not support all of them initially. - """ - - func: Callable[..., None] = MISSING - """The function to be called to fill the pre-allocated observation buffer.""" - - # Keep stable configuration fields for compatibility with existing task configs. - modifiers: list[ModifierCfg] | None = None # noqa: F405 - noise: NoiseCfg | NoiseModelCfg | None = None # noqa: F405 - clip: tuple[float, float] | None = None - scale: tuple[float, ...] | float | None = None - history_length: int = 0 - flatten_history_dim: bool = True From e842ea073fd50ee8467361945bccafea67db1acc Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 5 May 2026 09:37:55 +0000 Subject: [PATCH 03/58] Add WarpFrontend bridge: '--manager=warp' on any stable task MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Lets a stable manager-based RL env cfg run on the experimental warp runtime without a separate warp task registration. Replaces the duplicate per-robot warp env cfgs by adapting the stable cfg in place. Components: - isaaclab_experimental/envs/warp_frontend.py: 'WarpFrontend' walks the stable cfg, swaps each `term.func` to its same-named warp twin (only accepting candidates whose `__module__` lives under the warp packages — the warp mdp module re-exports stable terms via `from … import *`, so a naive `getattr` would silently keep the stable function), swaps each action `class_type`, picks the `newton` field of any PresetCfg, drops sensors with no warp counterpart (`height_scanner`), and in-place class-promotes `SceneEntityCfg` instances so warp kernels see the `joint_mask` / *\_ids_wp` cached fields. A `WarpAdaptReport` records every term that had no warp twin and surfaces them via the logger; `strict=True` makes the same condition raise `LookupError` instead. - scripts/reinforcement_learning/rsl_rl/train.py: '--manager=warp' flag routes the constructed env through `WarpFrontend.build` instead of `gym.make`. The flag also auto-injects `presets=newton` into Hydra's argv so that PresetCfg wrappers resolve to the newton preset (Hydra's preset resolution runs *before* the adapter). Validated: cartpole and Anymal-D Flat both pass a 3-way comparison — - 'Isaac-Cartpole-v0' (stable manager): trains. - 'Isaac-Cartpole-Warp-v0' (existing direct path): reward 0.06 / ep 76. - 'Isaac-Cartpole-v0 --manager=warp': reward 0.06 / ep 76 (matches direct exactly). - 'Isaac-Velocity-Flat-Anymal-D-v0 presets=newton': -8.45 / 191. - 'Isaac-Velocity-Flat-Anymal-D-Warp-v0': -7.47 / 168. - 'Isaac-Velocity-Flat-Anymal-D-v0 --manager=warp': -7.69 / 174 (within run-to-run variance of the direct warp path). Both stable and warp frontends remain functional for every task; this PR adds a flag-based selector without removing the existing direct warp registrations. --- .../reinforcement_learning/rsl_rl/train.py | 28 +- .../envs/warp_frontend.py | 340 ++++++++++++++++++ 2 files changed, 367 insertions(+), 1 deletion(-) create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py diff --git a/scripts/reinforcement_learning/rsl_rl/train.py b/scripts/reinforcement_learning/rsl_rl/train.py index eefc13a8aa2c..e66485d3a37d 100644 --- a/scripts/reinforcement_learning/rsl_rl/train.py +++ b/scripts/reinforcement_learning/rsl_rl/train.py @@ -67,6 +67,17 @@ "--ray-proc-id", "-rid", type=int, default=None, help="Automatically configured by Ray integration, otherwise None." ) parser.add_argument("--external_callback", default=None, help="Fully qualified path to an externally defined callback.") +parser.add_argument( + "--manager", + type=str, + default="stable", + choices=["stable", "warp"], + help=( + "Manager-based env runtime. 'stable' uses isaaclab.envs.ManagerBasedRLEnv (torch);" + " 'warp' adapts the same task cfg via isaaclab_experimental.envs.warp_frontend.WarpFrontend" + " and runs on isaaclab_experimental.envs.ManagerBasedRLEnvWarp." + ), +) cli_args.add_rsl_rl_args(parser) add_launcher_args(parser) args_cli, remaining_args = parser.parse_known_args() @@ -86,6 +97,14 @@ # The remaining arguments are the arguments that were not consumed by both this scripts # argparser and (optionally) the external callback function. remaining_args = list_intersection(remaining_args, remaining_args_env_registration) + +# When the warp manager runtime is selected, the env cfg must resolve any +# PresetCfg wrappers to their `newton` field (Hydra preset resolution runs +# *before* the WarpFrontend adapter, so we inject the override here unless +# the user already passed one explicitly). +if args_cli.manager == "warp" and not any(a.startswith("presets=") for a in remaining_args): + remaining_args.append("presets=newton") + sys.argv = [sys.argv[0]] + remaining_args # -- check RSL-RL version ---------------------------------------------------- @@ -171,7 +190,14 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen env_cfg.log_dir = log_dir # create isaac environment - env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + if args_cli.manager == "warp": + # Lazy: this is the first warp-side import. Calling it after + # SimulationApp is already alive avoids racing pxr extension init. + from isaaclab_experimental.envs.warp_frontend import WarpFrontend + + env = WarpFrontend().build(env_cfg, task_id=args_cli.task) + else: + env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) # convert to single-agent instance if required by the RL algorithm if isinstance(env.unwrapped.cfg, DirectMARLEnvCfg): diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py new file mode 100644 index 000000000000..36170e991340 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py @@ -0,0 +1,340 @@ +# 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 + +"""Adapt a stable manager-based RL env cfg to run on the warp runtime. + +Module-level imports are deliberately limited to lightweight stdlib modules. +Warp-library imports fire only inside :meth:`WarpFrontend.build`, which must +be called *after* :class:`SimulationApp` is alive — otherwise the warp lib +loading will race with Kit's USD/pxr extension initialisation. +""" + +from __future__ import annotations + +import importlib +import logging +from collections.abc import Iterable +from dataclasses import dataclass, field +from types import ModuleType +from typing import Any + +logger = logging.getLogger(__name__) + + +@dataclass +class MissingItem: + """A symbol the adapter looked for but couldn't resolve to a warp twin.""" + + group: str + """Where it was looked up — e.g. ``rewards``, ``observations.policy``, ``actions``.""" + term_name: str + """Field name on the group, e.g. ``track_lin_vel_xy_exp``.""" + expected_name: str + """The symbol name searched for in the warp modules (the stable ``__name__``).""" + kind: str + """Either ``"func"`` (mdp function) or ``"class_type"`` (action runtime class).""" + searched: tuple[str, ...] + """Module dotted-paths searched, in order.""" + action: str + """Outcome: ``"dropped"`` (term set to None) or ``"raised"`` (caller saw an error).""" + + +@dataclass +class WarpAdaptReport: + """Result of :meth:`WarpFrontend.adapt`. Empty ``missing`` means clean run.""" + + missing: list[MissingItem] = field(default_factory=list) + + def __bool__(self) -> bool: # truthy when there is something to report + return bool(self.missing) + + def format(self) -> str: + if not self.missing: + return "WarpFrontend: no missing items." + lines = [f"WarpFrontend: {len(self.missing)} missing warp twin(s):"] + for m in self.missing: + lines.append( + f" - {m.group}.{m.term_name}: stable {m.kind}={m.expected_name!r} → " + f"no twin in {list(m.searched)} ({m.action})" + ) + return "\n".join(lines) + + +class _NameSwap: + """Replace ``term.`` with the same-named symbol in ``modules``. + + Drop the term entirely if no twin exists and ``drop_if_missing`` is set; + otherwise raise :class:`LookupError`. Either way, append a record to + ``report.missing`` so the caller can summarise what was lost. + + Used for both ``term.func`` (mdp swap) and ``action.class_type`` swaps — + the algorithm is identical, only the attribute name and policy differ. + """ + + def __init__( + self, + modules: Iterable[ModuleType], + attr: str, + drop_if_missing: bool, + report: WarpAdaptReport, + group_label: str, + warp_module_prefixes: tuple[str, ...] = ("isaaclab_experimental", "isaaclab_tasks_experimental"), + ): + self._modules = tuple(modules) + self._attr = attr + self._drop = drop_if_missing + self._report = report + self._group = group_label + # The warp mdp __init__ re-exports stable mdp terms via ``from isaaclab.envs.mdp import *`` + # before adding its own overrides. ``getattr(warp_mdp, name)`` therefore can return the + # *stable* implementation untouched. We accept a candidate as a real warp twin only if + # its ``__module__`` lives under one of these prefixes. + self._warp_module_prefixes = warp_module_prefixes + + def apply_to(self, group: Any) -> None: + if group is None: + return + for name in (n for n in dir(group) if not n.startswith("_")): + term = getattr(group, name, None) + if term is None or not hasattr(term, self._attr): + continue + stable = getattr(term, self._attr) + if not callable(stable): + continue + twin = self._resolve_twin(stable.__name__) + if twin is not None: + setattr(term, self._attr, twin) + continue + + # No twin found. + missing = MissingItem( + group=self._group, + term_name=name, + expected_name=stable.__name__, + kind=self._attr, + searched=tuple(m.__name__ for m in self._modules), + action="dropped" if self._drop else "raised", + ) + self._report.missing.append(missing) + if self._drop: + setattr(group, name, None) + else: + raise LookupError( + f"WarpFrontend: cannot adapt {self._group}.{name} — no warp twin for" + f" stable {self._attr}={stable.__name__!r}. Searched {missing.searched}." + ) + + def _resolve_twin(self, name: str) -> Any | None: + """Return the warp implementation of ``name`` from the search modules. + + Skips candidates that are merely re-exports of stable code: the warp mdp + package does ``from isaaclab.envs.mdp import *`` and a ``getattr`` lookup + will happily return the stable symbol untouched. We confirm a candidate + is actually a warp implementation by checking its ``__module__``. + """ + for module in self._modules: + candidate = getattr(module, name, None) + if candidate is None: + continue + origin = getattr(candidate, "__module__", "") or "" + if origin.startswith(self._warp_module_prefixes): + return candidate + return None + + +class WarpFrontend: + """Adapt a stable env cfg in place; build a warp-runtime env. + + Call :meth:`build` only after ``SimulationApp`` is alive — warp lib loads + lazily inside ``build``. All knobs are constructor args; subclass only if + you need to change the algorithm itself, not its parameters. + + Example:: + + with launch_simulation(env_cfg, args_cli): + env = WarpFrontend().build(env_cfg, task_id=args_cli.task) + """ + + DEFAULT_TERM_GROUPS: tuple[tuple[str, ...], ...] = ( + ("observations", "policy"), + ("events",), + ("rewards",), + ("terminations",), + ("curriculum",), + ("actions",), # included for SceneEntityCfg upgrade in action params + ) + + def __init__( + self, + stable_root: str = "isaaclab_tasks", + warp_root: str = "isaaclab_tasks_experimental", + fallback_mdp: str = "isaaclab_experimental.envs.mdp", + warp_actions: str = "isaaclab_experimental.envs.mdp.actions.joint_actions", + drop_sensors: Iterable[str] = ("height_scanner",), + term_groups: Iterable[tuple[str, ...]] = DEFAULT_TERM_GROUPS, + strict: bool = False, + ): + self._stable_root = stable_root + self._warp_root = warp_root + self._fallback_mdp = fallback_mdp + self._warp_actions = warp_actions + self._drop_sensors = tuple(drop_sensors) + self._term_groups = tuple(term_groups) + self._strict = strict + + # -- public --------------------------------------------------------------- + + def adapt(self, cfg: Any, task_id: str) -> WarpAdaptReport: + """Mutate ``cfg`` in place; return a report of what couldn't be adapted. + + With ``strict=True``, any missing twin (term func or action class) raises + :class:`LookupError` immediately so the caller sees the failure in the + traceback instead of a silently-dropped term. + """ + report = WarpAdaptReport() + + # 1. PresetCfg → newton (also dodges the lazy PhysxCfg materialisation). + # Note: by the time we run, Hydra's `resolve_presets` has already + # collapsed any PresetCfg wrapper to a concrete preset. The auto + # ``presets=newton`` injected from train.py for ``--manager=warp`` + # ensures that concrete preset is the newton one. If the preset + # somehow still has a ``newton`` attribute (custom subclass), pick it. + physics = getattr(cfg.sim, "physics", None) + if physics is not None and hasattr(physics, "newton"): + cfg.sim.physics = physics.newton + + # 2. Drop sensors warp can't run yet. + scene = getattr(cfg, "scene", None) + for sensor in self._drop_sensors: + if scene is not None and getattr(scene, sensor, None) is not None: + setattr(scene, sensor, None) + + # 3. Upgrade SceneEntityCfg instances inside term.params dicts to the + # warp-extended variant (adds joint_mask / *_ids_wp wp.array fields + # that warp mdp kernels read). + self._upgrade_scene_entity_cfgs(cfg) + + # 4. Swap term.func across all groups. + warp_mdp_modules = self._mdp_modules_for(task_id) + # Always log which warp mdp modules were resolved so it's easy to debug + # missing-twin reports. + logger.info( + "WarpFrontend: warp mdp modules for %r → %s", + task_id, + [m.__name__ for m in warp_mdp_modules], + ) + for path in self._term_groups: + if path == ("actions",): + continue # actions don't have ``func`` — handled below + group = self._walk(cfg, path) + label = ".".join(path) + swap = _NameSwap(warp_mdp_modules, "func", drop_if_missing=not self._strict, report=report, group_label=label) + swap.apply_to(group) + + # 5. Swap action class_type. An action with no warp class_type can't + # run on the warp manager, so this is always strict (raise). + action_swap = _NameSwap( + [importlib.import_module(self._warp_actions)], + attr="class_type", + drop_if_missing=False, + report=report, + group_label="actions", + ) + action_swap.apply_to(getattr(cfg, "actions", None)) + + # 6. Surface the report. + if report.missing: + level = logging.ERROR if self._strict else logging.WARNING + logger.log(level, report.format()) + + return report + + def build(self, cfg: Any, task_id: str): + """Adapt ``cfg`` and return a :class:`ManagerBasedRLEnvWarp` instance. + + Always logs the missing-twin report (if any). With ``strict=True``, + any missing twin raises before the env is constructed. + """ + # Lazy: this is the first warp-lib load. Caller must already be inside + # the SimulationApp context, i.e. inside ``launch_simulation``. + from isaaclab_experimental.envs import ManagerBasedRLEnvWarp + + self.adapt(cfg, task_id) + return ManagerBasedRLEnvWarp(cfg=cfg) + + # -- helpers -------------------------------------------------------------- + + def _upgrade_scene_entity_cfgs(self, cfg: Any) -> None: + """Upgrade every stable :class:`SceneEntityCfg` to the warp variant. + + Walks ``term.params`` dicts inside each term group and re-classes any + stable :class:`SceneEntityCfg` instance into the warp subclass so warp + mdp kernels can read ``joint_mask`` / ``joint_ids_wp`` / ``body_ids_wp`` + after :meth:`resolve`. + """ + from isaaclab.managers.scene_entity_cfg import SceneEntityCfg as _StableSE + from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as _WarpSE + + def upgrade(obj: Any) -> None: + if isinstance(obj, _WarpSE) or not isinstance(obj, _StableSE): + return + # In-place class promotion + initialise warp-only fields. The class + # hierarchy guarantees compatible memory layout (warp inherits stable). + obj.__class__ = _WarpSE + obj.joint_mask = None + obj.joint_ids_wp = None + obj.body_ids_wp = None + + for path in self._term_groups: + group = self._walk(cfg, path) + if group is None: + continue + for name in (n for n in dir(group) if not n.startswith("_")): + term = getattr(group, name, None) + if term is None: + continue + params = getattr(term, "params", None) + if isinstance(params, dict): + for v in params.values(): + upgrade(v) + + @staticmethod + def _walk(root: Any, path: tuple[str, ...]) -> Any: + node = root + for attr in path: + node = getattr(node, attr, None) + if node is None: + return None + return node + + def _mdp_modules_for(self, task_id: str) -> list[ModuleType]: + """Locate warp mdp modules that mirror the stable task's mdp. + + Convention: replace ``isaaclab_tasks`` with ``isaaclab_tasks_experimental`` + in the stable cfg's module path; walk up looking for an ``mdp`` + sub-module on the warp side; always append the package-wide warp mdp + as a fallback for cross-task base terms. + """ + modules: list[ModuleType] = [] + try: + import gymnasium as gym + + entry = gym.spec(task_id).kwargs.get("env_cfg_entry_point") + except Exception: + entry = None + if isinstance(entry, str) and entry.startswith(self._stable_root): + warp_pkg = entry.rsplit(".", 1)[0].replace(self._stable_root, self._warp_root, 1) + parts = warp_pkg.split(".") + for depth in range(len(parts), 0, -1): + try: + modules.append(importlib.import_module(".".join(parts[:depth] + ["mdp"]))) + break + except ImportError: + continue + try: + modules.append(importlib.import_module(self._fallback_mdp)) + except ImportError: + logger.warning("WarpFrontend: fallback mdp module %r not importable", self._fallback_mdp) + return modules From 1cab7027c398664f5c750b18970f02ab2fd1052a Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 5 May 2026 16:41:24 +0000 Subject: [PATCH 04/58] Add changelog and apply ruff formatting for warp bridge Adds the isaaclab_experimental changelog fragment for the WarpFrontend adapter and the --manager flag, and applies the ruff-format pass that the pre-commit hook produced on warp_frontend.py. --- .../changelog.d/warp-manager-bridge.rst | 26 +++++++++++++++++++ .../envs/warp_frontend.py | 5 +++- 2 files changed, 30 insertions(+), 1 deletion(-) create mode 100644 source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst diff --git a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst new file mode 100644 index 000000000000..758f4223a5f5 --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst @@ -0,0 +1,26 @@ +Added +^^^^^ + +* Added :class:`~isaaclab_experimental.envs.warp_frontend.WarpFrontend`, a + runtime adapter that lets any stable manager-based RL task config run on the + experimental warp manager runtime (:class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp`) + without a parallel ``-Warp-v0`` registration. The adapter resolves + :class:`~isaaclab_physx.preset.PresetCfg` to its ``newton`` field, swaps + ``term.func`` references to same-named warp twins discovered in the warp + ``mdp`` modules (skipping stable re-exports), promotes ``SceneEntityCfg`` + instances in-place to the warp variant, and reports any missing twins + before the env is built. + +* Added a ``--manager={stable,warp}`` flag to ``rsl_rl/train.py``. When set + to ``warp`` the script auto-injects ``presets=newton`` (so Hydra picks the + Newton physics preset before the adapter runs) and dispatches the env + through ``WarpFrontend`` instead of ``gym.make``. + +Fixed +^^^^^ + +* Fixed a regression in :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` + introduced when the ``SimulationContext.get_setting`` API was reshaped: + the warp env now mirrors the stable env and probes + :meth:`~isaaclab.sim.SimulationContext.has_active_visualizers` instead of + splitting a string setting that no longer exists. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py index 36170e991340..9da4f9c80e0e 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py @@ -230,7 +230,9 @@ def adapt(self, cfg: Any, task_id: str) -> WarpAdaptReport: continue # actions don't have ``func`` — handled below group = self._walk(cfg, path) label = ".".join(path) - swap = _NameSwap(warp_mdp_modules, "func", drop_if_missing=not self._strict, report=report, group_label=label) + swap = _NameSwap( + warp_mdp_modules, "func", drop_if_missing=not self._strict, report=report, group_label=label + ) swap.apply_to(group) # 5. Swap action class_type. An action with no warp class_type can't @@ -275,6 +277,7 @@ def _upgrade_scene_entity_cfgs(self, cfg: Any) -> None: after :meth:`resolve`. """ from isaaclab.managers.scene_entity_cfg import SceneEntityCfg as _StableSE + from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as _WarpSE def upgrade(obj: Any) -> None: From 4876327a782a012a2fc1a7755986ce93b0de69fd Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 5 May 2026 17:24:06 +0000 Subject: [PATCH 05/58] Refactor warp bridge into pluggable rule pipeline; rename to --frontend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The adapter is now a sequence of CompatRule objects (resolve preset, drop sensors, promote SceneEntityCfg, swap mdp funcs, swap action class). New incompatibilities are added by writing a small rule subclass instead of editing the dispatcher. The CLI flag is renamed --manager → --frontend because the dispatch also covers direct envs: a stable manager-based cfg is adapted onto ManagerBasedRLEnvWarp; a direct task is verified to point at a warp env class and dispatched via gym.make. A stable direct cfg + --frontend=warp raises IncompatibleEnvError with the offending entry_point and a hint at the *-Direct-Warp-v0 alternative. Other fixes: - Forward render_mode through build() so --video keeps working. - Attach the CompatReport on env.unwrapped.warp_compat_report so callers can inspect what was dropped or left unresolved. - Assert the warp SceneEntityCfg subclasses the stable one before doing the in-place __class__ promotion; the rule fails loudly if the hierarchy is ever broken. - Narrow the bare except in mdp-module discovery so real ImportErrors from broken cfgs propagate. - presets=newton is now only auto-injected for stable manager-based tasks; direct warp tasks (which don't carry presets) are left alone. - Warn when the user passes presets= with --frontend=warp. - Add the commands group to the rule that promotes SceneEntityCfg. --- .../reinforcement_learning/rsl_rl/train.py | 49 +- .../changelog.d/warp-manager-bridge.rst | 26 +- .../envs/warp_frontend.py | 720 ++++++++++++------ 3 files changed, 517 insertions(+), 278 deletions(-) diff --git a/scripts/reinforcement_learning/rsl_rl/train.py b/scripts/reinforcement_learning/rsl_rl/train.py index e66485d3a37d..15326eb4b196 100644 --- a/scripts/reinforcement_learning/rsl_rl/train.py +++ b/scripts/reinforcement_learning/rsl_rl/train.py @@ -68,14 +68,15 @@ ) parser.add_argument("--external_callback", default=None, help="Fully qualified path to an externally defined callback.") parser.add_argument( - "--manager", + "--frontend", type=str, default="stable", choices=["stable", "warp"], help=( - "Manager-based env runtime. 'stable' uses isaaclab.envs.ManagerBasedRLEnv (torch);" - " 'warp' adapts the same task cfg via isaaclab_experimental.envs.warp_frontend.WarpFrontend" - " and runs on isaaclab_experimental.envs.ManagerBasedRLEnvWarp." + "Runtime backend for the env. 'stable' uses isaaclab.envs.* (torch);" + " 'warp' routes through isaaclab_experimental.envs.warp_frontend.WarpFrontend," + " which adapts a manager-based stable cfg onto ManagerBasedRLEnvWarp or dispatches" + " a direct task to its registered warp env class." ), ) cli_args.add_rsl_rl_args(parser) @@ -98,12 +99,29 @@ # argparser and (optionally) the external callback function. remaining_args = list_intersection(remaining_args, remaining_args_env_registration) -# When the warp manager runtime is selected, the env cfg must resolve any -# PresetCfg wrappers to their `newton` field (Hydra preset resolution runs -# *before* the WarpFrontend adapter, so we inject the override here unless -# the user already passed one explicitly). -if args_cli.manager == "warp" and not any(a.startswith("presets=") for a in remaining_args): - remaining_args.append("presets=newton") +# When the warp frontend is selected on a stable manager-based task, the cfg +# must resolve any PresetCfg wrappers to their ``newton`` field. Hydra +# resolves presets *before* the WarpFrontend runs, so we inject +# ``presets=newton`` here. We only inject for tasks registered under +# ``isaaclab_tasks.manager_based`` — direct tasks and pre-warp registrations +# don't carry a preset system, and injecting ``presets=newton`` against them +# causes Hydra to error before the frontend can produce its own diagnostic. +if args_cli.frontend == "warp" and args_cli.task is not None: + try: + _spec = gym.spec(args_cli.task) + except gym.error.NameNotFound: + _spec = None + _entry = _spec.entry_point if _spec is not None else None + _is_stable_manager = isinstance(_entry, str) and _entry.startswith("isaaclab_tasks.manager_based") + _explicit_preset = next((a for a in remaining_args if a.startswith("presets=")), None) + if _is_stable_manager and _explicit_preset is None: + remaining_args.append("presets=newton") + elif _is_stable_manager and _explicit_preset != "presets=newton": + logger.warning( + "--frontend=warp on %r expects presets=newton; got %r — adapter may fail to find a Newton physics cfg.", + args_cli.task, + _explicit_preset, + ) sys.argv = [sys.argv[0]] + remaining_args @@ -190,14 +208,15 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen env_cfg.log_dir = log_dir # create isaac environment - if args_cli.manager == "warp": - # Lazy: this is the first warp-side import. Calling it after - # SimulationApp is already alive avoids racing pxr extension init. + render_mode = "rgb_array" if args_cli.video else None + if args_cli.frontend == "warp": + # Lazy: first warp-side import. Calling this after SimulationApp + # is already alive avoids racing pxr extension init. from isaaclab_experimental.envs.warp_frontend import WarpFrontend - env = WarpFrontend().build(env_cfg, task_id=args_cli.task) + env = WarpFrontend().build(env_cfg, task_id=args_cli.task, render_mode=render_mode) else: - env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + env = gym.make(args_cli.task, cfg=env_cfg, render_mode=render_mode) # convert to single-agent instance if required by the RL algorithm if isinstance(env.unwrapped.cfg, DirectMARLEnvCfg): diff --git a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst index 758f4223a5f5..881450b1312d 100644 --- a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst +++ b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst @@ -3,18 +3,24 @@ Added * Added :class:`~isaaclab_experimental.envs.warp_frontend.WarpFrontend`, a runtime adapter that lets any stable manager-based RL task config run on the - experimental warp manager runtime (:class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp`) - without a parallel ``-Warp-v0`` registration. The adapter resolves - :class:`~isaaclab_physx.preset.PresetCfg` to its ``newton`` field, swaps - ``term.func`` references to same-named warp twins discovered in the warp - ``mdp`` modules (skipping stable re-exports), promotes ``SceneEntityCfg`` - instances in-place to the warp variant, and reports any missing twins - before the env is built. + experimental warp runtime (:class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp`) + without a parallel ``-Warp-v0`` registration. The adapter is built on a + pluggable :class:`~isaaclab_experimental.envs.warp_frontend.CompatRule` + pipeline; new incompatibilities (sensor types, term-cfg fields, action + classes) are added by writing a small rule subclass instead of editing the + dispatcher. The default rules cover physics-preset resolution, dropping + unsupported sensors, in-place :class:`SceneEntityCfg` promotion, mdp + function swaps, and action-class swaps. The frontend also dispatches + direct envs by verifying their registered entry-point class lives under + ``isaaclab_experimental`` / ``isaaclab_tasks_experimental`` and routing + through :func:`gym.make` unchanged. -* Added a ``--manager={stable,warp}`` flag to ``rsl_rl/train.py``. When set +* Added a ``--frontend={stable,warp}`` flag to ``rsl_rl/train.py``. When set to ``warp`` the script auto-injects ``presets=newton`` (so Hydra picks the - Newton physics preset before the adapter runs) and dispatches the env - through ``WarpFrontend`` instead of ``gym.make``. + Newton physics preset before the adapter runs), warns on conflicting + ``presets=`` overrides, and dispatches the env through ``WarpFrontend`` + instead of :func:`gym.make`. ``render_mode`` is forwarded so ``--video`` + keeps working under the warp frontend. Fixed ^^^^^ diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py index 9da4f9c80e0e..23ea519f7cea 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py @@ -3,295 +3,304 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Adapt a stable manager-based RL env cfg to run on the warp runtime. +"""Run a stable env cfg on the warp runtime via a pluggable rule pipeline. + +The frontend classifies the cfg as manager-based or direct, runs the +applicable :class:`CompatRule` instances against it, and constructs the +matching warp env class. New incompatibilities (e.g. a new sensor type warp +can't run, a new term-cfg field that needs upgrading) are added by writing a +small :class:`CompatRule` subclass and passing it through the constructor — +no surgery in the dispatcher. Module-level imports are deliberately limited to lightweight stdlib modules. -Warp-library imports fire only inside :meth:`WarpFrontend.build`, which must -be called *after* :class:`SimulationApp` is alive — otherwise the warp lib -loading will race with Kit's USD/pxr extension initialisation. +Warp library imports fire only inside :meth:`WarpFrontend.build`, which must +be called *after* :class:`SimulationApp` is alive, otherwise warp's lib load +races with Kit's USD/pxr extension initialisation. """ from __future__ import annotations import importlib import logging -from collections.abc import Iterable +from abc import ABC, abstractmethod +from collections.abc import Iterable, Sequence from dataclasses import dataclass, field +from enum import Enum from types import ModuleType -from typing import Any +from typing import Any, ClassVar logger = logging.getLogger(__name__) -@dataclass -class MissingItem: - """A symbol the adapter looked for but couldn't resolve to a warp twin.""" - - group: str - """Where it was looked up — e.g. ``rewards``, ``observations.policy``, ``actions``.""" - term_name: str - """Field name on the group, e.g. ``track_lin_vel_xy_exp``.""" - expected_name: str - """The symbol name searched for in the warp modules (the stable ``__name__``).""" - kind: str - """Either ``"func"`` (mdp function) or ``"class_type"`` (action runtime class).""" - searched: tuple[str, ...] - """Module dotted-paths searched, in order.""" - action: str - """Outcome: ``"dropped"`` (term set to None) or ``"raised"`` (caller saw an error).""" +# --------------------------------------------------------------------------- +# Compatibility-rule framework +# --------------------------------------------------------------------------- -@dataclass -class WarpAdaptReport: - """Result of :meth:`WarpFrontend.adapt`. Empty ``missing`` means clean run.""" +class CfgKind(str, Enum): + """Workflow type of an env cfg the frontend is asked to adapt.""" - missing: list[MissingItem] = field(default_factory=list) + MANAGER_BASED = "manager_based" + DIRECT = "direct" - def __bool__(self) -> bool: # truthy when there is something to report - return bool(self.missing) - def format(self) -> str: - if not self.missing: - return "WarpFrontend: no missing items." - lines = [f"WarpFrontend: {len(self.missing)} missing warp twin(s):"] - for m in self.missing: - lines.append( - f" - {m.group}.{m.term_name}: stable {m.kind}={m.expected_name!r} → " - f"no twin in {list(m.searched)} ({m.action})" - ) - return "\n".join(lines) +_BOTH = frozenset({CfgKind.MANAGER_BASED, CfgKind.DIRECT}) +_MANAGER_ONLY = frozenset({CfgKind.MANAGER_BASED}) + +_WARP_MODULE_PREFIXES: tuple[str, ...] = ("isaaclab_experimental", "isaaclab_tasks_experimental") +"""Module-path prefixes that distinguish a real warp twin from a stable +re-export. The warp ``mdp`` packages do ``from isaaclab.envs.mdp import *``, +so a plain ``getattr`` lookup can return the stable symbol untouched.""" + +_TERM_GROUPS_FOR_PROMOTION: tuple[tuple[str, ...], ...] = ( + ("observations", "policy"), + ("events",), + ("rewards",), + ("terminations",), + ("commands",), + ("curriculum",), + ("actions",), +) +"""Cfg groups walked when promoting :class:`SceneEntityCfg` instances. Includes +``actions`` because action params often carry a :class:`SceneEntityCfg`.""" -class _NameSwap: - """Replace ``term.`` with the same-named symbol in ``modules``. +class IncompatibleEnvError(RuntimeError): + """Raised when the frontend can't run on the given cfg / task pair.""" - Drop the term entirely if no twin exists and ``drop_if_missing`` is set; - otherwise raise :class:`LookupError`. Either way, append a record to - ``report.missing`` so the caller can summarise what was lost. + def __init__(self, message: str, report: CompatReport | None = None): + super().__init__(message) + self.report = report - Used for both ``term.func`` (mdp swap) and ``action.class_type`` swaps — - the algorithm is identical, only the attribute name and policy differ. + +@dataclass +class CompatIssue: + """One thing the adapter looked for but couldn't resolve to a warp twin.""" + + rule: str + """Rule that produced this issue (matches :attr:`CompatRule.name`).""" + location: str + """Cfg path the issue was found at, e.g. ``rewards`` / ``actions``.""" + item: str + """Attribute name on the group, e.g. ``track_lin_vel_xy_exp``.""" + expected: str + """The symbol the rule searched for (the stable ``__name__``).""" + searched: tuple[str, ...] + """Module dotted-paths inspected, in order.""" + action: str + """Outcome: ``"dropped"``, ``"raised"``, or ``"left-stable"``.""" + + +@dataclass +class CompatReport: + """Outcome of running the rule pipeline against a cfg. + + ``changes`` is a human-readable list of mutations applied; ``issues`` + records every term the adapter couldn't resolve. A clean run has both + empty. """ - def __init__( - self, - modules: Iterable[ModuleType], - attr: str, - drop_if_missing: bool, - report: WarpAdaptReport, - group_label: str, - warp_module_prefixes: tuple[str, ...] = ("isaaclab_experimental", "isaaclab_tasks_experimental"), - ): - self._modules = tuple(modules) - self._attr = attr - self._drop = drop_if_missing - self._report = report - self._group = group_label - # The warp mdp __init__ re-exports stable mdp terms via ``from isaaclab.envs.mdp import *`` - # before adding its own overrides. ``getattr(warp_mdp, name)`` therefore can return the - # *stable* implementation untouched. We accept a candidate as a real warp twin only if - # its ``__module__`` lives under one of these prefixes. - self._warp_module_prefixes = warp_module_prefixes - - def apply_to(self, group: Any) -> None: - if group is None: - return - for name in (n for n in dir(group) if not n.startswith("_")): - term = getattr(group, name, None) - if term is None or not hasattr(term, self._attr): - continue - stable = getattr(term, self._attr) - if not callable(stable): - continue - twin = self._resolve_twin(stable.__name__) - if twin is not None: - setattr(term, self._attr, twin) - continue + rules_applied: list[str] = field(default_factory=list) + changes: list[str] = field(default_factory=list) + issues: list[CompatIssue] = field(default_factory=list) - # No twin found. - missing = MissingItem( - group=self._group, - term_name=name, - expected_name=stable.__name__, - kind=self._attr, - searched=tuple(m.__name__ for m in self._modules), - action="dropped" if self._drop else "raised", - ) - self._report.missing.append(missing) - if self._drop: - setattr(group, name, None) - else: - raise LookupError( - f"WarpFrontend: cannot adapt {self._group}.{name} — no warp twin for" - f" stable {self._attr}={stable.__name__!r}. Searched {missing.searched}." + def has_issues(self) -> bool: + return bool(self.issues) + + def has_fatal(self) -> bool: + return any(i.action == "raised" for i in self.issues) + + def __bool__(self) -> bool: + return bool(self.issues) or bool(self.changes) + + def format(self) -> str: + lines: list[str] = [] + if self.changes: + lines.append(f"WarpFrontend applied {len(self.changes)} change(s):") + lines.extend(f" - {c}" for c in self.changes) + if self.issues: + lines.append(f"WarpFrontend has {len(self.issues)} unresolved item(s):") + for i in self.issues: + lines.append( + f" - [{i.rule}] {i.location}.{i.item}: expected {i.expected!r}" + f" (searched {list(i.searched)}) → {i.action}" ) + return "\n".join(lines) if lines else "WarpFrontend: no changes." - def _resolve_twin(self, name: str) -> Any | None: - """Return the warp implementation of ``name`` from the search modules. - Skips candidates that are merely re-exports of stable code: the warp mdp - package does ``from isaaclab.envs.mdp import *`` and a ``getattr`` lookup - will happily return the stable symbol untouched. We confirm a candidate - is actually a warp implementation by checking its ``__module__``. - """ - for module in self._modules: - candidate = getattr(module, name, None) - if candidate is None: - continue - origin = getattr(candidate, "__module__", "") or "" - if origin.startswith(self._warp_module_prefixes): - return candidate - return None +@dataclass +class AdaptContext: + """Per-call context handed to each rule.""" + task_id: str + kind: CfgKind + strict: bool -class WarpFrontend: - """Adapt a stable env cfg in place; build a warp-runtime env. - Call :meth:`build` only after ``SimulationApp`` is alive — warp lib loads - lazily inside ``build``. All knobs are constructor args; subclass only if - you need to change the algorithm itself, not its parameters. +class CompatRule(ABC): + """One step in the warp adaptation pipeline. - Example:: + A rule mutates ``cfg`` in place and appends to ``report``. To handle a + new incompatibility, subclass and implement :meth:`apply`. Rules declare + which workflow they apply to via :attr:`applies_to`; the pipeline skips + any rule whose ``applies_to`` doesn't include the cfg kind. + """ - with launch_simulation(env_cfg, args_cli): - env = WarpFrontend().build(env_cfg, task_id=args_cli.task) + name: ClassVar[str] + applies_to: ClassVar[frozenset[CfgKind]] + + @abstractmethod + def apply(self, cfg: Any, ctx: AdaptContext, report: CompatReport) -> None: + """Mutate ``cfg`` in place and append issues / changes to ``report``.""" + + +# --------------------------------------------------------------------------- +# Helpers shared by rules +# --------------------------------------------------------------------------- + + +def _walk(root: Any, path: tuple[str, ...]) -> Any: + node = root + for attr in path: + node = getattr(node, attr, None) + if node is None: + return None + return node + + +def _resolve_warp_twin(name: str, modules: Sequence[ModuleType]) -> Any | None: + """Return ``name`` from ``modules`` if it lives under a warp prefix, else None.""" + for module in modules: + candidate = getattr(module, name, None) + if candidate is None: + continue + origin = getattr(candidate, "__module__", "") or "" + if origin.startswith(_WARP_MODULE_PREFIXES): + return candidate + return None + + +def _swap_named_attr( + group: Any, + location: str, + attr: str, + modules: Sequence[ModuleType], + *, + rule: str, + strict: bool, + report: CompatReport, +) -> None: + """Replace ``term.`` with the same-named callable from ``modules``. + + For each public attribute on ``group`` whose value has a callable ``attr`` + field, find a same-named symbol in ``modules`` whose ``__module__`` lives + under :data:`_WARP_MODULE_PREFIXES` and assign it. If no twin is found, + drop the term (set to ``None``) when ``strict`` is False, raise + :class:`LookupError` when True. Either way append a :class:`CompatIssue`. """ + if group is None: + return + searched = tuple(m.__name__ for m in modules) + # Snapshot the attribute names so mutating ``group`` mid-loop is safe. + for name in [n for n in dir(group) if not n.startswith("_")]: + term = getattr(group, name, None) + if term is None or not hasattr(term, attr): + continue + stable = getattr(term, attr) + if not callable(stable): + continue + twin = _resolve_warp_twin(stable.__name__, modules) + if twin is not None: + setattr(term, attr, twin) + continue + issue = CompatIssue( + rule=rule, + location=location, + item=name, + expected=stable.__name__, + searched=searched, + action="raised" if strict else "dropped", + ) + report.issues.append(issue) + if strict: + raise LookupError( + f"WarpFrontend: cannot adapt {location}.{name} — no warp twin for " + f"stable {attr}={stable.__name__!r}. Searched {list(searched)}." + ) + setattr(group, name, None) - DEFAULT_TERM_GROUPS: tuple[tuple[str, ...], ...] = ( - ("observations", "policy"), - ("events",), - ("rewards",), - ("terminations",), - ("curriculum",), - ("actions",), # included for SceneEntityCfg upgrade in action params - ) - def __init__( - self, - stable_root: str = "isaaclab_tasks", - warp_root: str = "isaaclab_tasks_experimental", - fallback_mdp: str = "isaaclab_experimental.envs.mdp", - warp_actions: str = "isaaclab_experimental.envs.mdp.actions.joint_actions", - drop_sensors: Iterable[str] = ("height_scanner",), - term_groups: Iterable[tuple[str, ...]] = DEFAULT_TERM_GROUPS, - strict: bool = False, - ): - self._stable_root = stable_root - self._warp_root = warp_root - self._fallback_mdp = fallback_mdp - self._warp_actions = warp_actions - self._drop_sensors = tuple(drop_sensors) - self._term_groups = tuple(term_groups) - self._strict = strict +# --------------------------------------------------------------------------- +# Concrete rules +# --------------------------------------------------------------------------- - # -- public --------------------------------------------------------------- - def adapt(self, cfg: Any, task_id: str) -> WarpAdaptReport: - """Mutate ``cfg`` in place; return a report of what couldn't be adapted. +class ResolvePhysicsPresetRule(CompatRule): + """Collapse ``cfg.sim.physics`` from a :class:`PresetCfg` to its ``newton`` field. - With ``strict=True``, any missing twin (term func or action class) raises - :class:`LookupError` immediately so the caller sees the failure in the - traceback instead of a silently-dropped term. - """ - report = WarpAdaptReport() - - # 1. PresetCfg → newton (also dodges the lazy PhysxCfg materialisation). - # Note: by the time we run, Hydra's `resolve_presets` has already - # collapsed any PresetCfg wrapper to a concrete preset. The auto - # ``presets=newton`` injected from train.py for ``--manager=warp`` - # ensures that concrete preset is the newton one. If the preset - # somehow still has a ``newton`` attribute (custom subclass), pick it. - physics = getattr(cfg.sim, "physics", None) + Hydra's preset resolution normally collapses the wrapper before the + frontend runs. This rule covers two residual cases: programmatic use of + :meth:`WarpFrontend.adapt` (no Hydra), and custom physics cfgs that still + expose a ``newton`` attribute. + """ + + name = "resolve_physics_preset" + applies_to = _BOTH + + def apply(self, cfg, ctx, report): + physics = getattr(getattr(cfg, "sim", None), "physics", None) if physics is not None and hasattr(physics, "newton"): cfg.sim.physics = physics.newton + report.changes.append(f"sim.physics → {type(physics.newton).__name__}") - # 2. Drop sensors warp can't run yet. - scene = getattr(cfg, "scene", None) - for sensor in self._drop_sensors: - if scene is not None and getattr(scene, sensor, None) is not None: - setattr(scene, sensor, None) - # 3. Upgrade SceneEntityCfg instances inside term.params dicts to the - # warp-extended variant (adds joint_mask / *_ids_wp wp.array fields - # that warp mdp kernels read). - self._upgrade_scene_entity_cfgs(cfg) - - # 4. Swap term.func across all groups. - warp_mdp_modules = self._mdp_modules_for(task_id) - # Always log which warp mdp modules were resolved so it's easy to debug - # missing-twin reports. - logger.info( - "WarpFrontend: warp mdp modules for %r → %s", - task_id, - [m.__name__ for m in warp_mdp_modules], - ) - for path in self._term_groups: - if path == ("actions",): - continue # actions don't have ``func`` — handled below - group = self._walk(cfg, path) - label = ".".join(path) - swap = _NameSwap( - warp_mdp_modules, "func", drop_if_missing=not self._strict, report=report, group_label=label - ) - swap.apply_to(group) +class DropUnsupportedSensorsRule(CompatRule): + """Drop scene sensors the warp runtime can't run yet.""" - # 5. Swap action class_type. An action with no warp class_type can't - # run on the warp manager, so this is always strict (raise). - action_swap = _NameSwap( - [importlib.import_module(self._warp_actions)], - attr="class_type", - drop_if_missing=False, - report=report, - group_label="actions", - ) - action_swap.apply_to(getattr(cfg, "actions", None)) - - # 6. Surface the report. - if report.missing: - level = logging.ERROR if self._strict else logging.WARNING - logger.log(level, report.format()) + name = "drop_unsupported_sensors" + applies_to = _BOTH - return report + def __init__(self, sensors: Iterable[str] = ("height_scanner",)): + self.sensors = tuple(sensors) - def build(self, cfg: Any, task_id: str): - """Adapt ``cfg`` and return a :class:`ManagerBasedRLEnvWarp` instance. + def apply(self, cfg, ctx, report): + scene = getattr(cfg, "scene", None) + if scene is None: + return + for sensor in self.sensors: + if getattr(scene, sensor, None) is not None: + setattr(scene, sensor, None) + report.changes.append(f"scene.{sensor} → None") - Always logs the missing-twin report (if any). With ``strict=True``, - any missing twin raises before the env is constructed. - """ - # Lazy: this is the first warp-lib load. Caller must already be inside - # the SimulationApp context, i.e. inside ``launch_simulation``. - from isaaclab_experimental.envs import ManagerBasedRLEnvWarp - self.adapt(cfg, task_id) - return ManagerBasedRLEnvWarp(cfg=cfg) +class PromoteSceneEntityCfgRule(CompatRule): + """Promote :class:`SceneEntityCfg` instances under term params to the warp variant. - # -- helpers -------------------------------------------------------------- + The warp variant adds cached ``joint_mask`` / ``joint_ids_wp`` / + ``body_ids_wp`` fields the warp mdp kernels read after :meth:`resolve`. + The class hierarchy is asserted at apply time — if the warp class no + longer subclasses the stable class (refactor or divergence), the rule + raises :class:`TypeError` so the assumption can never silently corrupt + instances. + """ - def _upgrade_scene_entity_cfgs(self, cfg: Any) -> None: - """Upgrade every stable :class:`SceneEntityCfg` to the warp variant. + name = "promote_scene_entity_cfg" + applies_to = _MANAGER_ONLY - Walks ``term.params`` dicts inside each term group and re-classes any - stable :class:`SceneEntityCfg` instance into the warp subclass so warp - mdp kernels can read ``joint_mask`` / ``joint_ids_wp`` / ``body_ids_wp`` - after :meth:`resolve`. - """ + def apply(self, cfg, ctx, report): from isaaclab.managers.scene_entity_cfg import SceneEntityCfg as _StableSE from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as _WarpSE - def upgrade(obj: Any) -> None: - if isinstance(obj, _WarpSE) or not isinstance(obj, _StableSE): - return - # In-place class promotion + initialise warp-only fields. The class - # hierarchy guarantees compatible memory layout (warp inherits stable). - obj.__class__ = _WarpSE - obj.joint_mask = None - obj.joint_ids_wp = None - obj.body_ids_wp = None - - for path in self._term_groups: - group = self._walk(cfg, path) + if not issubclass(_WarpSE, _StableSE): + raise TypeError( + "PromoteSceneEntityCfgRule requires the warp SceneEntityCfg to subclass the stable one;" + f" got mro {[c.__name__ for c in _WarpSE.__mro__]}" + ) + + promoted = 0 + for path in _TERM_GROUPS_FOR_PROMOTION: + group = _walk(cfg, path) if group is None: continue for name in (n for n in dir(group) if not n.startswith("_")): @@ -299,36 +308,86 @@ def upgrade(obj: Any) -> None: if term is None: continue params = getattr(term, "params", None) - if isinstance(params, dict): - for v in params.values(): - upgrade(v) + if not isinstance(params, dict): + continue + for value in params.values(): + if isinstance(value, _WarpSE) or not isinstance(value, _StableSE): + continue + value.__class__ = _WarpSE + value.joint_mask = None + value.joint_ids_wp = None + value.body_ids_wp = None + promoted += 1 + if promoted: + report.changes.append(f"promoted {promoted} SceneEntityCfg instance(s) to warp variant") + + +class SwapMdpFunctionsRule(CompatRule): + """Replace ``term.func`` with the same-named callable from warp ``mdp`` modules. + + Walks observation / event / reward / termination / command / curriculum + groups. A re-export from stable code is rejected by checking + ``__module__``. Missing twins are dropped (term set to ``None``) unless + the context is strict, in which case they raise. + """ - @staticmethod - def _walk(root: Any, path: tuple[str, ...]) -> Any: - node = root - for attr in path: - node = getattr(node, attr, None) - if node is None: - return None - return node + name = "swap_mdp_functions" + applies_to = _MANAGER_ONLY + + GROUPS: ClassVar[tuple[tuple[str, ...], ...]] = ( + ("observations", "policy"), + ("events",), + ("rewards",), + ("terminations",), + ("commands",), + ("curriculum",), + ) + + def __init__( + self, + stable_root: str = "isaaclab_tasks", + warp_root: str = "isaaclab_tasks_experimental", + fallback_mdp: str = "isaaclab_experimental.envs.mdp", + ): + self.stable_root = stable_root + self.warp_root = warp_root + self.fallback_mdp = fallback_mdp + + def apply(self, cfg, ctx, report): + modules = self._mdp_modules_for(ctx.task_id) + logger.info("WarpFrontend: warp mdp modules → %s", [m.__name__ for m in modules]) + for path in self.GROUPS: + group = _walk(cfg, path) + _swap_named_attr( + group, + location=".".join(path), + attr="func", + modules=modules, + rule=self.name, + strict=ctx.strict, + report=report, + ) def _mdp_modules_for(self, task_id: str) -> list[ModuleType]: """Locate warp mdp modules that mirror the stable task's mdp. - Convention: replace ``isaaclab_tasks`` with ``isaaclab_tasks_experimental`` - in the stable cfg's module path; walk up looking for an ``mdp`` - sub-module on the warp side; always append the package-wide warp mdp - as a fallback for cross-task base terms. + Convention: replace :attr:`stable_root` with :attr:`warp_root` in the + stable cfg's module path, walk up looking for an ``mdp`` sub-module on + the warp side, and always append :attr:`fallback_mdp` for cross-task + base terms. """ modules: list[ModuleType] = [] try: import gymnasium as gym entry = gym.spec(task_id).kwargs.get("env_cfg_entry_point") - except Exception: + except (KeyError, AttributeError, ModuleNotFoundError, gym.error.NameNotFound): + # Spec is missing / kwargs missing the entry — fall back to the + # package-wide warp mdp module only. Real ImportError from a + # broken cfg module is *not* swallowed; it propagates. entry = None - if isinstance(entry, str) and entry.startswith(self._stable_root): - warp_pkg = entry.rsplit(".", 1)[0].replace(self._stable_root, self._warp_root, 1) + if isinstance(entry, str) and entry.startswith(self.stable_root): + warp_pkg = entry.rsplit(".", 1)[0].replace(self.stable_root, self.warp_root, 1) parts = warp_pkg.split(".") for depth in range(len(parts), 0, -1): try: @@ -337,7 +396,162 @@ def _mdp_modules_for(self, task_id: str) -> list[ModuleType]: except ImportError: continue try: - modules.append(importlib.import_module(self._fallback_mdp)) + modules.append(importlib.import_module(self.fallback_mdp)) except ImportError: - logger.warning("WarpFrontend: fallback mdp module %r not importable", self._fallback_mdp) + logger.warning("WarpFrontend: fallback mdp module %r not importable", self.fallback_mdp) return modules + + +class SwapActionClassTypeRule(CompatRule): + """Swap ``cfg.actions..class_type`` with the same-named warp class. + + Always strict: an action with no warp class can't run on the warp + runtime, so a missing twin raises :class:`LookupError`. + """ + + name = "swap_action_class_type" + applies_to = _MANAGER_ONLY + + def __init__(self, warp_actions_module: str = "isaaclab_experimental.envs.mdp.actions.joint_actions"): + self.warp_actions_module = warp_actions_module + + def apply(self, cfg, ctx, report): + actions = getattr(cfg, "actions", None) + if actions is None: + return + try: + module = importlib.import_module(self.warp_actions_module) + except ImportError as exc: + raise IncompatibleEnvError( + f"WarpFrontend: warp action module {self.warp_actions_module!r} not importable" + ) from exc + _swap_named_attr( + actions, + location="actions", + attr="class_type", + modules=[module], + rule=self.name, + strict=True, + report=report, + ) + + +# --------------------------------------------------------------------------- +# Frontend +# --------------------------------------------------------------------------- + + +class WarpFrontend: + """Run a stable env cfg on the warp runtime via a pluggable rule pipeline. + + The frontend classifies the input cfg as manager-based or direct, runs + the applicable :class:`CompatRule` instances, then constructs the + matching warp env class. + + Direct envs encode their runtime in the env class itself (each warp + direct env is a separate class with its own kernels), so the frontend + doesn't transform direct cfgs — it only verifies the registered + ``entry_point`` lives under :data:`_WARP_MODULE_PREFIXES` and dispatches + via :func:`gym.make`. A stable direct cfg + ``--frontend=warp`` raises + :class:`IncompatibleEnvError` with a clear message. + + Example:: + + with launch_simulation(cfg, args_cli): + env = WarpFrontend().build(cfg, task_id=args_cli.task) + if env.unwrapped.warp_compat_report.has_issues(): + print(env.unwrapped.warp_compat_report.format()) + """ + + DEFAULT_RULES: ClassVar[tuple[type[CompatRule], ...]] = ( + ResolvePhysicsPresetRule, + DropUnsupportedSensorsRule, + PromoteSceneEntityCfgRule, + SwapMdpFunctionsRule, + SwapActionClassTypeRule, + ) + + def __init__(self, rules: Sequence[CompatRule] | None = None, strict: bool = False): + self.rules: list[CompatRule] = list(rules) if rules is not None else [r() for r in self.DEFAULT_RULES] + self.strict = strict + + # -- public --------------------------------------------------------------- + + def adapt(self, cfg: Any, task_id: str) -> CompatReport: + """Run the rule pipeline; mutate ``cfg`` in place; return a report. + + The report is logged at ``WARNING`` (or ``ERROR`` in strict mode when + any issue is fatal). Missing twins raise :class:`LookupError` only + when ``strict=True``. + """ + kind = self._classify(cfg) + ctx = AdaptContext(task_id=task_id, kind=kind, strict=self.strict) + report = CompatReport() + for rule in self.rules: + if kind not in rule.applies_to: + continue + rule.apply(cfg, ctx, report) + report.rules_applied.append(rule.name) + if report: + level = logging.ERROR if self.strict and report.has_fatal() else logging.WARNING + logger.log(level, report.format()) + return report + + def build(self, cfg: Any, task_id: str, render_mode: str | None = None): + """Adapt ``cfg`` and construct a warp env. + + Returns an env instance with ``warp_compat_report`` attached on the + unwrapped env. ``render_mode`` is forwarded so ``--video`` keeps + working when the frontend is selected at the train-script level. + """ + # Lazy: this is the first warp-lib load. The caller must already be + # inside the SimulationApp context (i.e. inside ``launch_simulation``). + from isaaclab.envs import DirectRLEnvCfg, ManagerBasedRLEnvCfg + + report = self.adapt(cfg, task_id) + if isinstance(cfg, ManagerBasedRLEnvCfg): + from isaaclab_experimental.envs import ManagerBasedRLEnvWarp + + env = ManagerBasedRLEnvWarp(cfg=cfg, render_mode=render_mode) + elif isinstance(cfg, DirectRLEnvCfg): + self._verify_direct_warp(task_id) + import gymnasium as gym + + env = gym.make(task_id, cfg=cfg, render_mode=render_mode) + else: + raise IncompatibleEnvError( + f"WarpFrontend supports ManagerBasedRLEnvCfg or DirectRLEnvCfg, got {type(cfg).__name__}", + report=report, + ) + env.unwrapped.warp_compat_report = report + return env + + # -- helpers -------------------------------------------------------------- + + @staticmethod + def _classify(cfg: Any) -> CfgKind: + from isaaclab.envs import DirectRLEnvCfg, ManagerBasedRLEnvCfg + + if isinstance(cfg, ManagerBasedRLEnvCfg): + return CfgKind.MANAGER_BASED + if isinstance(cfg, DirectRLEnvCfg): + return CfgKind.DIRECT + raise IncompatibleEnvError( + f"WarpFrontend supports ManagerBasedRLEnvCfg or DirectRLEnvCfg, got {type(cfg).__name__}" + ) + + @staticmethod + def _verify_direct_warp(task_id: str) -> None: + import gymnasium as gym + + try: + spec = gym.spec(task_id) + except gym.error.NameNotFound as exc: + raise IncompatibleEnvError(f"Task {task_id!r} is not registered with gymnasium") from exc + ep = spec.entry_point + if isinstance(ep, str) and not ep.startswith(_WARP_MODULE_PREFIXES): + raise IncompatibleEnvError( + f"Direct env {task_id!r} entry_point {ep!r} is not a warp implementation." + f" --frontend=warp on a direct task requires the registered class to live under" + f" {list(_WARP_MODULE_PREFIXES)} (e.g. *-Direct-Warp-v0)." + ) From 0a9f2cd99e08378369bfa6fc7ae73a4b66ee4954 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 5 May 2026 17:45:36 +0000 Subject: [PATCH 06/58] Fix presets=newton injection for stable manager-based tasks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The earlier check inspected gym.spec(task).entry_point, which for stable manager-based tasks is "isaaclab.envs:ManagerBasedRLEnv" — the env class path, not the cfg path. So the startswith("isaaclab_tasks.manager_based") test always failed and presets=newton was never injected. Hydra then resolved every PresetCfg in the cfg tree (physics, contact_forces, etc.) to its default field, leaving the warp runtime with PhysX class_types it can't load. Switch to spec.kwargs["env_cfg_entry_point"], which actually points at the task cfg module (e.g. "isaaclab_tasks.manager_based.locomotion.velocity. config.anymal_d.flat_env_cfg:AnymalDFlatEnvCfg"), and the prefix check selects the right tasks. --- scripts/reinforcement_learning/rsl_rl/train.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/scripts/reinforcement_learning/rsl_rl/train.py b/scripts/reinforcement_learning/rsl_rl/train.py index 15326eb4b196..9bb76d884555 100644 --- a/scripts/reinforcement_learning/rsl_rl/train.py +++ b/scripts/reinforcement_learning/rsl_rl/train.py @@ -102,17 +102,20 @@ # When the warp frontend is selected on a stable manager-based task, the cfg # must resolve any PresetCfg wrappers to their ``newton`` field. Hydra # resolves presets *before* the WarpFrontend runs, so we inject -# ``presets=newton`` here. We only inject for tasks registered under -# ``isaaclab_tasks.manager_based`` — direct tasks and pre-warp registrations -# don't carry a preset system, and injecting ``presets=newton`` against them -# causes Hydra to error before the frontend can produce its own diagnostic. +# ``presets=newton`` here. We only inject for tasks whose env_cfg_entry_point +# is under ``isaaclab_tasks.manager_based``; direct tasks and the pre-warp +# ``*-Warp-v0`` registrations don't carry the preset system, and injecting +# ``presets=newton`` against them causes Hydra to error before the frontend +# can produce its own diagnostic. Note: ``spec.entry_point`` is the env +# *class* path (e.g. ``isaaclab.envs:ManagerBasedRLEnv``) — we check the +# *cfg* entry point in ``spec.kwargs``. if args_cli.frontend == "warp" and args_cli.task is not None: try: _spec = gym.spec(args_cli.task) except gym.error.NameNotFound: _spec = None - _entry = _spec.entry_point if _spec is not None else None - _is_stable_manager = isinstance(_entry, str) and _entry.startswith("isaaclab_tasks.manager_based") + _cfg_entry = _spec.kwargs.get("env_cfg_entry_point") if _spec is not None else None + _is_stable_manager = isinstance(_cfg_entry, str) and _cfg_entry.startswith("isaaclab_tasks.manager_based") _explicit_preset = next((a for a in remaining_args if a.startswith("presets=")), None) if _is_stable_manager and _explicit_preset is None: remaining_args.append("presets=newton") From d558741194c923a2357c9f15c7c26b08efa0fdc4 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 6 May 2026 07:40:55 +0000 Subject: [PATCH 07/58] Refactor warp bridge into a frontend framework package The single-file warp_frontend.py grew into a real subsystem worth splitting out, so move it into a frontend/ package with explicit abstractions: - frontend/base.py: Frontend ABC, CompatRule (check + transform via a unified run() method), TaskResolver (centralised gym.spec introspection -> TaskMeta), Workflow / Runtime / Severity enums, Issue / Change / Report record types, register_frontend / get_frontend registry. Helpers walk_attrs / resolve_warp_twin / iter_term_attrs are shared utilities used by rules. - frontend/torch.py: TorchFrontend, the default. Pass-through to gym.make with one rule (WarnIfTaskIsWarpRegistered) for the contradiction case. - frontend/warp.py: WarpFrontend with the full rule pipeline. Includes a new CheckPhysicsIsNewton blocking rule that surfaces the PhysX-with-warp incompatibility (asset class_type strings resolve to isaaclab_physx.* classes that depend on omni.physics.tensors.api, which the warp runtime does not initialise). CLI: rename --frontend stable -> torch since the axis is *runtime*, not *stability tier*. The frontend selector now reads cleanly: --frontend torch -> default gym.make path --frontend warp -> experimental warp runtime via WarpFrontend train.py becomes a thin dispatcher: get_frontend(name) gives a Frontend instance, frontend.preprocess_hydra_args(...) handles preset injection, frontend.build(cfg, task) returns the env. No more inline conditional imports; no more inline preset-injection logic. env.unwrapped.frontend_report is the inspection point - callers and tests can read what changed and what was missing without re-running adapt(). To add a new compatibility check, write a CompatRule subclass and append it to the relevant frontend's `rules` tuple. To add a new runtime, subclass Frontend and call register_frontend(name, cls). --- .../reinforcement_learning/rsl_rl/train.py | 63 +- .../changelog.d/warp-manager-bridge.rst | 73 ++- .../envs/frontend/__init__.py | 100 ++++ .../envs/frontend/base.py | 486 +++++++++++++++ .../envs/frontend/torch.py | 72 +++ .../envs/frontend/warp.py | 453 ++++++++++++++ .../envs/warp_frontend.py | 557 ------------------ 7 files changed, 1186 insertions(+), 618 deletions(-) create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/frontend/__init__.py create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/frontend/torch.py create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py delete mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py diff --git a/scripts/reinforcement_learning/rsl_rl/train.py b/scripts/reinforcement_learning/rsl_rl/train.py index 9bb76d884555..e146a0cc2575 100644 --- a/scripts/reinforcement_learning/rsl_rl/train.py +++ b/scripts/reinforcement_learning/rsl_rl/train.py @@ -70,13 +70,14 @@ parser.add_argument( "--frontend", type=str, - default="stable", - choices=["stable", "warp"], + default="torch", + choices=["torch", "warp"], help=( - "Runtime backend for the env. 'stable' uses isaaclab.envs.* (torch);" - " 'warp' routes through isaaclab_experimental.envs.warp_frontend.WarpFrontend," - " which adapts a manager-based stable cfg onto ManagerBasedRLEnvWarp or dispatches" - " a direct task to its registered warp env class." + "Runtime backend for the env. 'torch' uses isaaclab.envs.* (PhysX or Newton via" + " factory dispatch); 'warp' routes through the WarpFrontend, which adapts a" + " manager-based cfg onto ManagerBasedRLEnvWarp or dispatches a direct task to its" + " registered warp env class. See isaaclab_experimental.envs.frontend for the" + " pluggable rule pipeline." ), ) cli_args.add_rsl_rl_args(parser) @@ -99,32 +100,15 @@ # argparser and (optionally) the external callback function. remaining_args = list_intersection(remaining_args, remaining_args_env_registration) -# When the warp frontend is selected on a stable manager-based task, the cfg -# must resolve any PresetCfg wrappers to their ``newton`` field. Hydra -# resolves presets *before* the WarpFrontend runs, so we inject -# ``presets=newton`` here. We only inject for tasks whose env_cfg_entry_point -# is under ``isaaclab_tasks.manager_based``; direct tasks and the pre-warp -# ``*-Warp-v0`` registrations don't carry the preset system, and injecting -# ``presets=newton`` against them causes Hydra to error before the frontend -# can produce its own diagnostic. Note: ``spec.entry_point`` is the env -# *class* path (e.g. ``isaaclab.envs:ManagerBasedRLEnv``) — we check the -# *cfg* entry point in ``spec.kwargs``. -if args_cli.frontend == "warp" and args_cli.task is not None: - try: - _spec = gym.spec(args_cli.task) - except gym.error.NameNotFound: - _spec = None - _cfg_entry = _spec.kwargs.get("env_cfg_entry_point") if _spec is not None else None - _is_stable_manager = isinstance(_cfg_entry, str) and _cfg_entry.startswith("isaaclab_tasks.manager_based") - _explicit_preset = next((a for a in remaining_args if a.startswith("presets=")), None) - if _is_stable_manager and _explicit_preset is None: - remaining_args.append("presets=newton") - elif _is_stable_manager and _explicit_preset != "presets=newton": - logger.warning( - "--frontend=warp on %r expects presets=newton; got %r — adapter may fail to find a Newton physics cfg.", - args_cli.task, - _explicit_preset, - ) +# Build the chosen frontend and let it pre-process Hydra args. Frontends +# may inject preset selections (e.g. WarpFrontend appends ``presets=newton`` +# for stable manager-based tasks so PresetCfg wrappers resolve to the Newton +# field before Hydra builds the cfg). +from isaaclab_experimental.envs.frontend import get_frontend # noqa: E402 + +_frontend = get_frontend(args_cli.frontend) +if args_cli.task is not None: + remaining_args = _frontend.preprocess_hydra_args(args_cli.task, remaining_args) sys.argv = [sys.argv[0]] + remaining_args @@ -210,16 +194,13 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen # set the log directory for the environment (works for all environment types) env_cfg.log_dir = log_dir - # create isaac environment + # create isaac environment via the chosen frontend's rule pipeline. + # The frontend may mutate env_cfg in place (preset resolution, + # SceneEntityCfg promotion, mdp twin swaps, ...) before constructing + # the env. ``env.unwrapped.frontend_report`` carries what changed + # and what was missing. render_mode = "rgb_array" if args_cli.video else None - if args_cli.frontend == "warp": - # Lazy: first warp-side import. Calling this after SimulationApp - # is already alive avoids racing pxr extension init. - from isaaclab_experimental.envs.warp_frontend import WarpFrontend - - env = WarpFrontend().build(env_cfg, task_id=args_cli.task, render_mode=render_mode) - else: - env = gym.make(args_cli.task, cfg=env_cfg, render_mode=render_mode) + env = _frontend.build(env_cfg, task_id=args_cli.task, render_mode=render_mode) # convert to single-agent instance if required by the RL algorithm if isinstance(env.unwrapped.cfg, DirectMARLEnvCfg): diff --git a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst index 881450b1312d..d7197a0a756e 100644 --- a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst +++ b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst @@ -1,26 +1,59 @@ Added ^^^^^ -* Added :class:`~isaaclab_experimental.envs.warp_frontend.WarpFrontend`, a - runtime adapter that lets any stable manager-based RL task config run on the - experimental warp runtime (:class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp`) - without a parallel ``-Warp-v0`` registration. The adapter is built on a - pluggable :class:`~isaaclab_experimental.envs.warp_frontend.CompatRule` - pipeline; new incompatibilities (sensor types, term-cfg fields, action - classes) are added by writing a small rule subclass instead of editing the - dispatcher. The default rules cover physics-preset resolution, dropping - unsupported sensors, in-place :class:`SceneEntityCfg` promotion, mdp - function swaps, and action-class swaps. The frontend also dispatches - direct envs by verifying their registered entry-point class lives under - ``isaaclab_experimental`` / ``isaaclab_tasks_experimental`` and routing - through :func:`gym.make` unchanged. - -* Added a ``--frontend={stable,warp}`` flag to ``rsl_rl/train.py``. When set - to ``warp`` the script auto-injects ``presets=newton`` (so Hydra picks the - Newton physics preset before the adapter runs), warns on conflicting - ``presets=`` overrides, and dispatches the env through ``WarpFrontend`` - instead of :func:`gym.make`. ``render_mode`` is forwarded so ``--video`` - keeps working under the warp frontend. +* Added :mod:`isaaclab_experimental.envs.frontend`, a small framework that + selects the runtime backend for IsaacLab tasks. A :class:`Frontend` takes + a stable env cfg + task id, runs a pluggable :class:`CompatRule` pipeline + against the (cfg, task, frontend) triple, and constructs the env on the + matching runtime. New runtimes plug in by subclassing :class:`Frontend` + and calling :func:`register_frontend`; new compatibility checks plug in + by subclassing :class:`CompatRule` and listing it in a frontend's + :attr:`Frontend.rules`. + + The framework ships two built-in frontends: + + * :class:`TorchFrontend` (``--frontend=torch``, default) — passes through + to :func:`gym.make`. Emits a warning if the task is registered under + the warp runtime. + * :class:`WarpFrontend` (``--frontend=warp``) — adapts a manager-based + stable cfg onto :class:`ManagerBasedRLEnvWarp` via the rule pipeline, + or dispatches a direct task to its registered warp env class. + + The default warp rules are: + + * :class:`CheckPhysicsIsNewton` — blocking check; PhysX physics with the + warp runtime is a hard incompatibility (``isaaclab_physx`` classes + depend on ``omni.physics.tensors``, which the warp runtime does not + initialise). + * :class:`ResolvePhysicsPreset` — collapses ``PresetCfg`` to its + ``newton`` field for programmatic callers (Hydra's preset resolution + handles the CLI case). + * :class:`DropUnsupportedSensors` — drops scene sensors warp can't run. + * :class:`PromoteSceneEntityCfg` — in-place class promotion for warp + cached fields, with an :func:`issubclass` assertion to fail loudly if + the warp class ever stops subclassing the stable one. + * :class:`SwapMdpFunctions` — name-based ``term.func`` replacement + against the warp ``mdp`` modules; rejects stable re-exports by + inspecting ``__module__``. + * :class:`SwapActionClassType` — strict swap of action ``class_type`` + to the warp twin; missing twin is :attr:`Severity.BLOCKING`. + * :class:`VerifyDirectIsWarp` — for direct cfgs, blocks if the + registered entry-point isn't a warp class. + + All rules emit :class:`Issue` (incompatibility) and / or :class:`Change` + (transformation applied) records into a :class:`Report` accessible on + ``env.unwrapped.frontend_report`` after :meth:`Frontend.build`. + + :class:`TaskResolver` centralises ``gym.spec`` introspection and + classifies a task into a :class:`TaskMeta` (workflow, registered + runtime, cfg class). Rules read from :class:`TaskMeta` rather than + poking gym directly. + +* Added a ``--frontend={torch,warp}`` flag to ``rsl_rl/train.py``. The + script delegates Hydra-arg pre-processing (``presets=newton`` injection, + conflicting-preset warnings) and env construction to the selected + frontend via :func:`get_frontend`; ``render_mode`` is forwarded so + ``--video`` keeps working under the warp frontend. Fixed ^^^^^ diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/__init__.py new file mode 100644 index 000000000000..ad81605e6442 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/__init__.py @@ -0,0 +1,100 @@ +# 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 + +"""Frontend framework: pluggable runtime selectors for IsaacLab tasks. + +A frontend (chosen via ``--frontend {torch,warp}``) is a thin dispatcher that +takes a stable env cfg + task id, runs a pluggable :class:`CompatRule` +pipeline against the (cfg, task, frontend) triple, and constructs the env on +the matching runtime. New runtimes plug in by subclassing :class:`Frontend` +and calling :func:`register_frontend`; new compatibility checks plug in by +subclassing :class:`CompatRule` and adding it to a frontend's +:attr:`Frontend.rules`. + +Public API:: + + from isaaclab_experimental.envs.frontend import get_frontend + frontend = get_frontend("warp") + env = frontend.build(env_cfg, task_id="Isaac-Cartpole-v0") + print(env.unwrapped.frontend_report.format()) +""" + +from __future__ import annotations + +from .base import ( + WARP_ROOT_PREFIXES, + Change, + CompatRule, + Frontend, + FrontendIncompatibleError, + Issue, + Report, + ResolveContext, + Runtime, + Severity, + TaskMeta, + TaskResolver, + Workflow, + available_frontends, + get_frontend, + iter_term_attrs, + register_frontend, + resolve_warp_twin, + walk_attrs, +) +from .torch import TorchFrontend, WarnIfTaskIsWarpRegistered +from .warp import ( + CheckPhysicsIsNewton, + DropUnsupportedSensors, + PromoteSceneEntityCfg, + ResolvePhysicsPreset, + SwapActionClassType, + SwapMdpFunctions, + VerifyDirectIsWarp, + WarpFrontend, +) + +# Register built-ins so ``get_frontend("torch" | "warp")`` works without +# the caller importing the concrete classes. +register_frontend(TorchFrontend.name, TorchFrontend) +register_frontend(WarpFrontend.name, WarpFrontend) + + +__all__ = [ + # core abstractions + "Frontend", + "CompatRule", + "TaskMeta", + "TaskResolver", + "Report", + "Issue", + "Change", + "Severity", + "Workflow", + "Runtime", + "ResolveContext", + "FrontendIncompatibleError", + "WARP_ROOT_PREFIXES", + # registry + "register_frontend", + "get_frontend", + "available_frontends", + # rule helpers + "walk_attrs", + "resolve_warp_twin", + "iter_term_attrs", + # built-in frontends + "TorchFrontend", + "WarpFrontend", + # built-in rules + "WarnIfTaskIsWarpRegistered", + "CheckPhysicsIsNewton", + "ResolvePhysicsPreset", + "DropUnsupportedSensors", + "PromoteSceneEntityCfg", + "SwapMdpFunctions", + "SwapActionClassType", + "VerifyDirectIsWarp", +] diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py new file mode 100644 index 000000000000..81fc21a33cf7 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py @@ -0,0 +1,486 @@ +# 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 + +"""Frontend framework — abstractions for runtime selection. + +A *frontend* is a user-facing runtime selector (chosen via ``--frontend +{torch,warp}``). It takes a stable env cfg and a registered gym task id and +returns a runnable env, having first run a pluggable pipeline of +compatibility checks and transforms against the (cfg, task, frontend) +triple. + +The framework is intentionally narrow: + +* :class:`TaskResolver` is the single point that reads ``gym.spec`` and + classifies a task into a :class:`TaskMeta` (workflow, registered + runtime, cfg class). Rules and frontends consume :class:`TaskMeta` + instead of poking gym directly, so the registration format can evolve + in one place. +* :class:`CompatRule` is a single check or transform applied during + resolution. Each rule yields :class:`Issue` records (incompatibilities + the frontend can't paper over) and / or :class:`Change` records + (transformations applied to the cfg). Subclasses implement + :meth:`CompatRule.run` and may override :meth:`CompatRule.applies_to` + to scope themselves by workflow / runtime. +* :class:`Frontend` is the dispatcher. Subclasses declare a name, a list + of rule classes, and a :meth:`Frontend.construct` strategy that builds + the env once rules pass. +* :func:`register_frontend` / :func:`get_frontend` form the registry CLI + hooks read. + +To add a new compatibility check, write a small :class:`CompatRule` and +list it in the relevant frontend's :attr:`Frontend.rules`. To add a new +runtime, subclass :class:`Frontend` and call :func:`register_frontend`. +""" + +from __future__ import annotations + +import logging +from abc import ABC, abstractmethod +from collections.abc import Iterable, Iterator +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, ClassVar + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + + +WARP_ROOT_PREFIXES: tuple[str, ...] = ("isaaclab_experimental", "isaaclab_tasks_experimental") +"""Module prefixes that mark a class as a warp implementation. + +The warp ``mdp`` packages do ``from isaaclab.envs.mdp import *``, so a +plain ``getattr`` lookup can return a stable symbol untouched. We accept +a candidate as a real warp twin only if its ``__module__`` lives under +one of these prefixes.""" + + +# --------------------------------------------------------------------------- +# Enums +# --------------------------------------------------------------------------- + + +class Workflow(str, Enum): + """Task workflow type.""" + + MANAGER_BASED = "manager_based" + DIRECT = "direct" + DIRECT_MARL = "direct_marl" + + +class Runtime(str, Enum): + """Runtime backend a task is registered against. + + A task whose ``entry_point`` lives under :data:`WARP_ROOT_PREFIXES` is + classified as ``WARP``. Tasks registered against ``isaaclab.envs.*`` + (the standard manager-based and direct env classes) are classified + as ``TORCH``; those env classes use ``FactoryBase`` to dispatch on + the active physics backend at construction time, so PhysX and Newton + physics both flow through the torch runtime path. + """ + + TORCH = "torch" + WARP = "warp" + UNKNOWN = "unknown" + + +class Severity(str, Enum): + """Severity of a compatibility :class:`Issue`.""" + + BLOCKING = "blocking" + """Frontend cannot build the env. :meth:`Frontend.build` raises.""" + WARNING = "warning" + """Surfaced to the user but not fatal (e.g. a term was dropped).""" + + +# --------------------------------------------------------------------------- +# Records +# --------------------------------------------------------------------------- + + +@dataclass +class TaskMeta: + """Static description of a registered gym task, produced by :class:`TaskResolver`.""" + + task_id: str + """Gym task id, e.g. ``Isaac-Cartpole-v0``.""" + spec: Any + """The :class:`gymnasium.envs.registration.EnvSpec` returned by ``gym.spec``.""" + cfg_class: type + """Concrete cfg type passed to the frontend.""" + workflow: Workflow + """Manager-based / direct / direct-marl.""" + runtime: Runtime + """Runtime the task is registered against (independent of the frontend asked for).""" + + @property + def env_cfg_entry_point(self) -> str | None: + """Module path of the cfg class as registered in ``spec.kwargs``.""" + return self.spec.kwargs.get("env_cfg_entry_point") if self.spec is not None else None + + @property + def entry_point(self) -> Any: + """Env class entry point as registered (string or class).""" + return self.spec.entry_point if self.spec is not None else None + + +@dataclass +class Issue: + """An incompatibility surfaced by a rule.""" + + rule: str + """Rule that produced this issue (matches :attr:`CompatRule.name`).""" + severity: Severity + message: str + """One-line human-readable description.""" + location: str = "" + """Optional cfg path the issue was found at (e.g. ``rewards.foo``).""" + detail: dict[str, Any] = field(default_factory=dict) + """Free-form rule-specific data (searched modules, expected names, ...).""" + + +@dataclass +class Change: + """A transformation a rule applied to the cfg.""" + + rule: str + description: str + + +@dataclass +class Report: + """Outcome of a frontend's rule pipeline.""" + + frontend: str + """Name of the frontend that produced this report.""" + task: TaskMeta + rules_run: list[str] = field(default_factory=list) + issues: list[Issue] = field(default_factory=list) + changes: list[Change] = field(default_factory=list) + + @property + def blocking(self) -> list[Issue]: + return [i for i in self.issues if i.severity == Severity.BLOCKING] + + def has_blocking(self) -> bool: + return any(i.severity == Severity.BLOCKING for i in self.issues) + + def __bool__(self) -> bool: + return bool(self.changes) or bool(self.issues) + + def format(self) -> str: + lines = [ + f"Frontend {self.frontend!r} on task {self.task.task_id!r}" + f" ({self.task.workflow.value} / registered runtime {self.task.runtime.value})" + ] + if self.changes: + lines.append(f" changes ({len(self.changes)}):") + lines.extend(f" [{c.rule}] {c.description}" for c in self.changes) + if self.issues: + lines.append(f" issues ({len(self.issues)}):") + for i in self.issues: + head = f" [{i.severity.value}/{i.rule}]" + where = f" {i.location}" if i.location else "" + lines.append(f"{head}{where} {i.message}") + return "\n".join(lines) + + +@dataclass +class ResolveContext: + """Per-call context handed to rules.""" + + frontend: str + """Name of the running frontend.""" + task: TaskMeta + strict: bool = False + """When True, rules may escalate warnings to blocking issues.""" + + +# --------------------------------------------------------------------------- +# Exceptions +# --------------------------------------------------------------------------- + + +class FrontendIncompatibleError(RuntimeError): + """Raised when a frontend cannot build the requested (task, cfg) pair.""" + + def __init__(self, message: str, report: Report | None = None): + super().__init__(message) + self.report = report + + +# --------------------------------------------------------------------------- +# CompatRule +# --------------------------------------------------------------------------- + + +class CompatRule(ABC): + """A single check / transform applied during frontend resolution. + + A rule's :meth:`run` method may yield :class:`Issue` records (the rule + detected an incompatibility), :class:`Change` records (the rule mutated + the cfg), or both. Pure-check and pure-transform rules are common; mixed + rules are useful when a transformation has a partial-failure mode (e.g. + a name swap that found 5 of 6 twins). + + Subclass and: + + * Set :attr:`name` to a short identifier used in the report. + * Implement :meth:`run` to do the work. + * Override :meth:`applies_to` to skip the rule when irrelevant. + """ + + name: ClassVar[str] + + def applies_to(self, ctx: ResolveContext) -> bool: + """Return True if this rule should run on the given context. + + Default: always applies. Override to scope by ``ctx.task.workflow``, + ``ctx.task.runtime``, ``ctx.frontend``, etc. + """ + return True + + @abstractmethod + def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: + """Inspect / mutate ``cfg`` and yield :class:`Issue` / :class:`Change` records.""" + + +# --------------------------------------------------------------------------- +# TaskResolver +# --------------------------------------------------------------------------- + + +class TaskResolver: + """Inspect a registered gym task and classify its workflow + runtime. + + All rules and frontends should call :meth:`resolve` rather than reading + ``gym.spec`` directly. Centralising the read keeps registration-format + knowledge in one place. + """ + + @classmethod + def resolve(cls, task_id: str, cfg: Any) -> TaskMeta: + """Return a :class:`TaskMeta` for ``(task_id, cfg)``.""" + import gymnasium as gym + + try: + spec = gym.spec(task_id) + except gym.error.NameNotFound: + spec = None + return TaskMeta( + task_id=task_id, + spec=spec, + cfg_class=type(cfg), + workflow=cls._classify_workflow(cfg), + runtime=cls._classify_runtime(spec), + ) + + @staticmethod + def _classify_workflow(cfg: Any) -> Workflow: + from isaaclab.envs import DirectMARLEnvCfg, DirectRLEnvCfg, ManagerBasedRLEnvCfg + + if isinstance(cfg, ManagerBasedRLEnvCfg): + return Workflow.MANAGER_BASED + if isinstance(cfg, DirectMARLEnvCfg): + return Workflow.DIRECT_MARL + if isinstance(cfg, DirectRLEnvCfg): + return Workflow.DIRECT + raise FrontendIncompatibleError( + f"Cfg type {type(cfg).__name__!r} is not a recognised env cfg" + f" (expected ManagerBasedRLEnvCfg / DirectRLEnvCfg / DirectMARLEnvCfg subclass)." + ) + + @staticmethod + def _classify_runtime(spec: Any) -> Runtime: + if spec is None: + return Runtime.UNKNOWN + ep = spec.entry_point + if isinstance(ep, str): + if ep.startswith(WARP_ROOT_PREFIXES): + return Runtime.WARP + if ep.startswith(("isaaclab.envs", "isaaclab_tasks")): + return Runtime.TORCH + return Runtime.UNKNOWN + + +# --------------------------------------------------------------------------- +# Frontend +# --------------------------------------------------------------------------- + + +class Frontend(ABC): + """User-facing runtime selector. + + Subclasses declare: + + * :attr:`name`: CLI identifier used by ``--frontend`` and the registry. + * :attr:`rules`: rule classes the pipeline instantiates by default. + * :meth:`construct`: how to build the env once rules pass. + + Subclasses may also override: + + * :meth:`preprocess_hydra_args`: CLI-time Hydra arg munging (e.g. inject + a preset selection so Hydra resolves ``PresetCfg`` to the right field + before the cfg ever reaches us). + * :meth:`preprocess_cfg`: cfg pre-processing run before rules. + """ + + name: ClassVar[str] + rules: ClassVar[tuple[type[CompatRule], ...]] = () + + def __init__(self, rules: Iterable[CompatRule] | None = None, strict: bool = False): + self._rules: list[CompatRule] = [r() for r in type(self).rules] if rules is None else list(rules) + self.strict = strict + + # -- public --------------------------------------------------------------- + + def build(self, cfg: Any, task_id: str, **construct_kwargs: Any) -> Any: + """Resolve the task, run rules, and construct the env. + + Raises :class:`FrontendIncompatibleError` if any rule emits a + :class:`Severity.BLOCKING` issue. The returned env carries the + :class:`Report` on ``env.unwrapped.frontend_report`` for inspection. + """ + report = self.resolve(cfg, task_id) + self._log_report(report) + if report.has_blocking(): + raise FrontendIncompatibleError( + f"Frontend {self.name!r} cannot build {task_id!r}:\n{report.format()}", + report=report, + ) + env = self.construct(cfg, report.task, **construct_kwargs) + env.unwrapped.frontend_report = report + return env + + def resolve(self, cfg: Any, task_id: str) -> Report: + """Run the rule pipeline; mutate ``cfg`` in place; return a :class:`Report`. + + Use this directly for dry-run validation (tests, CLI ``--check``). + """ + meta = TaskResolver.resolve(task_id, cfg) + ctx = ResolveContext(frontend=self.name, task=meta, strict=self.strict) + report = Report(frontend=self.name, task=meta) + self.preprocess_cfg(cfg, ctx) + for rule in self._rules: + if not rule.applies_to(ctx): + continue + report.rules_run.append(rule.name) + for record in rule.run(cfg, ctx): + if isinstance(record, Issue): + report.issues.append(record) + elif isinstance(record, Change): + report.changes.append(record) + else: + raise TypeError(f"Rule {rule.name!r} yielded {type(record).__name__}; expected Issue or Change.") + return report + + # -- subclass hooks ------------------------------------------------------- + + @abstractmethod + def construct(self, cfg: Any, meta: TaskMeta, **kwargs: Any) -> Any: + """Build the env. Called only when no blocking issues were raised.""" + + def preprocess_cfg(self, cfg: Any, ctx: ResolveContext) -> None: + """Subclass hook for cfg pre-processing run before rules. Default: no-op.""" + + def preprocess_hydra_args(self, task_id: str, args: list[str]) -> list[str]: + """Subclass hook for CLI-time Hydra arg munging. Default: no-op. + + Called from the train script so frontends can inject preset + selections (or refuse incompatible ones) before Hydra builds the + cfg. Returns the (possibly modified) args list. + """ + return args + + # -- helpers -------------------------------------------------------------- + + def _log_report(self, report: Report) -> None: + if not report: + return + level = logging.ERROR if report.has_blocking() else logging.WARNING + logger.log(level, report.format()) + + +# --------------------------------------------------------------------------- +# Registry +# --------------------------------------------------------------------------- + + +_REGISTRY: dict[str, type[Frontend]] = {} + + +def register_frontend(name: str, cls: type[Frontend]) -> None: + """Register a frontend implementation under ``name``. + + The CLI ``--frontend`` flag and :func:`get_frontend` look up by this name. + Re-registering the same class is idempotent; re-registering a different + class raises :class:`ValueError`. + """ + existing = _REGISTRY.get(name) + if existing is not None and existing is not cls: + raise ValueError( + f"Frontend {name!r} is already registered to {existing.__name__};" + f" refusing to re-register to {cls.__name__}." + ) + _REGISTRY[name] = cls + + +def get_frontend(name: str, **kwargs: Any) -> Frontend: + """Look up a frontend by name and return a fresh instance. + + Extra kwargs are forwarded to :meth:`Frontend.__init__` (e.g. ``strict=True``). + """ + try: + cls = _REGISTRY[name] + except KeyError as exc: + raise ValueError(f"Unknown frontend {name!r}; available: {available_frontends()}.") from exc + return cls(**kwargs) + + +def available_frontends() -> list[str]: + """Return registered frontend names, sorted.""" + return sorted(_REGISTRY) + + +# --------------------------------------------------------------------------- +# Helpers shared by rules +# --------------------------------------------------------------------------- + + +def walk_attrs(root: Any, path: tuple[str, ...]) -> Any: + """Walk ``root.....`` returning ``None`` on any miss.""" + node = root + for attr in path: + node = getattr(node, attr, None) + if node is None: + return None + return node + + +def resolve_warp_twin(name: str, modules: Iterable[Any]) -> Any | None: + """Return ``name`` from ``modules`` if its ``__module__`` lives under a warp prefix.""" + for module in modules: + candidate = getattr(module, name, None) + if candidate is None: + continue + origin = getattr(candidate, "__module__", "") or "" + if origin.startswith(WARP_ROOT_PREFIXES): + return candidate + return None + + +def iter_term_attrs(group: Any) -> Iterator[tuple[str, Any]]: + """Yield ``(name, term)`` pairs from a manager-cfg group, skipping dunders.""" + if group is None: + return + for name in [n for n in dir(group) if not n.startswith("_")]: + term = getattr(group, name, None) + if term is None: + continue + yield name, term diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/torch.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/torch.py new file mode 100644 index 000000000000..1c4054e79cda --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/torch.py @@ -0,0 +1,72 @@ +# 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 + +"""Torch (default) frontend. + +The torch frontend is a thin wrapper around :func:`gym.make`. It does not +mutate the cfg — the standard ``isaaclab.envs.*`` env classes use +``FactoryBase`` to dispatch on the active physics backend at construction +time, so PhysX and Newton physics both run through this single path. + +Its rule set is small: one warning when the user asks for the torch +frontend on a task that's *registered* against the warp runtime +(``entry_point`` under :data:`WARP_ROOT_PREFIXES`). That's a contradiction +of intent — the env class is a warp env regardless of the frontend flag — +so we surface it but don't block. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +from .base import ( + Change, + CompatRule, + Frontend, + Issue, + ResolveContext, + Runtime, + Severity, + TaskMeta, +) + + +class WarnIfTaskIsWarpRegistered(CompatRule): + """Warn when a warp-registered task is asked to run on the torch frontend. + + The env class itself is a warp implementation; ``--frontend=torch`` will + still happily build it via ``gym.make`` (since the entry-point class is + what runs), but the user's stated intent says torch. We emit a warning + so the contradiction is visible. + """ + + name = "warn_if_task_is_warp_registered" + + def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: + if ctx.task.runtime == Runtime.WARP: + yield Issue( + rule=self.name, + severity=Severity.WARNING, + message=( + f"task is registered against the warp runtime" + f" (entry_point={ctx.task.entry_point!r}); --frontend=torch" + f" will still build the warp env class. Consider --frontend=warp." + ), + location="task.entry_point", + detail={"runtime": ctx.task.runtime.value}, + ) + + +class TorchFrontend(Frontend): + """Default frontend: ``gym.make`` against the registered ``entry_point``.""" + + name = "torch" + rules = (WarnIfTaskIsWarpRegistered,) + + def construct(self, cfg: Any, meta: TaskMeta, **kwargs: Any) -> Any: + import gymnasium as gym + + return gym.make(meta.task_id, cfg=cfg, **kwargs) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py new file mode 100644 index 000000000000..1ce0c5b21704 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py @@ -0,0 +1,453 @@ +# 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 + +"""Warp frontend. + +Routes a stable manager-based task cfg onto :class:`ManagerBasedRLEnvWarp` +by running a pluggable :class:`CompatRule` pipeline. Direct envs encode +their runtime in the env class itself (each warp direct env is its own +class), so for direct cfgs the frontend just verifies the registered +entry-point lives under :data:`WARP_ROOT_PREFIXES` and forwards to +:func:`gym.make`. + +Rule pipeline (manager-based path): + +* :class:`CheckPhysicsIsNewton` — blocking issue if ``cfg.sim.physics`` + is PhysX-flavoured, since the warp runtime needs Newton physics. +* :class:`ResolvePhysicsPreset` — collapses ``PresetCfg`` wrappers (no-op + when Hydra already did so). +* :class:`DropUnsupportedSensors` — removes sensors warp can't run yet. +* :class:`PromoteSceneEntityCfg` — in-place class promotion to the warp + variant so warp mdp kernels see ``joint_mask`` / ``joint_ids_wp``. +* :class:`SwapMdpFunctions` — name-based ``term.func`` replacement against + the warp ``mdp`` modules. +* :class:`SwapActionClassType` — strict swap of action ``class_type`` to + the warp twin. +* :class:`VerifyDirectIsWarp` — only fires for direct cfgs; blocks if + the registered entry-point isn't a warp class. +""" + +from __future__ import annotations + +import importlib +import logging +from collections.abc import Iterable +from types import ModuleType +from typing import Any, ClassVar + +from .base import ( + WARP_ROOT_PREFIXES, + Change, + CompatRule, + Frontend, + FrontendIncompatibleError, + Issue, + ResolveContext, + Runtime, + Severity, + TaskMeta, + Workflow, + iter_term_attrs, + resolve_warp_twin, + walk_attrs, +) + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Rules +# --------------------------------------------------------------------------- + + +class CheckPhysicsIsNewton(CompatRule): + """Block if ``cfg.sim.physics`` isn't Newton. + + The warp runtime resolves asset / sensor ``class_type`` via the active + physics backend's module tree (``isaaclab_newton.*`` for Newton, + ``isaaclab_physx.*`` for PhysX). Loading a PhysX class under the warp + runtime fails on ``omni.physics.tensors.api`` (a Kit module the warp + runtime doesn't initialise), so a PhysX physics cfg is fundamentally + incompatible with this frontend. + + Hydra's preset resolution + :meth:`WarpFrontend.preprocess_hydra_args` + normally collapse ``PresetCfg`` wrappers to the ``newton`` field + before this rule runs; we still check, so a programmatic caller or a + misconfigured task fails fast and loudly here instead of crashing + deep in scene init. + """ + + name = "check_physics_is_newton" + + def applies_to(self, ctx: ResolveContext) -> bool: + return ctx.task.workflow == Workflow.MANAGER_BASED + + def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: + physics = getattr(getattr(cfg, "sim", None), "physics", None) + if physics is None: + return + cls = type(physics) + module = getattr(cls, "__module__", "") or "" + # Skip PresetCfg wrappers — they'll be unwrapped by ResolvePhysicsPreset. + if hasattr(physics, "newton") and not hasattr(physics, "class_type"): + return + if module.startswith("isaaclab_physx"): + yield Issue( + rule=self.name, + severity=Severity.BLOCKING, + message=( + f"sim.physics is {cls.__name__} ({module}); the warp runtime" + f" needs Newton physics. Pass `presets=newton` or use a Newton" + f" physics cfg explicitly." + ), + location="cfg.sim.physics", + detail={"physics_class": cls.__name__, "physics_module": module}, + ) + + +class ResolvePhysicsPreset(CompatRule): + """Collapse ``cfg.sim.physics`` from ``PresetCfg`` to its ``newton`` field. + + Hydra's preset resolution normally collapses the wrapper before the + frontend runs. This rule covers the residual case of programmatic use + of :meth:`Frontend.build` without Hydra, and custom physics cfgs that + still expose a ``newton`` attribute. + """ + + name = "resolve_physics_preset" + + def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: + physics = getattr(getattr(cfg, "sim", None), "physics", None) + if physics is None or not hasattr(physics, "newton"): + return + cfg.sim.physics = physics.newton + yield Change( + rule=self.name, + description=f"sim.physics → {type(physics.newton).__name__}", + ) + + +class DropUnsupportedSensors(CompatRule): + """Drop scene sensors the warp runtime can't run yet. + + Default: ``("height_scanner",)``. Pass ``sensors=...`` at construction + time to override. + """ + + name = "drop_unsupported_sensors" + + def __init__(self, sensors: Iterable[str] = ("height_scanner",)): + self.sensors = tuple(sensors) + + def applies_to(self, ctx: ResolveContext) -> bool: + return ctx.task.workflow == Workflow.MANAGER_BASED + + def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: + scene = getattr(cfg, "scene", None) + if scene is None: + return + for sensor in self.sensors: + if getattr(scene, sensor, None) is not None: + setattr(scene, sensor, None) + yield Change(rule=self.name, description=f"scene.{sensor} → None") + + +class PromoteSceneEntityCfg(CompatRule): + """Promote :class:`SceneEntityCfg` instances under term params to the warp variant. + + The warp variant adds cached ``joint_mask`` / ``joint_ids_wp`` / + ``body_ids_wp`` fields that the warp mdp kernels read after + :meth:`resolve`. The class hierarchy is asserted at apply time — if + the warp class no longer subclasses the stable class, the rule raises + :class:`FrontendIncompatibleError` so the in-place ``__class__`` + promotion can never silently corrupt instances. + """ + + name = "promote_scene_entity_cfg" + + GROUPS: ClassVar[tuple[tuple[str, ...], ...]] = ( + ("observations", "policy"), + ("events",), + ("rewards",), + ("terminations",), + ("commands",), + ("curriculum",), + ("actions",), + ) + + def applies_to(self, ctx: ResolveContext) -> bool: + return ctx.task.workflow == Workflow.MANAGER_BASED + + def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: + from isaaclab.managers.scene_entity_cfg import SceneEntityCfg as Stable + + from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as Warp + + if not issubclass(Warp, Stable): + raise FrontendIncompatibleError( + f"Warp SceneEntityCfg must subclass stable SceneEntityCfg; got mro {[c.__name__ for c in Warp.__mro__]}" + ) + + promoted = 0 + for path in self.GROUPS: + group = walk_attrs(cfg, path) + if group is None: + continue + for _name, term in iter_term_attrs(group): + params = getattr(term, "params", None) + if not isinstance(params, dict): + continue + for value in params.values(): + if isinstance(value, Warp) or not isinstance(value, Stable): + continue + value.__class__ = Warp + value.joint_mask = None + value.joint_ids_wp = None + value.body_ids_wp = None + promoted += 1 + if promoted: + yield Change( + rule=self.name, + description=f"promoted {promoted} SceneEntityCfg instance(s) to warp variant", + ) + + +class SwapMdpFunctions(CompatRule): + """Replace ``term.func`` with the same-named callable from warp ``mdp`` modules. + + Walks observation / event / reward / termination / command / curriculum + groups. A re-export from stable code is rejected by checking + ``__module__``. Missing twins are dropped (term set to ``None``) and + surfaced as :attr:`Severity.WARNING`; in strict mode they're escalated + to :attr:`Severity.BLOCKING`. + """ + + name = "swap_mdp_functions" + + GROUPS: ClassVar[tuple[tuple[str, ...], ...]] = ( + ("observations", "policy"), + ("events",), + ("rewards",), + ("terminations",), + ("commands",), + ("curriculum",), + ) + + def __init__( + self, + stable_root: str = "isaaclab_tasks", + warp_root: str = "isaaclab_tasks_experimental", + fallback_mdp: str = "isaaclab_experimental.envs.mdp", + ): + self.stable_root = stable_root + self.warp_root = warp_root + self.fallback_mdp = fallback_mdp + + def applies_to(self, ctx: ResolveContext) -> bool: + return ctx.task.workflow == Workflow.MANAGER_BASED + + def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: + modules = self._mdp_modules(ctx.task) + logger.info("WarpFrontend: warp mdp modules → %s", [m.__name__ for m in modules]) + searched = tuple(m.__name__ for m in modules) + for path in self.GROUPS: + group = walk_attrs(cfg, path) + if group is None: + continue + location_prefix = ".".join(path) + for name, term in iter_term_attrs(group): + if not hasattr(term, "func"): + continue + stable = term.func + if not callable(stable): + continue + twin = resolve_warp_twin(stable.__name__, modules) + if twin is not None: + term.func = twin + continue + # No twin found: drop the term and report. + setattr(group, name, None) + yield Issue( + rule=self.name, + severity=Severity.BLOCKING if ctx.strict else Severity.WARNING, + message=(f"no warp twin for stable func={stable.__name__!r}; term dropped."), + location=f"{location_prefix}.{name}", + detail={"expected": stable.__name__, "searched": list(searched)}, + ) + + def _mdp_modules(self, task: TaskMeta) -> list[ModuleType]: + """Locate warp mdp modules that mirror the task's stable mdp.""" + modules: list[ModuleType] = [] + entry = task.env_cfg_entry_point + if isinstance(entry, str) and entry.startswith(self.stable_root): + warp_pkg = entry.rsplit(".", 1)[0].replace(self.stable_root, self.warp_root, 1) + parts = warp_pkg.split(".") + for depth in range(len(parts), 0, -1): + try: + modules.append(importlib.import_module(".".join(parts[:depth] + ["mdp"]))) + break + except ImportError: + continue + try: + modules.append(importlib.import_module(self.fallback_mdp)) + except ImportError: + logger.warning("WarpFrontend: fallback mdp module %r not importable", self.fallback_mdp) + return modules + + +class SwapActionClassType(CompatRule): + """Swap ``cfg.actions..class_type`` with the same-named warp class. + + Always strict at the term level: an action with no warp twin can't run + on the warp runtime, so a missing twin is :attr:`Severity.BLOCKING` + regardless of ``ctx.strict``. + """ + + name = "swap_action_class_type" + + def __init__(self, warp_actions_module: str = "isaaclab_experimental.envs.mdp.actions.joint_actions"): + self.warp_actions_module = warp_actions_module + + def applies_to(self, ctx: ResolveContext) -> bool: + return ctx.task.workflow == Workflow.MANAGER_BASED + + def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: + actions = getattr(cfg, "actions", None) + if actions is None: + return + try: + module = importlib.import_module(self.warp_actions_module) + except ImportError as exc: + yield Issue( + rule=self.name, + severity=Severity.BLOCKING, + message=(f"warp action module {self.warp_actions_module!r} not importable: {exc}"), + location="actions", + ) + return + searched = (module.__name__,) + for name, term in iter_term_attrs(actions): + if not hasattr(term, "class_type"): + continue + stable = term.class_type + if not callable(stable): + continue + twin = resolve_warp_twin(stable.__name__, [module]) + if twin is not None: + term.class_type = twin + continue + yield Issue( + rule=self.name, + severity=Severity.BLOCKING, + message=f"no warp twin for stable class_type={stable.__name__!r}", + location=f"actions.{name}", + detail={"expected": stable.__name__, "searched": list(searched)}, + ) + + +class VerifyDirectIsWarp(CompatRule): + """For direct envs, block if the registered entry-point isn't a warp class. + + Direct envs encode their runtime in the env class itself (each warp + direct env is a separate class with its own kernels), so the warp + frontend can't *adapt* a stable direct cfg — it can only refuse it + with a clear message. A direct task whose entry-point already lives + under :data:`WARP_ROOT_PREFIXES` is a pass-through. + """ + + name = "verify_direct_is_warp" + + def applies_to(self, ctx: ResolveContext) -> bool: + return ctx.task.workflow == Workflow.DIRECT + + def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: + if ctx.task.runtime == Runtime.WARP: + return + ep = ctx.task.entry_point + yield Issue( + rule=self.name, + severity=Severity.BLOCKING, + message=( + f"direct env entry_point {ep!r} is not a warp implementation." + f" --frontend=warp on a direct task requires the registered class to live under" + f" {list(WARP_ROOT_PREFIXES)} (e.g. *-Direct-Warp-v0)." + ), + location="task.entry_point", + detail={"entry_point": ep}, + ) + + +# --------------------------------------------------------------------------- +# WarpFrontend +# --------------------------------------------------------------------------- + + +class WarpFrontend(Frontend): + """Run a stable env cfg on the warp runtime. + + Manager-based cfgs get the full rule pipeline and are constructed on + :class:`ManagerBasedRLEnvWarp`. Direct cfgs run :class:`VerifyDirectIsWarp` + only and dispatch through :func:`gym.make`. + """ + + name = "warp" + rules = ( + CheckPhysicsIsNewton, + ResolvePhysicsPreset, + DropUnsupportedSensors, + PromoteSceneEntityCfg, + SwapMdpFunctions, + SwapActionClassType, + VerifyDirectIsWarp, + ) + + # -- construction --------------------------------------------------------- + + def construct(self, cfg: Any, meta: TaskMeta, **kwargs: Any) -> Any: + if meta.workflow == Workflow.MANAGER_BASED: + from isaaclab_experimental.envs import ManagerBasedRLEnvWarp + + return ManagerBasedRLEnvWarp(cfg=cfg, **kwargs) + if meta.workflow == Workflow.DIRECT: + import gymnasium as gym + + return gym.make(meta.task_id, cfg=cfg, **kwargs) + raise FrontendIncompatibleError( + f"WarpFrontend does not support workflow {meta.workflow.value!r} (task {meta.task_id!r})." + ) + + # -- CLI hook ------------------------------------------------------------- + + def preprocess_hydra_args(self, task_id: str, args: list[str]) -> list[str]: + """Inject ``presets=newton`` for stable manager-based tasks. + + Hydra resolves ``presets=...`` *before* the frontend runs, so the + injection has to happen at the CLI layer. We only inject for tasks + whose ``env_cfg_entry_point`` is under ``isaaclab_tasks.manager_based``; + other tasks (direct, pre-warp ``*-Warp-v0``) don't carry a preset + system and would error out. + + If the user already passed an explicit ``presets=`` we don't override, + but we warn when their preset isn't ``newton``. + """ + import gymnasium as gym + + try: + spec = gym.spec(task_id) + except gym.error.NameNotFound: + return args + cfg_ep = spec.kwargs.get("env_cfg_entry_point") + if not isinstance(cfg_ep, str) or not cfg_ep.startswith("isaaclab_tasks.manager_based"): + return args + explicit = next((a for a in args if a.startswith("presets=")), None) + if explicit is None: + return [*args, "presets=newton"] + if explicit != "presets=newton": + logger.warning( + "--frontend=warp on %r expects presets=newton; got %r — adapter may fail.", + task_id, + explicit, + ) + return args diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py deleted file mode 100644 index 23ea519f7cea..000000000000 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/warp_frontend.py +++ /dev/null @@ -1,557 +0,0 @@ -# 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 - -"""Run a stable env cfg on the warp runtime via a pluggable rule pipeline. - -The frontend classifies the cfg as manager-based or direct, runs the -applicable :class:`CompatRule` instances against it, and constructs the -matching warp env class. New incompatibilities (e.g. a new sensor type warp -can't run, a new term-cfg field that needs upgrading) are added by writing a -small :class:`CompatRule` subclass and passing it through the constructor — -no surgery in the dispatcher. - -Module-level imports are deliberately limited to lightweight stdlib modules. -Warp library imports fire only inside :meth:`WarpFrontend.build`, which must -be called *after* :class:`SimulationApp` is alive, otherwise warp's lib load -races with Kit's USD/pxr extension initialisation. -""" - -from __future__ import annotations - -import importlib -import logging -from abc import ABC, abstractmethod -from collections.abc import Iterable, Sequence -from dataclasses import dataclass, field -from enum import Enum -from types import ModuleType -from typing import Any, ClassVar - -logger = logging.getLogger(__name__) - - -# --------------------------------------------------------------------------- -# Compatibility-rule framework -# --------------------------------------------------------------------------- - - -class CfgKind(str, Enum): - """Workflow type of an env cfg the frontend is asked to adapt.""" - - MANAGER_BASED = "manager_based" - DIRECT = "direct" - - -_BOTH = frozenset({CfgKind.MANAGER_BASED, CfgKind.DIRECT}) -_MANAGER_ONLY = frozenset({CfgKind.MANAGER_BASED}) - -_WARP_MODULE_PREFIXES: tuple[str, ...] = ("isaaclab_experimental", "isaaclab_tasks_experimental") -"""Module-path prefixes that distinguish a real warp twin from a stable -re-export. The warp ``mdp`` packages do ``from isaaclab.envs.mdp import *``, -so a plain ``getattr`` lookup can return the stable symbol untouched.""" - -_TERM_GROUPS_FOR_PROMOTION: tuple[tuple[str, ...], ...] = ( - ("observations", "policy"), - ("events",), - ("rewards",), - ("terminations",), - ("commands",), - ("curriculum",), - ("actions",), -) -"""Cfg groups walked when promoting :class:`SceneEntityCfg` instances. Includes -``actions`` because action params often carry a :class:`SceneEntityCfg`.""" - - -class IncompatibleEnvError(RuntimeError): - """Raised when the frontend can't run on the given cfg / task pair.""" - - def __init__(self, message: str, report: CompatReport | None = None): - super().__init__(message) - self.report = report - - -@dataclass -class CompatIssue: - """One thing the adapter looked for but couldn't resolve to a warp twin.""" - - rule: str - """Rule that produced this issue (matches :attr:`CompatRule.name`).""" - location: str - """Cfg path the issue was found at, e.g. ``rewards`` / ``actions``.""" - item: str - """Attribute name on the group, e.g. ``track_lin_vel_xy_exp``.""" - expected: str - """The symbol the rule searched for (the stable ``__name__``).""" - searched: tuple[str, ...] - """Module dotted-paths inspected, in order.""" - action: str - """Outcome: ``"dropped"``, ``"raised"``, or ``"left-stable"``.""" - - -@dataclass -class CompatReport: - """Outcome of running the rule pipeline against a cfg. - - ``changes`` is a human-readable list of mutations applied; ``issues`` - records every term the adapter couldn't resolve. A clean run has both - empty. - """ - - rules_applied: list[str] = field(default_factory=list) - changes: list[str] = field(default_factory=list) - issues: list[CompatIssue] = field(default_factory=list) - - def has_issues(self) -> bool: - return bool(self.issues) - - def has_fatal(self) -> bool: - return any(i.action == "raised" for i in self.issues) - - def __bool__(self) -> bool: - return bool(self.issues) or bool(self.changes) - - def format(self) -> str: - lines: list[str] = [] - if self.changes: - lines.append(f"WarpFrontend applied {len(self.changes)} change(s):") - lines.extend(f" - {c}" for c in self.changes) - if self.issues: - lines.append(f"WarpFrontend has {len(self.issues)} unresolved item(s):") - for i in self.issues: - lines.append( - f" - [{i.rule}] {i.location}.{i.item}: expected {i.expected!r}" - f" (searched {list(i.searched)}) → {i.action}" - ) - return "\n".join(lines) if lines else "WarpFrontend: no changes." - - -@dataclass -class AdaptContext: - """Per-call context handed to each rule.""" - - task_id: str - kind: CfgKind - strict: bool - - -class CompatRule(ABC): - """One step in the warp adaptation pipeline. - - A rule mutates ``cfg`` in place and appends to ``report``. To handle a - new incompatibility, subclass and implement :meth:`apply`. Rules declare - which workflow they apply to via :attr:`applies_to`; the pipeline skips - any rule whose ``applies_to`` doesn't include the cfg kind. - """ - - name: ClassVar[str] - applies_to: ClassVar[frozenset[CfgKind]] - - @abstractmethod - def apply(self, cfg: Any, ctx: AdaptContext, report: CompatReport) -> None: - """Mutate ``cfg`` in place and append issues / changes to ``report``.""" - - -# --------------------------------------------------------------------------- -# Helpers shared by rules -# --------------------------------------------------------------------------- - - -def _walk(root: Any, path: tuple[str, ...]) -> Any: - node = root - for attr in path: - node = getattr(node, attr, None) - if node is None: - return None - return node - - -def _resolve_warp_twin(name: str, modules: Sequence[ModuleType]) -> Any | None: - """Return ``name`` from ``modules`` if it lives under a warp prefix, else None.""" - for module in modules: - candidate = getattr(module, name, None) - if candidate is None: - continue - origin = getattr(candidate, "__module__", "") or "" - if origin.startswith(_WARP_MODULE_PREFIXES): - return candidate - return None - - -def _swap_named_attr( - group: Any, - location: str, - attr: str, - modules: Sequence[ModuleType], - *, - rule: str, - strict: bool, - report: CompatReport, -) -> None: - """Replace ``term.`` with the same-named callable from ``modules``. - - For each public attribute on ``group`` whose value has a callable ``attr`` - field, find a same-named symbol in ``modules`` whose ``__module__`` lives - under :data:`_WARP_MODULE_PREFIXES` and assign it. If no twin is found, - drop the term (set to ``None``) when ``strict`` is False, raise - :class:`LookupError` when True. Either way append a :class:`CompatIssue`. - """ - if group is None: - return - searched = tuple(m.__name__ for m in modules) - # Snapshot the attribute names so mutating ``group`` mid-loop is safe. - for name in [n for n in dir(group) if not n.startswith("_")]: - term = getattr(group, name, None) - if term is None or not hasattr(term, attr): - continue - stable = getattr(term, attr) - if not callable(stable): - continue - twin = _resolve_warp_twin(stable.__name__, modules) - if twin is not None: - setattr(term, attr, twin) - continue - issue = CompatIssue( - rule=rule, - location=location, - item=name, - expected=stable.__name__, - searched=searched, - action="raised" if strict else "dropped", - ) - report.issues.append(issue) - if strict: - raise LookupError( - f"WarpFrontend: cannot adapt {location}.{name} — no warp twin for " - f"stable {attr}={stable.__name__!r}. Searched {list(searched)}." - ) - setattr(group, name, None) - - -# --------------------------------------------------------------------------- -# Concrete rules -# --------------------------------------------------------------------------- - - -class ResolvePhysicsPresetRule(CompatRule): - """Collapse ``cfg.sim.physics`` from a :class:`PresetCfg` to its ``newton`` field. - - Hydra's preset resolution normally collapses the wrapper before the - frontend runs. This rule covers two residual cases: programmatic use of - :meth:`WarpFrontend.adapt` (no Hydra), and custom physics cfgs that still - expose a ``newton`` attribute. - """ - - name = "resolve_physics_preset" - applies_to = _BOTH - - def apply(self, cfg, ctx, report): - physics = getattr(getattr(cfg, "sim", None), "physics", None) - if physics is not None and hasattr(physics, "newton"): - cfg.sim.physics = physics.newton - report.changes.append(f"sim.physics → {type(physics.newton).__name__}") - - -class DropUnsupportedSensorsRule(CompatRule): - """Drop scene sensors the warp runtime can't run yet.""" - - name = "drop_unsupported_sensors" - applies_to = _BOTH - - def __init__(self, sensors: Iterable[str] = ("height_scanner",)): - self.sensors = tuple(sensors) - - def apply(self, cfg, ctx, report): - scene = getattr(cfg, "scene", None) - if scene is None: - return - for sensor in self.sensors: - if getattr(scene, sensor, None) is not None: - setattr(scene, sensor, None) - report.changes.append(f"scene.{sensor} → None") - - -class PromoteSceneEntityCfgRule(CompatRule): - """Promote :class:`SceneEntityCfg` instances under term params to the warp variant. - - The warp variant adds cached ``joint_mask`` / ``joint_ids_wp`` / - ``body_ids_wp`` fields the warp mdp kernels read after :meth:`resolve`. - The class hierarchy is asserted at apply time — if the warp class no - longer subclasses the stable class (refactor or divergence), the rule - raises :class:`TypeError` so the assumption can never silently corrupt - instances. - """ - - name = "promote_scene_entity_cfg" - applies_to = _MANAGER_ONLY - - def apply(self, cfg, ctx, report): - from isaaclab.managers.scene_entity_cfg import SceneEntityCfg as _StableSE - - from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as _WarpSE - - if not issubclass(_WarpSE, _StableSE): - raise TypeError( - "PromoteSceneEntityCfgRule requires the warp SceneEntityCfg to subclass the stable one;" - f" got mro {[c.__name__ for c in _WarpSE.__mro__]}" - ) - - promoted = 0 - for path in _TERM_GROUPS_FOR_PROMOTION: - group = _walk(cfg, path) - if group is None: - continue - for name in (n for n in dir(group) if not n.startswith("_")): - term = getattr(group, name, None) - if term is None: - continue - params = getattr(term, "params", None) - if not isinstance(params, dict): - continue - for value in params.values(): - if isinstance(value, _WarpSE) or not isinstance(value, _StableSE): - continue - value.__class__ = _WarpSE - value.joint_mask = None - value.joint_ids_wp = None - value.body_ids_wp = None - promoted += 1 - if promoted: - report.changes.append(f"promoted {promoted} SceneEntityCfg instance(s) to warp variant") - - -class SwapMdpFunctionsRule(CompatRule): - """Replace ``term.func`` with the same-named callable from warp ``mdp`` modules. - - Walks observation / event / reward / termination / command / curriculum - groups. A re-export from stable code is rejected by checking - ``__module__``. Missing twins are dropped (term set to ``None``) unless - the context is strict, in which case they raise. - """ - - name = "swap_mdp_functions" - applies_to = _MANAGER_ONLY - - GROUPS: ClassVar[tuple[tuple[str, ...], ...]] = ( - ("observations", "policy"), - ("events",), - ("rewards",), - ("terminations",), - ("commands",), - ("curriculum",), - ) - - def __init__( - self, - stable_root: str = "isaaclab_tasks", - warp_root: str = "isaaclab_tasks_experimental", - fallback_mdp: str = "isaaclab_experimental.envs.mdp", - ): - self.stable_root = stable_root - self.warp_root = warp_root - self.fallback_mdp = fallback_mdp - - def apply(self, cfg, ctx, report): - modules = self._mdp_modules_for(ctx.task_id) - logger.info("WarpFrontend: warp mdp modules → %s", [m.__name__ for m in modules]) - for path in self.GROUPS: - group = _walk(cfg, path) - _swap_named_attr( - group, - location=".".join(path), - attr="func", - modules=modules, - rule=self.name, - strict=ctx.strict, - report=report, - ) - - def _mdp_modules_for(self, task_id: str) -> list[ModuleType]: - """Locate warp mdp modules that mirror the stable task's mdp. - - Convention: replace :attr:`stable_root` with :attr:`warp_root` in the - stable cfg's module path, walk up looking for an ``mdp`` sub-module on - the warp side, and always append :attr:`fallback_mdp` for cross-task - base terms. - """ - modules: list[ModuleType] = [] - try: - import gymnasium as gym - - entry = gym.spec(task_id).kwargs.get("env_cfg_entry_point") - except (KeyError, AttributeError, ModuleNotFoundError, gym.error.NameNotFound): - # Spec is missing / kwargs missing the entry — fall back to the - # package-wide warp mdp module only. Real ImportError from a - # broken cfg module is *not* swallowed; it propagates. - entry = None - if isinstance(entry, str) and entry.startswith(self.stable_root): - warp_pkg = entry.rsplit(".", 1)[0].replace(self.stable_root, self.warp_root, 1) - parts = warp_pkg.split(".") - for depth in range(len(parts), 0, -1): - try: - modules.append(importlib.import_module(".".join(parts[:depth] + ["mdp"]))) - break - except ImportError: - continue - try: - modules.append(importlib.import_module(self.fallback_mdp)) - except ImportError: - logger.warning("WarpFrontend: fallback mdp module %r not importable", self.fallback_mdp) - return modules - - -class SwapActionClassTypeRule(CompatRule): - """Swap ``cfg.actions..class_type`` with the same-named warp class. - - Always strict: an action with no warp class can't run on the warp - runtime, so a missing twin raises :class:`LookupError`. - """ - - name = "swap_action_class_type" - applies_to = _MANAGER_ONLY - - def __init__(self, warp_actions_module: str = "isaaclab_experimental.envs.mdp.actions.joint_actions"): - self.warp_actions_module = warp_actions_module - - def apply(self, cfg, ctx, report): - actions = getattr(cfg, "actions", None) - if actions is None: - return - try: - module = importlib.import_module(self.warp_actions_module) - except ImportError as exc: - raise IncompatibleEnvError( - f"WarpFrontend: warp action module {self.warp_actions_module!r} not importable" - ) from exc - _swap_named_attr( - actions, - location="actions", - attr="class_type", - modules=[module], - rule=self.name, - strict=True, - report=report, - ) - - -# --------------------------------------------------------------------------- -# Frontend -# --------------------------------------------------------------------------- - - -class WarpFrontend: - """Run a stable env cfg on the warp runtime via a pluggable rule pipeline. - - The frontend classifies the input cfg as manager-based or direct, runs - the applicable :class:`CompatRule` instances, then constructs the - matching warp env class. - - Direct envs encode their runtime in the env class itself (each warp - direct env is a separate class with its own kernels), so the frontend - doesn't transform direct cfgs — it only verifies the registered - ``entry_point`` lives under :data:`_WARP_MODULE_PREFIXES` and dispatches - via :func:`gym.make`. A stable direct cfg + ``--frontend=warp`` raises - :class:`IncompatibleEnvError` with a clear message. - - Example:: - - with launch_simulation(cfg, args_cli): - env = WarpFrontend().build(cfg, task_id=args_cli.task) - if env.unwrapped.warp_compat_report.has_issues(): - print(env.unwrapped.warp_compat_report.format()) - """ - - DEFAULT_RULES: ClassVar[tuple[type[CompatRule], ...]] = ( - ResolvePhysicsPresetRule, - DropUnsupportedSensorsRule, - PromoteSceneEntityCfgRule, - SwapMdpFunctionsRule, - SwapActionClassTypeRule, - ) - - def __init__(self, rules: Sequence[CompatRule] | None = None, strict: bool = False): - self.rules: list[CompatRule] = list(rules) if rules is not None else [r() for r in self.DEFAULT_RULES] - self.strict = strict - - # -- public --------------------------------------------------------------- - - def adapt(self, cfg: Any, task_id: str) -> CompatReport: - """Run the rule pipeline; mutate ``cfg`` in place; return a report. - - The report is logged at ``WARNING`` (or ``ERROR`` in strict mode when - any issue is fatal). Missing twins raise :class:`LookupError` only - when ``strict=True``. - """ - kind = self._classify(cfg) - ctx = AdaptContext(task_id=task_id, kind=kind, strict=self.strict) - report = CompatReport() - for rule in self.rules: - if kind not in rule.applies_to: - continue - rule.apply(cfg, ctx, report) - report.rules_applied.append(rule.name) - if report: - level = logging.ERROR if self.strict and report.has_fatal() else logging.WARNING - logger.log(level, report.format()) - return report - - def build(self, cfg: Any, task_id: str, render_mode: str | None = None): - """Adapt ``cfg`` and construct a warp env. - - Returns an env instance with ``warp_compat_report`` attached on the - unwrapped env. ``render_mode`` is forwarded so ``--video`` keeps - working when the frontend is selected at the train-script level. - """ - # Lazy: this is the first warp-lib load. The caller must already be - # inside the SimulationApp context (i.e. inside ``launch_simulation``). - from isaaclab.envs import DirectRLEnvCfg, ManagerBasedRLEnvCfg - - report = self.adapt(cfg, task_id) - if isinstance(cfg, ManagerBasedRLEnvCfg): - from isaaclab_experimental.envs import ManagerBasedRLEnvWarp - - env = ManagerBasedRLEnvWarp(cfg=cfg, render_mode=render_mode) - elif isinstance(cfg, DirectRLEnvCfg): - self._verify_direct_warp(task_id) - import gymnasium as gym - - env = gym.make(task_id, cfg=cfg, render_mode=render_mode) - else: - raise IncompatibleEnvError( - f"WarpFrontend supports ManagerBasedRLEnvCfg or DirectRLEnvCfg, got {type(cfg).__name__}", - report=report, - ) - env.unwrapped.warp_compat_report = report - return env - - # -- helpers -------------------------------------------------------------- - - @staticmethod - def _classify(cfg: Any) -> CfgKind: - from isaaclab.envs import DirectRLEnvCfg, ManagerBasedRLEnvCfg - - if isinstance(cfg, ManagerBasedRLEnvCfg): - return CfgKind.MANAGER_BASED - if isinstance(cfg, DirectRLEnvCfg): - return CfgKind.DIRECT - raise IncompatibleEnvError( - f"WarpFrontend supports ManagerBasedRLEnvCfg or DirectRLEnvCfg, got {type(cfg).__name__}" - ) - - @staticmethod - def _verify_direct_warp(task_id: str) -> None: - import gymnasium as gym - - try: - spec = gym.spec(task_id) - except gym.error.NameNotFound as exc: - raise IncompatibleEnvError(f"Task {task_id!r} is not registered with gymnasium") from exc - ep = spec.entry_point - if isinstance(ep, str) and not ep.startswith(_WARP_MODULE_PREFIXES): - raise IncompatibleEnvError( - f"Direct env {task_id!r} entry_point {ep!r} is not a warp implementation." - f" --frontend=warp on a direct task requires the registered class to live under" - f" {list(_WARP_MODULE_PREFIXES)} (e.g. *-Direct-Warp-v0)." - ) From 8934fe649a607bb066aab39abe7787a4264cc775 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 6 May 2026 07:52:27 +0000 Subject: [PATCH 08/58] Tighten frontend framework against review findings - SwapMdpFunctions: skip terms whose ``func.__module__`` is already under the warp prefixes. Without this, running the bridge against a task already registered under ``isaaclab_tasks_experimental`` (e.g. ``Isaac-Cartpole-Warp-v0 --frontend=warp``) would silently drop terms whose warp twin happens not to live in the resolved fallback module. Also tighten ``_mdp_modules`` to require the trailing dot when matching ``isaaclab_tasks`` so we don't double-replace the prefix and end up importing ``isaaclab_tasks_experimental_experimental.*``. - ResolvePhysicsPreset: scope to MANAGER_BASED via ``applies_to``. Direct cfgs aren't expected to carry ``PresetCfg`` wrappers; running this rule on them was a no-op but the scoping makes the contract explicit. - CheckPhysicsIsNewton: positively accept ``isaaclab_newton.*`` modules, block on ``isaaclab_physx.*``, warn on anything else. The previous rule only rejected PhysX, so a custom or third-party physics cfg in an unrelated module would slip through silently. - PromoteSceneEntityCfg: catch ``TypeError`` from the in-place ``__class__`` reassignment and surface it as a blocking issue. ``issubclass`` does not guarantee Python permits the layout change (slots, layout flags); failing loud at this seam is better than a cryptic crash mid-pipeline. - WarpFrontend.preprocess_hydra_args: normalise leading dashes when inspecting ``presets=``, so ``--presets=foo`` is treated the same as ``presets=foo`` (Hydra accepts both forms). - Frontend.resolve: hard-block when ``gym.spec`` returned no spec. Previously ``meta.runtime`` was ``UNKNOWN`` and ``construct`` would fail later with a less specific error; now the block fires before any rule runs. - train.py: import the frontend lazily and tolerate ``ImportError`` when ``--frontend=torch`` (the default). The experimental package is optional, so a missing install used to break the default path; now it falls back to ``gym.make`` for torch and only fails for ``warp``. - frontend/__init__.py: trim ``__all__`` and the wildcard re-export so helpers (``walk_attrs``, ``iter_term_attrs``, ``resolve_warp_twin``, ``WARP_ROOT_PREFIXES``) are no longer advertised as the public framework surface. They remain importable from ``frontend.base`` for users writing their own rules. --- .../reinforcement_learning/rsl_rl/train.py | 33 ++++++-- .../envs/frontend/__init__.py | 13 +-- .../envs/frontend/base.py | 16 ++++ .../envs/frontend/warp.py | 83 +++++++++++++++---- 4 files changed, 114 insertions(+), 31 deletions(-) diff --git a/scripts/reinforcement_learning/rsl_rl/train.py b/scripts/reinforcement_learning/rsl_rl/train.py index e146a0cc2575..e3801a08f88b 100644 --- a/scripts/reinforcement_learning/rsl_rl/train.py +++ b/scripts/reinforcement_learning/rsl_rl/train.py @@ -104,10 +104,26 @@ # may inject preset selections (e.g. WarpFrontend appends ``presets=newton`` # for stable manager-based tasks so PresetCfg wrappers resolve to the Newton # field before Hydra builds the cfg). -from isaaclab_experimental.envs.frontend import get_frontend # noqa: E402 - -_frontend = get_frontend(args_cli.frontend) -if args_cli.task is not None: +# +# The frontend package lives under ``isaaclab_experimental``, which is +# optional. If it isn't installed, ``--frontend=torch`` (the default) still +# needs to work — fall back to plain ``gym.make`` in that case. The warp +# selection requires the experimental package, so import failure there is +# fatal. +_frontend = None +try: + from isaaclab_experimental.envs.frontend import get_frontend # noqa: E402 + + _frontend = get_frontend(args_cli.frontend) +except ImportError as _frontend_import_err: + if args_cli.frontend != "torch": + raise SystemExit( + f"--frontend={args_cli.frontend} requires isaaclab_experimental to be installed," + f" but the import failed: {_frontend_import_err}" + ) from _frontend_import_err + logger.info("isaaclab_experimental not available; --frontend=torch will dispatch via plain gym.make.") + +if _frontend is not None and args_cli.task is not None: remaining_args = _frontend.preprocess_hydra_args(args_cli.task, remaining_args) sys.argv = [sys.argv[0]] + remaining_args @@ -198,9 +214,14 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen # The frontend may mutate env_cfg in place (preset resolution, # SceneEntityCfg promotion, mdp twin swaps, ...) before constructing # the env. ``env.unwrapped.frontend_report`` carries what changed - # and what was missing. + # and what was missing. When the experimental frontend package isn't + # available we already validated --frontend=torch above, so a plain + # gym.make is the documented fallback. render_mode = "rgb_array" if args_cli.video else None - env = _frontend.build(env_cfg, task_id=args_cli.task, render_mode=render_mode) + if _frontend is not None: + env = _frontend.build(env_cfg, task_id=args_cli.task, render_mode=render_mode) + else: + env = gym.make(args_cli.task, cfg=env_cfg, render_mode=render_mode) # convert to single-agent instance if required by the RL algorithm if isinstance(env.unwrapped.cfg, DirectMARLEnvCfg): diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/__init__.py index ad81605e6442..0f6e98b33974 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/__init__.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/__init__.py @@ -24,7 +24,6 @@ from __future__ import annotations from .base import ( - WARP_ROOT_PREFIXES, Change, CompatRule, Frontend, @@ -39,10 +38,7 @@ Workflow, available_frontends, get_frontend, - iter_term_attrs, register_frontend, - resolve_warp_twin, - walk_attrs, ) from .torch import TorchFrontend, WarnIfTaskIsWarpRegistered from .warp import ( @@ -62,6 +58,10 @@ register_frontend(WarpFrontend.name, WarpFrontend) +# Public API. Internal helpers (``walk_attrs``, ``iter_term_attrs``, +# ``resolve_warp_twin``, ``WARP_ROOT_PREFIXES``) are importable from +# :mod:`isaaclab_experimental.envs.frontend.base` for users writing their own +# rules, but they are not advertised here as the supported framework surface. __all__ = [ # core abstractions "Frontend", @@ -76,15 +76,10 @@ "Runtime", "ResolveContext", "FrontendIncompatibleError", - "WARP_ROOT_PREFIXES", # registry "register_frontend", "get_frontend", "available_frontends", - # rule helpers - "walk_attrs", - "resolve_warp_twin", - "iter_term_attrs", # built-in frontends "TorchFrontend", "WarpFrontend", diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py index 81fc21a33cf7..dece605f3aaa 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py @@ -366,6 +366,22 @@ def resolve(self, cfg: Any, task_id: str) -> Report: meta = TaskResolver.resolve(task_id, cfg) ctx = ResolveContext(frontend=self.name, task=meta, strict=self.strict) report = Report(frontend=self.name, task=meta) + # If gym couldn't resolve the spec, the frontend has nothing to dispatch + # against. Block early with a clear message rather than letting + # downstream construct() fail with a NameNotFound. + if meta.spec is None: + report.issues.append( + Issue( + rule="task_resolver", + severity=Severity.BLOCKING, + message=( + f"task {task_id!r} is not registered with gymnasium." + " Make sure the task package is imported before the frontend runs." + ), + location="task.spec", + ) + ) + return report self.preprocess_cfg(cfg, ctx) for rule in self._rules: if not rule.applies_to(ctx): diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py index 1ce0c5b21704..9dbd3652400b 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py @@ -81,6 +81,12 @@ class CheckPhysicsIsNewton(CompatRule): name = "check_physics_is_newton" + NEWTON_MODULE_PREFIXES: ClassVar[tuple[str, ...]] = ("isaaclab_newton",) + """Module prefixes that mark a physics cfg as Newton-flavoured.""" + + PHYSX_MODULE_PREFIXES: ClassVar[tuple[str, ...]] = ("isaaclab_physx",) + """Module prefixes that mark a physics cfg as PhysX-flavoured (definitely incompatible).""" + def applies_to(self, ctx: ResolveContext) -> bool: return ctx.task.workflow == Workflow.MANAGER_BASED @@ -93,18 +99,21 @@ def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: # Skip PresetCfg wrappers — they'll be unwrapped by ResolvePhysicsPreset. if hasattr(physics, "newton") and not hasattr(physics, "class_type"): return - if module.startswith("isaaclab_physx"): - yield Issue( - rule=self.name, - severity=Severity.BLOCKING, - message=( - f"sim.physics is {cls.__name__} ({module}); the warp runtime" - f" needs Newton physics. Pass `presets=newton` or use a Newton" - f" physics cfg explicitly." - ), - location="cfg.sim.physics", - detail={"physics_class": cls.__name__, "physics_module": module}, - ) + if module.startswith(self.NEWTON_MODULE_PREFIXES): + return # OK + severity = Severity.BLOCKING if module.startswith(self.PHYSX_MODULE_PREFIXES) else Severity.WARNING + backend = "PhysX" if severity == Severity.BLOCKING else "unknown" + yield Issue( + rule=self.name, + severity=severity, + message=( + f"sim.physics is {cls.__name__} ({module}); the warp runtime" + f" needs a Newton physics cfg ({backend} detected). Pass" + f" `presets=newton` or use a Newton physics cfg explicitly." + ), + location="cfg.sim.physics", + detail={"physics_class": cls.__name__, "physics_module": module}, + ) class ResolvePhysicsPreset(CompatRule): @@ -118,6 +127,9 @@ class ResolvePhysicsPreset(CompatRule): name = "resolve_physics_preset" + def applies_to(self, ctx: ResolveContext) -> bool: + return ctx.task.workflow == Workflow.MANAGER_BASED + def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: physics = getattr(getattr(cfg, "sim", None), "physics", None) if physics is None or not hasattr(physics, "newton"): @@ -202,7 +214,28 @@ def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: for value in params.values(): if isinstance(value, Warp) or not isinstance(value, Stable): continue - value.__class__ = Warp + # ``__class__ =`` is permitted only when the layouts + # match (e.g. neither side adds ``__slots__`` over the + # other). ``issubclass`` does *not* guarantee this. If + # the runtime refuses, surface a blocking issue rather + # than crash mid-pipeline; the cfg is still safe because + # we set warp-only fields *after* the assignment. + try: + value.__class__ = Warp + except TypeError as exc: + yield Issue( + rule=self.name, + severity=Severity.BLOCKING, + message=( + f"in-place __class__ promotion to {Warp.__name__} failed:" + f" {exc}. The warp variant likely diverged in slots/layout" + f" from the stable one; rebuild cfg explicitly with the warp" + f" SceneEntityCfg." + ), + location=".".join(path), + detail={"warp_class": Warp.__name__, "stable_class": Stable.__name__}, + ) + return value.joint_mask = None value.joint_ids_wp = None value.body_ids_wp = None @@ -263,6 +296,14 @@ def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: stable = term.func if not callable(stable): continue + # Idempotency: if the term already references a warp-native + # callable (e.g. running the bridge against an already-warp + # task), don't try to swap. Without this guard the rule would + # ``resolve_warp_twin`` against modules that might not contain + # the same symbol and silently drop the term. + origin = getattr(stable, "__module__", "") or "" + if origin.startswith(WARP_ROOT_PREFIXES): + continue twin = resolve_warp_twin(stable.__name__, modules) if twin is not None: term.func = twin @@ -281,7 +322,12 @@ def _mdp_modules(self, task: TaskMeta) -> list[ModuleType]: """Locate warp mdp modules that mirror the task's stable mdp.""" modules: list[ModuleType] = [] entry = task.env_cfg_entry_point - if isinstance(entry, str) and entry.startswith(self.stable_root): + # Use a trailing dot when matching the stable root so a task registered + # under ``isaaclab_tasks_experimental.*`` (the warp side) doesn't match + # ``isaaclab_tasks`` and end up with a double-replaced + # ``isaaclab_tasks_experimental_experimental.*`` import path. + stable_prefix = self.stable_root + "." + if isinstance(entry, str) and entry.startswith(stable_prefix): warp_pkg = entry.rsplit(".", 1)[0].replace(self.stable_root, self.warp_root, 1) parts = warp_pkg.split(".") for depth in range(len(parts), 0, -1): @@ -441,10 +487,15 @@ def preprocess_hydra_args(self, task_id: str, args: list[str]) -> list[str]: cfg_ep = spec.kwargs.get("env_cfg_entry_point") if not isinstance(cfg_ep, str) or not cfg_ep.startswith("isaaclab_tasks.manager_based"): return args - explicit = next((a for a in args if a.startswith("presets=")), None) + + # Match both ``presets=foo`` and ``--presets=foo`` — Hydra accepts either. + def _is_preset_arg(arg: str) -> bool: + return arg.lstrip("-").startswith("presets=") + + explicit = next((a for a in args if _is_preset_arg(a)), None) if explicit is None: return [*args, "presets=newton"] - if explicit != "presets=newton": + if explicit.lstrip("-") != "presets=newton": logger.warning( "--frontend=warp on %r expects presets=newton; got %r — adapter may fail.", task_id, From 67050a863d5e295a730ca5a49edac3ffc909abe7 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 6 May 2026 07:58:25 +0000 Subject: [PATCH 09/58] Tighten frontend framework against second review pass - TaskResolver._classify_runtime: also accept class/callable entry points by inspecting __module__. gym.register accepts both ``"module:Class"`` strings and class objects; the old check only handled strings, so a warp-registered task using the class form classified as Runtime.UNKNOWN and the warn / verify rules silently disengaged. - SwapMdpFunctions._mdp_modules: narrow the exception match from ImportError to ModuleNotFoundError where ``exc.name`` matches the module being looked up. Previously a real ImportError raised inside an existing mdp module (broken import inside the package) would be silently swallowed and the rule would fall through to the fallback module, producing misleading "no warp twin" reports. --- .../isaaclab_experimental/envs/frontend/base.py | 15 +++++++++++---- .../isaaclab_experimental/envs/frontend/warp.py | 17 +++++++++++++---- 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py index dece605f3aaa..120daf274aa7 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py @@ -301,11 +301,18 @@ def _classify_runtime(spec: Any) -> Runtime: if spec is None: return Runtime.UNKNOWN ep = spec.entry_point + # ``entry_point`` may be a ``"module:Class"`` string (the common case) + # or a class/callable object. Inspect ``__module__`` for the latter so + # warp-registered tasks classified as such regardless of the + # registration form. if isinstance(ep, str): - if ep.startswith(WARP_ROOT_PREFIXES): - return Runtime.WARP - if ep.startswith(("isaaclab.envs", "isaaclab_tasks")): - return Runtime.TORCH + module = ep + else: + module = getattr(ep, "__module__", "") or "" + if module.startswith(WARP_ROOT_PREFIXES): + return Runtime.WARP + if module.startswith(("isaaclab.envs", "isaaclab_tasks")): + return Runtime.TORCH return Runtime.UNKNOWN diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py index 9dbd3652400b..cb7257944406 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py @@ -331,14 +331,23 @@ def _mdp_modules(self, task: TaskMeta) -> list[ModuleType]: warp_pkg = entry.rsplit(".", 1)[0].replace(self.stable_root, self.warp_root, 1) parts = warp_pkg.split(".") for depth in range(len(parts), 0, -1): + target = ".".join(parts[:depth] + ["mdp"]) try: - modules.append(importlib.import_module(".".join(parts[:depth] + ["mdp"]))) + modules.append(importlib.import_module(target)) break - except ImportError: - continue + except ModuleNotFoundError as exc: + # Only swallow "this candidate module does not exist". + # An ImportError raised *inside* a module that does exist + # is a real bug we should surface, not paper over by + # falling back. + if exc.name == target: + continue + raise try: modules.append(importlib.import_module(self.fallback_mdp)) - except ImportError: + except ModuleNotFoundError as exc: + if exc.name != self.fallback_mdp: + raise logger.warning("WarpFrontend: fallback mdp module %r not importable", self.fallback_mdp) return modules From e2cb7444bd87284d52fc27120d4cb9bb40367aa5 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 19 May 2026 07:26:18 +0000 Subject: [PATCH 10/58] Replace frontend rule pipeline with single adapter function MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Collapses the CompatRule / Frontend / Report / Issue / Change framework (509+72+513+95 LOC across base/torch/warp/__init__) into one module (frontend.py, ~370 LOC). Net diff: 1270 → 864 LOC. Behavior changes: - Drop the pre-Hydra presets=newton injection. The warp frontend now hard-checks cfg.sim.physics is NewtonCfg at build time and tells the user to pass presets=newton on the CLI when it isn't. - Drop DropUnsupportedSensors. The Newton RayCaster (#5510) makes the height_scanner case obsolete; any future incompatible sensor should fail loudly with the sensor name rather than silently set None. - Merge SwapMdpFunctions + SwapActionClassType into one pass that swaps term.func or term.class_type uniformly for every group including actions. - Missing warp twin is always a hard failure (no strict / non-strict toggle). Partial swaps would leave torch funcs in a cfg consumed by warp managers, which only accept the kernel-style signature. - Replace SceneEntityCfg.__class__ = WarpSceneEntityCfg with a proper WarpSceneEntityCfg.from_stable() classmethod that copies every selection field through __init__. Train.py is unchanged in shape — two call sites, now using the top-level build(frontend, cfg, task_id) function. --- .../reinforcement_learning/rsl_rl/train.py | 43 +- .../changelog.d/warp-manager-bridge.rst | 88 ++- .../isaaclab_experimental/envs/frontend.py | 373 +++++++++++++ .../envs/frontend/__init__.py | 95 ---- .../envs/frontend/base.py | 509 ----------------- .../envs/frontend/torch.py | 72 --- .../envs/frontend/warp.py | 513 ------------------ .../managers/scene_entity_cfg.py | 24 + .../test/envs/test_frontend.py | 417 ++++++++++++++ 9 files changed, 864 insertions(+), 1270 deletions(-) create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py delete mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/frontend/__init__.py delete mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py delete mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/frontend/torch.py delete mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py create mode 100644 source/isaaclab_experimental/test/envs/test_frontend.py diff --git a/scripts/reinforcement_learning/rsl_rl/train.py b/scripts/reinforcement_learning/rsl_rl/train.py index f8f332bd18cb..23bff4218629 100644 --- a/scripts/reinforcement_learning/rsl_rl/train.py +++ b/scripts/reinforcement_learning/rsl_rl/train.py @@ -118,23 +118,14 @@ # intersection are pre-fold (the callback reads the user's original sys.argv), so # preset tokens like ``physics=NAME`` compare correctly here. Fold runs after. remaining_args = list_intersection(remaining_args, remaining_args_env_registration) -# Build the chosen frontend and let it pre-process Hydra args. Frontends -# may inject preset selections (e.g. WarpFrontend appends ``presets=newton`` -# for stable manager-based tasks so PresetCfg wrappers resolve to the Newton -# field before Hydra builds the cfg). Frontend injection runs *before* -# ``fold_preset_tokens`` so the injected tokens fold alongside any -# user-supplied ``physics=`` / ``presets=`` selectors. -# -# The frontend package lives under ``isaaclab_experimental``, which is -# optional. If it isn't installed, ``--frontend=torch`` (the default) still -# needs to work — fall back to plain ``gym.make`` in that case. The warp -# selection requires the experimental package, so import failure there is -# fatal. -_frontend = None -try: - from isaaclab_experimental.envs.frontend import get_frontend # noqa: E402 - _frontend = get_frontend(args_cli.frontend) +# Build path for ``--frontend=warp`` lives in ``isaaclab_experimental``, which is an +# optional package — torch (the default) must still work without it, so import lazily. +# Warp callers are required to pass ``presets=newton`` themselves; the frontend hard-checks +# it at build time rather than mutating Hydra args here. +_frontend_build = None +try: + from isaaclab_experimental.envs.frontend import build as _frontend_build # noqa: E402 except ImportError as _frontend_import_err: if args_cli.frontend != "torch": raise SystemExit( @@ -143,9 +134,6 @@ ) from _frontend_import_err logger.info("isaaclab_experimental not available; --frontend=torch will dispatch via plain gym.make.") -if _frontend is not None and args_cli.task is not None: - remaining_args = _frontend.preprocess_hydra_args(args_cli.task, remaining_args) - sys.argv = [sys.argv[0]] + fold_preset_tokens(remaining_args) # -- check RSL-RL version ---------------------------------------------------- @@ -230,16 +218,15 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen # set the log directory for the environment (works for all environment types) env_cfg.log_dir = log_dir - # create isaac environment via the chosen frontend's rule pipeline. - # The frontend may mutate env_cfg in place (preset resolution, - # SceneEntityCfg promotion, mdp twin swaps, ...) before constructing - # the env. ``env.unwrapped.frontend_report`` carries what changed - # and what was missing. When the experimental frontend package isn't - # available we already validated --frontend=torch above, so a plain - # gym.make is the documented fallback. + # Build the env via the selected frontend. The warp frontend mutates env_cfg + # in place (SceneEntityCfg promotion + MDP func/class_type swap) and then + # constructs ``ManagerBasedRLEnvWarp``; missing warp twins are a hard failure. + # When ``isaaclab_experimental`` isn't installed, the argparse step above + # has already constrained --frontend to ``torch``, so the plain ``gym.make`` + # fallback is safe. render_mode = "rgb_array" if args_cli.video else None - if _frontend is not None: - env = _frontend.build(env_cfg, task_id=args_cli.task, render_mode=render_mode) + if _frontend_build is not None: + env = _frontend_build(args_cli.frontend, env_cfg, args_cli.task, render_mode=render_mode) else: env = gym.make(args_cli.task, cfg=env_cfg, render_mode=render_mode) diff --git a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst index d7197a0a756e..92dfd4f06f7b 100644 --- a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst +++ b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst @@ -1,59 +1,41 @@ Added ^^^^^ -* Added :mod:`isaaclab_experimental.envs.frontend`, a small framework that - selects the runtime backend for IsaacLab tasks. A :class:`Frontend` takes - a stable env cfg + task id, runs a pluggable :class:`CompatRule` pipeline - against the (cfg, task, frontend) triple, and constructs the env on the - matching runtime. New runtimes plug in by subclassing :class:`Frontend` - and calling :func:`register_frontend`; new compatibility checks plug in - by subclassing :class:`CompatRule` and listing it in a frontend's - :attr:`Frontend.rules`. - - The framework ships two built-in frontends: - - * :class:`TorchFrontend` (``--frontend=torch``, default) — passes through - to :func:`gym.make`. Emits a warning if the task is registered under - the warp runtime. - * :class:`WarpFrontend` (``--frontend=warp``) — adapts a manager-based - stable cfg onto :class:`ManagerBasedRLEnvWarp` via the rule pipeline, - or dispatches a direct task to its registered warp env class. - - The default warp rules are: - - * :class:`CheckPhysicsIsNewton` — blocking check; PhysX physics with the - warp runtime is a hard incompatibility (``isaaclab_physx`` classes - depend on ``omni.physics.tensors``, which the warp runtime does not - initialise). - * :class:`ResolvePhysicsPreset` — collapses ``PresetCfg`` to its - ``newton`` field for programmatic callers (Hydra's preset resolution - handles the CLI case). - * :class:`DropUnsupportedSensors` — drops scene sensors warp can't run. - * :class:`PromoteSceneEntityCfg` — in-place class promotion for warp - cached fields, with an :func:`issubclass` assertion to fail loudly if - the warp class ever stops subclassing the stable one. - * :class:`SwapMdpFunctions` — name-based ``term.func`` replacement - against the warp ``mdp`` modules; rejects stable re-exports by - inspecting ``__module__``. - * :class:`SwapActionClassType` — strict swap of action ``class_type`` - to the warp twin; missing twin is :attr:`Severity.BLOCKING`. - * :class:`VerifyDirectIsWarp` — for direct cfgs, blocks if the - registered entry-point isn't a warp class. - - All rules emit :class:`Issue` (incompatibility) and / or :class:`Change` - (transformation applied) records into a :class:`Report` accessible on - ``env.unwrapped.frontend_report`` after :meth:`Frontend.build`. - - :class:`TaskResolver` centralises ``gym.spec`` introspection and - classifies a task into a :class:`TaskMeta` (workflow, registered - runtime, cfg class). Rules read from :class:`TaskMeta` rather than - poking gym directly. - -* Added a ``--frontend={torch,warp}`` flag to ``rsl_rl/train.py``. The - script delegates Hydra-arg pre-processing (``presets=newton`` injection, - conflicting-preset warnings) and env construction to the selected - frontend via :func:`get_frontend`; ``render_mode`` is forwarded so - ``--video`` keeps working under the warp frontend. +* Added :mod:`isaaclab_experimental.envs.frontend`, a small runtime selector + used by ``--frontend {torch,warp}`` to choose how a task is constructed. + + * ``torch`` (default) dispatches via :func:`gym.make` unchanged. + * ``warp`` adapts a stable manager-based cfg in place and constructs + :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp`. The adapter + does three things, all hard-failing on incompatibility: + + 1. Requires :class:`~isaaclab_newton.physics.NewtonCfg` as the active + physics — the user is responsible for selecting it via + ``presets=newton``; no Hydra-arg mutation happens behind the scenes. + 2. Promotes every stable :class:`~isaaclab.managers.SceneEntityCfg` + embedded under term ``params`` to + :class:`isaaclab_experimental.managers.SceneEntityCfg` via the new + :meth:`~isaaclab_experimental.managers.SceneEntityCfg.from_stable` + classmethod (no ``__class__`` reassignment). + 3. Swaps every stable ``term.func`` *and* ``term.class_type`` (one pass, + handles observations / events / rewards / terminations / commands / + curriculum / actions) with its same-named warp twin from + ``isaaclab_tasks_experimental..mdp`` or + :mod:`isaaclab_experimental.envs.mdp`. Any missing twin raises + :class:`FrontendIncompatibleError` listing the affected term paths. + + Direct workflows aren't adapted; ``--frontend=warp`` on a direct task + requires the task to be pre-registered under ``isaaclab_experimental`` / + ``isaaclab_tasks_experimental`` (e.g. ``*-Direct-Warp-v0``). + +* Added :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable`, + a classmethod that builds a warp scene-entity cfg from a stable one by + copying every selection field through ``__init__``. + +* Added a ``--frontend {torch,warp}`` flag to ``rsl_rl/train.py``. The flag + selects the runtime via :func:`isaaclab_experimental.envs.frontend.build`; + ``isaaclab_experimental`` is treated as optional and ``--frontend=torch`` + falls back to ``gym.make`` when it isn't installed. Fixed ^^^^^ diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py new file mode 100644 index 000000000000..4dfc312d6481 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -0,0 +1,373 @@ +# 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 + +"""Runtime selector for IsaacLab tasks (``--frontend {torch,warp}``). + +Two responsibilities only: + +1. Decide which runtime constructs the env (torch via ``gym.make``, or warp via + :class:`ManagerBasedRLEnvWarp`). +2. For the warp runtime, adapt a stable manager-based cfg in place so the + warp managers can consume it (Newton physics + warp-twin term funcs and + action classes + warp variant of :class:`SceneEntityCfg`). + +The adaptation is a single function (:func:`_adapt_cfg_for_warp`); there is no +plugin framework. All warp MDP twins (functions *and* classes) must exist for +the chosen task — a missing twin is a hard failure, not a silent drop. +""" + +from __future__ import annotations + +import importlib +import logging +from collections.abc import Iterable, Iterator +from enum import StrEnum +from types import ModuleType +from typing import Any + +import gymnasium as gym + +from isaaclab.envs import DirectMARLEnvCfg, DirectRLEnvCfg, ManagerBasedRLEnvCfg +from isaaclab.managers.scene_entity_cfg import SceneEntityCfg as _StableSceneEntityCfg + +logger = logging.getLogger(__name__) + + +__all__ = [ + "Frontend", + "FrontendIncompatibleError", + "Workflow", + "build", +] + + +# --------------------------------------------------------------------------- +# Enums + error +# --------------------------------------------------------------------------- + + +class Frontend(StrEnum): + """Runtime selector exposed by ``--frontend``.""" + + TORCH = "torch" + WARP = "warp" + + +class Workflow(StrEnum): + """Manager-based vs direct env workflow.""" + + MANAGER_BASED = "manager_based" + DIRECT = "direct" + + +class FrontendIncompatibleError(RuntimeError): + """Raised when the chosen frontend can't run the requested task.""" + + +# --------------------------------------------------------------------------- +# Public entry point +# --------------------------------------------------------------------------- + + +# Prefixes used to decide whether a registered task's entry-point lives under +# the warp packages. Used by the direct-workflow guard. +_WARP_ROOT_PREFIXES: tuple[str, ...] = ("isaaclab_experimental", "isaaclab_tasks_experimental") + +# Term containers walked by the cfg-adaptation step. Actions are included so +# the swap logic handles ``class_type`` and ``func`` in one pass. +_TERM_PATHS: tuple[tuple[str, ...], ...] = ( + ("observations", "policy"), + ("events",), + ("rewards",), + ("terminations",), + ("commands",), + ("curriculum",), + ("actions",), +) + + +def build( + frontend: Frontend | str, + env_cfg: Any, + task_id: str, + **construct_kwargs: Any, +) -> gym.Env: + """Construct the env on the selected runtime. + + Args: + frontend: ``"torch"`` (default IsaacLab path) or ``"warp"`` (warp managers + + :class:`ManagerBasedRLEnvWarp` for manager-based tasks). + env_cfg: Stable env cfg. Mutated in place when ``frontend == "warp"``. + task_id: Gym registration id, e.g. ``"Isaac-Cartpole-v0"``. + **construct_kwargs: Forwarded to the env constructor (``render_mode``, …). + + Returns: + The constructed :class:`gym.Env`. + + Raises: + FrontendIncompatibleError: If the warp runtime can't run the task + (wrong physics, missing MDP twins, direct task not registered + as a warp env). + """ + frontend = Frontend(frontend) + if frontend is Frontend.TORCH: + return gym.make(task_id, cfg=env_cfg, **construct_kwargs) + + workflow = _detect_workflow(env_cfg) + if workflow is Workflow.DIRECT: + # Direct workflows aren't adapted — they must already be registered + # under the warp packages (e.g. ``Isaac-Cartpole-Direct-Warp-v0``). + _require_direct_is_warp_task(task_id) + return gym.make(task_id, cfg=env_cfg, **construct_kwargs) + + _adapt_cfg_for_warp(env_cfg, task_id) + # Imported lazily so that ``--frontend=torch`` callers don't pay the + # ``isaaclab_experimental.envs`` import cost. + from isaaclab_experimental.envs import ManagerBasedRLEnvWarp + + return ManagerBasedRLEnvWarp(cfg=env_cfg, **construct_kwargs) + + +# --------------------------------------------------------------------------- +# Cfg adaptation (warp only) +# --------------------------------------------------------------------------- + + +def _adapt_cfg_for_warp(cfg: Any, task_id: str) -> None: + """Mutate ``cfg`` in place so warp managers can consume it. + + Three steps, each independently testable: + + 1. :func:`_require_newton_physics` — hard check that ``cfg.sim.physics`` is + :class:`~isaaclab_newton.physics.NewtonCfg`. The user is responsible for + selecting the Newton variant of the task's :class:`PresetCfg` via + ``presets=newton``; we don't auto-inject. + 2. :func:`_promote_scene_entity_cfgs` — replace stable + :class:`~isaaclab.managers.SceneEntityCfg` instances under each term's + ``params`` with the warp variant (which adds warp-cached ``joint_mask``, + ``joint_ids_wp``, ``body_ids_wp`` fields). + 3. :func:`_swap_mdp` — for every MDP term in every group (observations, + events, rewards, terminations, commands, curriculum, actions), replace + any stable ``func`` *or* ``class_type`` with its same-named warp twin. + A missing twin raises :class:`FrontendIncompatibleError` — partial + coverage is unsafe under the warp managers' kernel-only signature. + """ + _require_newton_physics(cfg, task_id) + _promote_scene_entity_cfgs(cfg) + _swap_mdp(cfg, task_id) + + +def _require_newton_physics(cfg: Any, task_id: str) -> None: + """Block unless ``cfg.sim.physics`` is :class:`NewtonCfg`. + + The warp managers' assets read state through :class:`NewtonManager`; + a :class:`PhysxCfg` (or unresolved :class:`PresetCfg`) is a hard + incompatibility. The fix is to pass ``presets=newton`` on the CLI so + Hydra resolves the task's :class:`PresetCfg` wrapper to the Newton field + before construction. + """ + from isaaclab_newton.physics import NewtonCfg + + physics = getattr(getattr(cfg, "sim", None), "physics", None) + if isinstance(physics, NewtonCfg): + return + raise FrontendIncompatibleError( + f"--frontend=warp on {task_id!r}: expected cfg.sim.physics to be NewtonCfg," + f" got {type(physics).__name__!r}. Pass `presets=newton` on the CLI so" + f" Hydra resolves the task's PresetCfg wrapper to the Newton variant." + ) + + +def _promote_scene_entity_cfgs(cfg: Any) -> None: + """Replace stable :class:`SceneEntityCfg` instances with the warp variant. + + Walks every ``term.params: dict`` under each term group and rebuilds any + stable :class:`SceneEntityCfg` value via + :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable`. The + warp variant subclasses the stable one, so type checks elsewhere stay + valid; the new fields (``joint_mask`` / ``joint_ids_wp`` / ``body_ids_wp``) + are filled at :meth:`resolve` time by the warp scene. + """ + from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as _WarpSceneEntityCfg + + promoted = 0 + for path in _TERM_PATHS: + group = _walk_attrs(cfg, path) + if group is None: + continue + for _name, term in _iter_term_attrs(group): + params = getattr(term, "params", None) + if not isinstance(params, dict): + continue + for key, value in list(params.items()): + if isinstance(value, _WarpSceneEntityCfg) or not isinstance(value, _StableSceneEntityCfg): + continue + params[key] = _WarpSceneEntityCfg.from_stable(value) + promoted += 1 + if promoted: + logger.info("frontend.warp: promoted %d SceneEntityCfg instance(s) to warp variant", promoted) + + +def _swap_mdp(cfg: Any, task_id: str) -> None: + """Replace ``term.func`` and ``term.class_type`` with their warp twins. + + Walks the same term paths used by :func:`_promote_scene_entity_cfgs` and + on each term swaps whichever of ``func`` / ``class_type`` is set to a + stable-origin symbol. Twin lookup is name-based against the task's + matching ``isaaclab_tasks_experimental.<...>.mdp`` module and, as a + fallback, :mod:`isaaclab_experimental.envs.mdp`. Any missing twin raises + :class:`FrontendIncompatibleError` listing every affected term — partial + swaps would leave torch-style callables in the cfg and the warp managers + would call them with the wrong signature. + """ + modules = _warp_mdp_modules(task_id) + searched = tuple(m.__name__ for m in modules) + logger.info("frontend.warp: searching warp mdp modules %s", list(searched)) + + swapped = 0 + missing: list[tuple[str, str, str]] = [] # (location, attr, symbol) + for path in _TERM_PATHS: + group = _walk_attrs(cfg, path) + if group is None: + continue + location_prefix = ".".join(path) + for name, term in _iter_term_attrs(group): + for attr in ("func", "class_type"): + stable = getattr(term, attr, None) + if stable is None or not _is_swap_candidate(stable): + continue + twin = _resolve_warp_twin(stable.__name__, modules) + if twin is None: + missing.append((f"{location_prefix}.{name}", attr, stable.__name__)) + continue + setattr(term, attr, twin) + swapped += 1 + + if missing: + lines = "\n ".join(f"{loc}.{attr}: no warp twin for {sym!r}" for loc, attr, sym in missing) + raise FrontendIncompatibleError( + f"--frontend=warp on {task_id!r}: missing warp MDP twins (searched {list(searched)}):\n {lines}" + ) + + logger.info("frontend.warp: swapped %d MDP symbol(s) to warp twins", swapped) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _detect_workflow(cfg: Any) -> Workflow: + """Classify the env cfg into manager-based or direct (used to pick build path).""" + if isinstance(cfg, ManagerBasedRLEnvCfg): + return Workflow.MANAGER_BASED + if isinstance(cfg, (DirectRLEnvCfg, DirectMARLEnvCfg)): + return Workflow.DIRECT + raise FrontendIncompatibleError( + f"Unrecognised env cfg type {type(cfg).__name__!r};" + f" expected ManagerBasedRLEnvCfg / DirectRLEnvCfg / DirectMARLEnvCfg subclass." + ) + + +def _require_direct_is_warp_task(task_id: str) -> None: + """For direct workflows, the task must be pre-registered under the warp packages.""" + try: + spec = gym.spec(task_id) + except gym.error.NameNotFound as exc: + raise FrontendIncompatibleError(f"--frontend=warp: task {task_id!r} is not registered with gymnasium.") from exc + ep = spec.entry_point + module = ep if isinstance(ep, str) else (getattr(ep, "__module__", "") or "") + if not module.startswith(_WARP_ROOT_PREFIXES): + raise FrontendIncompatibleError( + f"--frontend=warp on direct task {task_id!r}: entry_point {ep!r}" + f" is not under {list(_WARP_ROOT_PREFIXES)}. Direct tasks must be" + f" registered as a warp env class (e.g. *-Direct-Warp-v0)." + ) + + +def _warp_mdp_modules(task_id: str) -> list[ModuleType]: + """Locate warp MDP modules to consult for twin lookups. + + Order of preference: + + 1. The task-specific module, derived by replacing ``isaaclab_tasks`` with + ``isaaclab_tasks_experimental`` in the task's ``env_cfg_entry_point`` + package and walking up to the first existing ``.mdp`` submodule. + 2. The shared :mod:`isaaclab_experimental.envs.mdp` fallback (where + generic warp twins live). + """ + modules: list[ModuleType] = [] + try: + spec = gym.spec(task_id) + except gym.error.NameNotFound: + spec = None + entry = spec.kwargs.get("env_cfg_entry_point") if spec is not None else None + if isinstance(entry, str) and entry.startswith("isaaclab_tasks."): + warp_pkg = entry.rsplit(".", 1)[0].replace("isaaclab_tasks", "isaaclab_tasks_experimental", 1) + parts = warp_pkg.split(".") + for depth in range(len(parts), 0, -1): + target = ".".join(parts[:depth] + ["mdp"]) + try: + modules.append(importlib.import_module(target)) + break + except ModuleNotFoundError as exc: + # Only swallow "this candidate doesn't exist" — a real + # import error inside an existing module is a bug. + if exc.name == target: + continue + raise + # Generic fallback. + fallback = "isaaclab_experimental.envs.mdp" + try: + modules.append(importlib.import_module(fallback)) + except ModuleNotFoundError as exc: + if exc.name == fallback: + logger.warning("frontend.warp: fallback mdp module %r not importable", fallback) + else: + raise + return modules + + +def _resolve_warp_twin(name: str, modules: Iterable[ModuleType]) -> Any | None: + """Return the same-named symbol from ``modules`` that originates under the warp packages.""" + for module in modules: + candidate = getattr(module, name, None) + if candidate is None: + continue + origin = getattr(candidate, "__module__", "") or "" + if origin.startswith(_WARP_ROOT_PREFIXES): + return candidate + return None + + +def _is_swap_candidate(value: Any) -> bool: + """Heuristic: callable or class whose origin is *not already* warp.""" + if not callable(value): + return False + origin = getattr(value, "__module__", "") or "" + if origin.startswith(_WARP_ROOT_PREFIXES): + return False # already a warp twin (idempotent) + return True + + +def _walk_attrs(root: Any, path: tuple[str, ...]) -> Any: + """Walk ``root..…``; return ``None`` on any miss.""" + node = root + for attr in path: + node = getattr(node, attr, None) + if node is None: + return None + return node + + +def _iter_term_attrs(group: Any) -> Iterator[tuple[str, Any]]: + """Yield ``(name, term)`` pairs from a manager-cfg group, skipping dunders/Nones.""" + if group is None: + return + for name in [n for n in dir(group) if not n.startswith("_")]: + term = getattr(group, name, None) + if term is None: + continue + yield name, term diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/__init__.py deleted file mode 100644 index 0f6e98b33974..000000000000 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/__init__.py +++ /dev/null @@ -1,95 +0,0 @@ -# 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 - -"""Frontend framework: pluggable runtime selectors for IsaacLab tasks. - -A frontend (chosen via ``--frontend {torch,warp}``) is a thin dispatcher that -takes a stable env cfg + task id, runs a pluggable :class:`CompatRule` -pipeline against the (cfg, task, frontend) triple, and constructs the env on -the matching runtime. New runtimes plug in by subclassing :class:`Frontend` -and calling :func:`register_frontend`; new compatibility checks plug in by -subclassing :class:`CompatRule` and adding it to a frontend's -:attr:`Frontend.rules`. - -Public API:: - - from isaaclab_experimental.envs.frontend import get_frontend - frontend = get_frontend("warp") - env = frontend.build(env_cfg, task_id="Isaac-Cartpole-v0") - print(env.unwrapped.frontend_report.format()) -""" - -from __future__ import annotations - -from .base import ( - Change, - CompatRule, - Frontend, - FrontendIncompatibleError, - Issue, - Report, - ResolveContext, - Runtime, - Severity, - TaskMeta, - TaskResolver, - Workflow, - available_frontends, - get_frontend, - register_frontend, -) -from .torch import TorchFrontend, WarnIfTaskIsWarpRegistered -from .warp import ( - CheckPhysicsIsNewton, - DropUnsupportedSensors, - PromoteSceneEntityCfg, - ResolvePhysicsPreset, - SwapActionClassType, - SwapMdpFunctions, - VerifyDirectIsWarp, - WarpFrontend, -) - -# Register built-ins so ``get_frontend("torch" | "warp")`` works without -# the caller importing the concrete classes. -register_frontend(TorchFrontend.name, TorchFrontend) -register_frontend(WarpFrontend.name, WarpFrontend) - - -# Public API. Internal helpers (``walk_attrs``, ``iter_term_attrs``, -# ``resolve_warp_twin``, ``WARP_ROOT_PREFIXES``) are importable from -# :mod:`isaaclab_experimental.envs.frontend.base` for users writing their own -# rules, but they are not advertised here as the supported framework surface. -__all__ = [ - # core abstractions - "Frontend", - "CompatRule", - "TaskMeta", - "TaskResolver", - "Report", - "Issue", - "Change", - "Severity", - "Workflow", - "Runtime", - "ResolveContext", - "FrontendIncompatibleError", - # registry - "register_frontend", - "get_frontend", - "available_frontends", - # built-in frontends - "TorchFrontend", - "WarpFrontend", - # built-in rules - "WarnIfTaskIsWarpRegistered", - "CheckPhysicsIsNewton", - "ResolvePhysicsPreset", - "DropUnsupportedSensors", - "PromoteSceneEntityCfg", - "SwapMdpFunctions", - "SwapActionClassType", - "VerifyDirectIsWarp", -] diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py deleted file mode 100644 index 120daf274aa7..000000000000 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/base.py +++ /dev/null @@ -1,509 +0,0 @@ -# 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 - -"""Frontend framework — abstractions for runtime selection. - -A *frontend* is a user-facing runtime selector (chosen via ``--frontend -{torch,warp}``). It takes a stable env cfg and a registered gym task id and -returns a runnable env, having first run a pluggable pipeline of -compatibility checks and transforms against the (cfg, task, frontend) -triple. - -The framework is intentionally narrow: - -* :class:`TaskResolver` is the single point that reads ``gym.spec`` and - classifies a task into a :class:`TaskMeta` (workflow, registered - runtime, cfg class). Rules and frontends consume :class:`TaskMeta` - instead of poking gym directly, so the registration format can evolve - in one place. -* :class:`CompatRule` is a single check or transform applied during - resolution. Each rule yields :class:`Issue` records (incompatibilities - the frontend can't paper over) and / or :class:`Change` records - (transformations applied to the cfg). Subclasses implement - :meth:`CompatRule.run` and may override :meth:`CompatRule.applies_to` - to scope themselves by workflow / runtime. -* :class:`Frontend` is the dispatcher. Subclasses declare a name, a list - of rule classes, and a :meth:`Frontend.construct` strategy that builds - the env once rules pass. -* :func:`register_frontend` / :func:`get_frontend` form the registry CLI - hooks read. - -To add a new compatibility check, write a small :class:`CompatRule` and -list it in the relevant frontend's :attr:`Frontend.rules`. To add a new -runtime, subclass :class:`Frontend` and call :func:`register_frontend`. -""" - -from __future__ import annotations - -import logging -from abc import ABC, abstractmethod -from collections.abc import Iterable, Iterator -from dataclasses import dataclass, field -from enum import Enum -from typing import Any, ClassVar - -logger = logging.getLogger(__name__) - - -# --------------------------------------------------------------------------- -# Constants -# --------------------------------------------------------------------------- - - -WARP_ROOT_PREFIXES: tuple[str, ...] = ("isaaclab_experimental", "isaaclab_tasks_experimental") -"""Module prefixes that mark a class as a warp implementation. - -The warp ``mdp`` packages do ``from isaaclab.envs.mdp import *``, so a -plain ``getattr`` lookup can return a stable symbol untouched. We accept -a candidate as a real warp twin only if its ``__module__`` lives under -one of these prefixes.""" - - -# --------------------------------------------------------------------------- -# Enums -# --------------------------------------------------------------------------- - - -class Workflow(str, Enum): - """Task workflow type.""" - - MANAGER_BASED = "manager_based" - DIRECT = "direct" - DIRECT_MARL = "direct_marl" - - -class Runtime(str, Enum): - """Runtime backend a task is registered against. - - A task whose ``entry_point`` lives under :data:`WARP_ROOT_PREFIXES` is - classified as ``WARP``. Tasks registered against ``isaaclab.envs.*`` - (the standard manager-based and direct env classes) are classified - as ``TORCH``; those env classes use ``FactoryBase`` to dispatch on - the active physics backend at construction time, so PhysX and Newton - physics both flow through the torch runtime path. - """ - - TORCH = "torch" - WARP = "warp" - UNKNOWN = "unknown" - - -class Severity(str, Enum): - """Severity of a compatibility :class:`Issue`.""" - - BLOCKING = "blocking" - """Frontend cannot build the env. :meth:`Frontend.build` raises.""" - WARNING = "warning" - """Surfaced to the user but not fatal (e.g. a term was dropped).""" - - -# --------------------------------------------------------------------------- -# Records -# --------------------------------------------------------------------------- - - -@dataclass -class TaskMeta: - """Static description of a registered gym task, produced by :class:`TaskResolver`.""" - - task_id: str - """Gym task id, e.g. ``Isaac-Cartpole-v0``.""" - spec: Any - """The :class:`gymnasium.envs.registration.EnvSpec` returned by ``gym.spec``.""" - cfg_class: type - """Concrete cfg type passed to the frontend.""" - workflow: Workflow - """Manager-based / direct / direct-marl.""" - runtime: Runtime - """Runtime the task is registered against (independent of the frontend asked for).""" - - @property - def env_cfg_entry_point(self) -> str | None: - """Module path of the cfg class as registered in ``spec.kwargs``.""" - return self.spec.kwargs.get("env_cfg_entry_point") if self.spec is not None else None - - @property - def entry_point(self) -> Any: - """Env class entry point as registered (string or class).""" - return self.spec.entry_point if self.spec is not None else None - - -@dataclass -class Issue: - """An incompatibility surfaced by a rule.""" - - rule: str - """Rule that produced this issue (matches :attr:`CompatRule.name`).""" - severity: Severity - message: str - """One-line human-readable description.""" - location: str = "" - """Optional cfg path the issue was found at (e.g. ``rewards.foo``).""" - detail: dict[str, Any] = field(default_factory=dict) - """Free-form rule-specific data (searched modules, expected names, ...).""" - - -@dataclass -class Change: - """A transformation a rule applied to the cfg.""" - - rule: str - description: str - - -@dataclass -class Report: - """Outcome of a frontend's rule pipeline.""" - - frontend: str - """Name of the frontend that produced this report.""" - task: TaskMeta - rules_run: list[str] = field(default_factory=list) - issues: list[Issue] = field(default_factory=list) - changes: list[Change] = field(default_factory=list) - - @property - def blocking(self) -> list[Issue]: - return [i for i in self.issues if i.severity == Severity.BLOCKING] - - def has_blocking(self) -> bool: - return any(i.severity == Severity.BLOCKING for i in self.issues) - - def __bool__(self) -> bool: - return bool(self.changes) or bool(self.issues) - - def format(self) -> str: - lines = [ - f"Frontend {self.frontend!r} on task {self.task.task_id!r}" - f" ({self.task.workflow.value} / registered runtime {self.task.runtime.value})" - ] - if self.changes: - lines.append(f" changes ({len(self.changes)}):") - lines.extend(f" [{c.rule}] {c.description}" for c in self.changes) - if self.issues: - lines.append(f" issues ({len(self.issues)}):") - for i in self.issues: - head = f" [{i.severity.value}/{i.rule}]" - where = f" {i.location}" if i.location else "" - lines.append(f"{head}{where} {i.message}") - return "\n".join(lines) - - -@dataclass -class ResolveContext: - """Per-call context handed to rules.""" - - frontend: str - """Name of the running frontend.""" - task: TaskMeta - strict: bool = False - """When True, rules may escalate warnings to blocking issues.""" - - -# --------------------------------------------------------------------------- -# Exceptions -# --------------------------------------------------------------------------- - - -class FrontendIncompatibleError(RuntimeError): - """Raised when a frontend cannot build the requested (task, cfg) pair.""" - - def __init__(self, message: str, report: Report | None = None): - super().__init__(message) - self.report = report - - -# --------------------------------------------------------------------------- -# CompatRule -# --------------------------------------------------------------------------- - - -class CompatRule(ABC): - """A single check / transform applied during frontend resolution. - - A rule's :meth:`run` method may yield :class:`Issue` records (the rule - detected an incompatibility), :class:`Change` records (the rule mutated - the cfg), or both. Pure-check and pure-transform rules are common; mixed - rules are useful when a transformation has a partial-failure mode (e.g. - a name swap that found 5 of 6 twins). - - Subclass and: - - * Set :attr:`name` to a short identifier used in the report. - * Implement :meth:`run` to do the work. - * Override :meth:`applies_to` to skip the rule when irrelevant. - """ - - name: ClassVar[str] - - def applies_to(self, ctx: ResolveContext) -> bool: - """Return True if this rule should run on the given context. - - Default: always applies. Override to scope by ``ctx.task.workflow``, - ``ctx.task.runtime``, ``ctx.frontend``, etc. - """ - return True - - @abstractmethod - def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: - """Inspect / mutate ``cfg`` and yield :class:`Issue` / :class:`Change` records.""" - - -# --------------------------------------------------------------------------- -# TaskResolver -# --------------------------------------------------------------------------- - - -class TaskResolver: - """Inspect a registered gym task and classify its workflow + runtime. - - All rules and frontends should call :meth:`resolve` rather than reading - ``gym.spec`` directly. Centralising the read keeps registration-format - knowledge in one place. - """ - - @classmethod - def resolve(cls, task_id: str, cfg: Any) -> TaskMeta: - """Return a :class:`TaskMeta` for ``(task_id, cfg)``.""" - import gymnasium as gym - - try: - spec = gym.spec(task_id) - except gym.error.NameNotFound: - spec = None - return TaskMeta( - task_id=task_id, - spec=spec, - cfg_class=type(cfg), - workflow=cls._classify_workflow(cfg), - runtime=cls._classify_runtime(spec), - ) - - @staticmethod - def _classify_workflow(cfg: Any) -> Workflow: - from isaaclab.envs import DirectMARLEnvCfg, DirectRLEnvCfg, ManagerBasedRLEnvCfg - - if isinstance(cfg, ManagerBasedRLEnvCfg): - return Workflow.MANAGER_BASED - if isinstance(cfg, DirectMARLEnvCfg): - return Workflow.DIRECT_MARL - if isinstance(cfg, DirectRLEnvCfg): - return Workflow.DIRECT - raise FrontendIncompatibleError( - f"Cfg type {type(cfg).__name__!r} is not a recognised env cfg" - f" (expected ManagerBasedRLEnvCfg / DirectRLEnvCfg / DirectMARLEnvCfg subclass)." - ) - - @staticmethod - def _classify_runtime(spec: Any) -> Runtime: - if spec is None: - return Runtime.UNKNOWN - ep = spec.entry_point - # ``entry_point`` may be a ``"module:Class"`` string (the common case) - # or a class/callable object. Inspect ``__module__`` for the latter so - # warp-registered tasks classified as such regardless of the - # registration form. - if isinstance(ep, str): - module = ep - else: - module = getattr(ep, "__module__", "") or "" - if module.startswith(WARP_ROOT_PREFIXES): - return Runtime.WARP - if module.startswith(("isaaclab.envs", "isaaclab_tasks")): - return Runtime.TORCH - return Runtime.UNKNOWN - - -# --------------------------------------------------------------------------- -# Frontend -# --------------------------------------------------------------------------- - - -class Frontend(ABC): - """User-facing runtime selector. - - Subclasses declare: - - * :attr:`name`: CLI identifier used by ``--frontend`` and the registry. - * :attr:`rules`: rule classes the pipeline instantiates by default. - * :meth:`construct`: how to build the env once rules pass. - - Subclasses may also override: - - * :meth:`preprocess_hydra_args`: CLI-time Hydra arg munging (e.g. inject - a preset selection so Hydra resolves ``PresetCfg`` to the right field - before the cfg ever reaches us). - * :meth:`preprocess_cfg`: cfg pre-processing run before rules. - """ - - name: ClassVar[str] - rules: ClassVar[tuple[type[CompatRule], ...]] = () - - def __init__(self, rules: Iterable[CompatRule] | None = None, strict: bool = False): - self._rules: list[CompatRule] = [r() for r in type(self).rules] if rules is None else list(rules) - self.strict = strict - - # -- public --------------------------------------------------------------- - - def build(self, cfg: Any, task_id: str, **construct_kwargs: Any) -> Any: - """Resolve the task, run rules, and construct the env. - - Raises :class:`FrontendIncompatibleError` if any rule emits a - :class:`Severity.BLOCKING` issue. The returned env carries the - :class:`Report` on ``env.unwrapped.frontend_report`` for inspection. - """ - report = self.resolve(cfg, task_id) - self._log_report(report) - if report.has_blocking(): - raise FrontendIncompatibleError( - f"Frontend {self.name!r} cannot build {task_id!r}:\n{report.format()}", - report=report, - ) - env = self.construct(cfg, report.task, **construct_kwargs) - env.unwrapped.frontend_report = report - return env - - def resolve(self, cfg: Any, task_id: str) -> Report: - """Run the rule pipeline; mutate ``cfg`` in place; return a :class:`Report`. - - Use this directly for dry-run validation (tests, CLI ``--check``). - """ - meta = TaskResolver.resolve(task_id, cfg) - ctx = ResolveContext(frontend=self.name, task=meta, strict=self.strict) - report = Report(frontend=self.name, task=meta) - # If gym couldn't resolve the spec, the frontend has nothing to dispatch - # against. Block early with a clear message rather than letting - # downstream construct() fail with a NameNotFound. - if meta.spec is None: - report.issues.append( - Issue( - rule="task_resolver", - severity=Severity.BLOCKING, - message=( - f"task {task_id!r} is not registered with gymnasium." - " Make sure the task package is imported before the frontend runs." - ), - location="task.spec", - ) - ) - return report - self.preprocess_cfg(cfg, ctx) - for rule in self._rules: - if not rule.applies_to(ctx): - continue - report.rules_run.append(rule.name) - for record in rule.run(cfg, ctx): - if isinstance(record, Issue): - report.issues.append(record) - elif isinstance(record, Change): - report.changes.append(record) - else: - raise TypeError(f"Rule {rule.name!r} yielded {type(record).__name__}; expected Issue or Change.") - return report - - # -- subclass hooks ------------------------------------------------------- - - @abstractmethod - def construct(self, cfg: Any, meta: TaskMeta, **kwargs: Any) -> Any: - """Build the env. Called only when no blocking issues were raised.""" - - def preprocess_cfg(self, cfg: Any, ctx: ResolveContext) -> None: - """Subclass hook for cfg pre-processing run before rules. Default: no-op.""" - - def preprocess_hydra_args(self, task_id: str, args: list[str]) -> list[str]: - """Subclass hook for CLI-time Hydra arg munging. Default: no-op. - - Called from the train script so frontends can inject preset - selections (or refuse incompatible ones) before Hydra builds the - cfg. Returns the (possibly modified) args list. - """ - return args - - # -- helpers -------------------------------------------------------------- - - def _log_report(self, report: Report) -> None: - if not report: - return - level = logging.ERROR if report.has_blocking() else logging.WARNING - logger.log(level, report.format()) - - -# --------------------------------------------------------------------------- -# Registry -# --------------------------------------------------------------------------- - - -_REGISTRY: dict[str, type[Frontend]] = {} - - -def register_frontend(name: str, cls: type[Frontend]) -> None: - """Register a frontend implementation under ``name``. - - The CLI ``--frontend`` flag and :func:`get_frontend` look up by this name. - Re-registering the same class is idempotent; re-registering a different - class raises :class:`ValueError`. - """ - existing = _REGISTRY.get(name) - if existing is not None and existing is not cls: - raise ValueError( - f"Frontend {name!r} is already registered to {existing.__name__};" - f" refusing to re-register to {cls.__name__}." - ) - _REGISTRY[name] = cls - - -def get_frontend(name: str, **kwargs: Any) -> Frontend: - """Look up a frontend by name and return a fresh instance. - - Extra kwargs are forwarded to :meth:`Frontend.__init__` (e.g. ``strict=True``). - """ - try: - cls = _REGISTRY[name] - except KeyError as exc: - raise ValueError(f"Unknown frontend {name!r}; available: {available_frontends()}.") from exc - return cls(**kwargs) - - -def available_frontends() -> list[str]: - """Return registered frontend names, sorted.""" - return sorted(_REGISTRY) - - -# --------------------------------------------------------------------------- -# Helpers shared by rules -# --------------------------------------------------------------------------- - - -def walk_attrs(root: Any, path: tuple[str, ...]) -> Any: - """Walk ``root.....`` returning ``None`` on any miss.""" - node = root - for attr in path: - node = getattr(node, attr, None) - if node is None: - return None - return node - - -def resolve_warp_twin(name: str, modules: Iterable[Any]) -> Any | None: - """Return ``name`` from ``modules`` if its ``__module__`` lives under a warp prefix.""" - for module in modules: - candidate = getattr(module, name, None) - if candidate is None: - continue - origin = getattr(candidate, "__module__", "") or "" - if origin.startswith(WARP_ROOT_PREFIXES): - return candidate - return None - - -def iter_term_attrs(group: Any) -> Iterator[tuple[str, Any]]: - """Yield ``(name, term)`` pairs from a manager-cfg group, skipping dunders.""" - if group is None: - return - for name in [n for n in dir(group) if not n.startswith("_")]: - term = getattr(group, name, None) - if term is None: - continue - yield name, term diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/torch.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/torch.py deleted file mode 100644 index 1c4054e79cda..000000000000 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/torch.py +++ /dev/null @@ -1,72 +0,0 @@ -# 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 - -"""Torch (default) frontend. - -The torch frontend is a thin wrapper around :func:`gym.make`. It does not -mutate the cfg — the standard ``isaaclab.envs.*`` env classes use -``FactoryBase`` to dispatch on the active physics backend at construction -time, so PhysX and Newton physics both run through this single path. - -Its rule set is small: one warning when the user asks for the torch -frontend on a task that's *registered* against the warp runtime -(``entry_point`` under :data:`WARP_ROOT_PREFIXES`). That's a contradiction -of intent — the env class is a warp env regardless of the frontend flag — -so we surface it but don't block. -""" - -from __future__ import annotations - -from collections.abc import Iterable -from typing import Any - -from .base import ( - Change, - CompatRule, - Frontend, - Issue, - ResolveContext, - Runtime, - Severity, - TaskMeta, -) - - -class WarnIfTaskIsWarpRegistered(CompatRule): - """Warn when a warp-registered task is asked to run on the torch frontend. - - The env class itself is a warp implementation; ``--frontend=torch`` will - still happily build it via ``gym.make`` (since the entry-point class is - what runs), but the user's stated intent says torch. We emit a warning - so the contradiction is visible. - """ - - name = "warn_if_task_is_warp_registered" - - def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: - if ctx.task.runtime == Runtime.WARP: - yield Issue( - rule=self.name, - severity=Severity.WARNING, - message=( - f"task is registered against the warp runtime" - f" (entry_point={ctx.task.entry_point!r}); --frontend=torch" - f" will still build the warp env class. Consider --frontend=warp." - ), - location="task.entry_point", - detail={"runtime": ctx.task.runtime.value}, - ) - - -class TorchFrontend(Frontend): - """Default frontend: ``gym.make`` against the registered ``entry_point``.""" - - name = "torch" - rules = (WarnIfTaskIsWarpRegistered,) - - def construct(self, cfg: Any, meta: TaskMeta, **kwargs: Any) -> Any: - import gymnasium as gym - - return gym.make(meta.task_id, cfg=cfg, **kwargs) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py deleted file mode 100644 index cb7257944406..000000000000 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend/warp.py +++ /dev/null @@ -1,513 +0,0 @@ -# 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 - -"""Warp frontend. - -Routes a stable manager-based task cfg onto :class:`ManagerBasedRLEnvWarp` -by running a pluggable :class:`CompatRule` pipeline. Direct envs encode -their runtime in the env class itself (each warp direct env is its own -class), so for direct cfgs the frontend just verifies the registered -entry-point lives under :data:`WARP_ROOT_PREFIXES` and forwards to -:func:`gym.make`. - -Rule pipeline (manager-based path): - -* :class:`CheckPhysicsIsNewton` — blocking issue if ``cfg.sim.physics`` - is PhysX-flavoured, since the warp runtime needs Newton physics. -* :class:`ResolvePhysicsPreset` — collapses ``PresetCfg`` wrappers (no-op - when Hydra already did so). -* :class:`DropUnsupportedSensors` — removes sensors warp can't run yet. -* :class:`PromoteSceneEntityCfg` — in-place class promotion to the warp - variant so warp mdp kernels see ``joint_mask`` / ``joint_ids_wp``. -* :class:`SwapMdpFunctions` — name-based ``term.func`` replacement against - the warp ``mdp`` modules. -* :class:`SwapActionClassType` — strict swap of action ``class_type`` to - the warp twin. -* :class:`VerifyDirectIsWarp` — only fires for direct cfgs; blocks if - the registered entry-point isn't a warp class. -""" - -from __future__ import annotations - -import importlib -import logging -from collections.abc import Iterable -from types import ModuleType -from typing import Any, ClassVar - -from .base import ( - WARP_ROOT_PREFIXES, - Change, - CompatRule, - Frontend, - FrontendIncompatibleError, - Issue, - ResolveContext, - Runtime, - Severity, - TaskMeta, - Workflow, - iter_term_attrs, - resolve_warp_twin, - walk_attrs, -) - -logger = logging.getLogger(__name__) - - -# --------------------------------------------------------------------------- -# Rules -# --------------------------------------------------------------------------- - - -class CheckPhysicsIsNewton(CompatRule): - """Block if ``cfg.sim.physics`` isn't Newton. - - The warp runtime resolves asset / sensor ``class_type`` via the active - physics backend's module tree (``isaaclab_newton.*`` for Newton, - ``isaaclab_physx.*`` for PhysX). Loading a PhysX class under the warp - runtime fails on ``omni.physics.tensors.api`` (a Kit module the warp - runtime doesn't initialise), so a PhysX physics cfg is fundamentally - incompatible with this frontend. - - Hydra's preset resolution + :meth:`WarpFrontend.preprocess_hydra_args` - normally collapse ``PresetCfg`` wrappers to the ``newton`` field - before this rule runs; we still check, so a programmatic caller or a - misconfigured task fails fast and loudly here instead of crashing - deep in scene init. - """ - - name = "check_physics_is_newton" - - NEWTON_MODULE_PREFIXES: ClassVar[tuple[str, ...]] = ("isaaclab_newton",) - """Module prefixes that mark a physics cfg as Newton-flavoured.""" - - PHYSX_MODULE_PREFIXES: ClassVar[tuple[str, ...]] = ("isaaclab_physx",) - """Module prefixes that mark a physics cfg as PhysX-flavoured (definitely incompatible).""" - - def applies_to(self, ctx: ResolveContext) -> bool: - return ctx.task.workflow == Workflow.MANAGER_BASED - - def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: - physics = getattr(getattr(cfg, "sim", None), "physics", None) - if physics is None: - return - cls = type(physics) - module = getattr(cls, "__module__", "") or "" - # Skip PresetCfg wrappers — they'll be unwrapped by ResolvePhysicsPreset. - if hasattr(physics, "newton") and not hasattr(physics, "class_type"): - return - if module.startswith(self.NEWTON_MODULE_PREFIXES): - return # OK - severity = Severity.BLOCKING if module.startswith(self.PHYSX_MODULE_PREFIXES) else Severity.WARNING - backend = "PhysX" if severity == Severity.BLOCKING else "unknown" - yield Issue( - rule=self.name, - severity=severity, - message=( - f"sim.physics is {cls.__name__} ({module}); the warp runtime" - f" needs a Newton physics cfg ({backend} detected). Pass" - f" `presets=newton` or use a Newton physics cfg explicitly." - ), - location="cfg.sim.physics", - detail={"physics_class": cls.__name__, "physics_module": module}, - ) - - -class ResolvePhysicsPreset(CompatRule): - """Collapse ``cfg.sim.physics`` from ``PresetCfg`` to its ``newton`` field. - - Hydra's preset resolution normally collapses the wrapper before the - frontend runs. This rule covers the residual case of programmatic use - of :meth:`Frontend.build` without Hydra, and custom physics cfgs that - still expose a ``newton`` attribute. - """ - - name = "resolve_physics_preset" - - def applies_to(self, ctx: ResolveContext) -> bool: - return ctx.task.workflow == Workflow.MANAGER_BASED - - def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: - physics = getattr(getattr(cfg, "sim", None), "physics", None) - if physics is None or not hasattr(physics, "newton"): - return - cfg.sim.physics = physics.newton - yield Change( - rule=self.name, - description=f"sim.physics → {type(physics.newton).__name__}", - ) - - -class DropUnsupportedSensors(CompatRule): - """Drop scene sensors the warp runtime can't run yet. - - Default: ``("height_scanner",)``. Pass ``sensors=...`` at construction - time to override. - """ - - name = "drop_unsupported_sensors" - - def __init__(self, sensors: Iterable[str] = ("height_scanner",)): - self.sensors = tuple(sensors) - - def applies_to(self, ctx: ResolveContext) -> bool: - return ctx.task.workflow == Workflow.MANAGER_BASED - - def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: - scene = getattr(cfg, "scene", None) - if scene is None: - return - for sensor in self.sensors: - if getattr(scene, sensor, None) is not None: - setattr(scene, sensor, None) - yield Change(rule=self.name, description=f"scene.{sensor} → None") - - -class PromoteSceneEntityCfg(CompatRule): - """Promote :class:`SceneEntityCfg` instances under term params to the warp variant. - - The warp variant adds cached ``joint_mask`` / ``joint_ids_wp`` / - ``body_ids_wp`` fields that the warp mdp kernels read after - :meth:`resolve`. The class hierarchy is asserted at apply time — if - the warp class no longer subclasses the stable class, the rule raises - :class:`FrontendIncompatibleError` so the in-place ``__class__`` - promotion can never silently corrupt instances. - """ - - name = "promote_scene_entity_cfg" - - GROUPS: ClassVar[tuple[tuple[str, ...], ...]] = ( - ("observations", "policy"), - ("events",), - ("rewards",), - ("terminations",), - ("commands",), - ("curriculum",), - ("actions",), - ) - - def applies_to(self, ctx: ResolveContext) -> bool: - return ctx.task.workflow == Workflow.MANAGER_BASED - - def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: - from isaaclab.managers.scene_entity_cfg import SceneEntityCfg as Stable - - from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as Warp - - if not issubclass(Warp, Stable): - raise FrontendIncompatibleError( - f"Warp SceneEntityCfg must subclass stable SceneEntityCfg; got mro {[c.__name__ for c in Warp.__mro__]}" - ) - - promoted = 0 - for path in self.GROUPS: - group = walk_attrs(cfg, path) - if group is None: - continue - for _name, term in iter_term_attrs(group): - params = getattr(term, "params", None) - if not isinstance(params, dict): - continue - for value in params.values(): - if isinstance(value, Warp) or not isinstance(value, Stable): - continue - # ``__class__ =`` is permitted only when the layouts - # match (e.g. neither side adds ``__slots__`` over the - # other). ``issubclass`` does *not* guarantee this. If - # the runtime refuses, surface a blocking issue rather - # than crash mid-pipeline; the cfg is still safe because - # we set warp-only fields *after* the assignment. - try: - value.__class__ = Warp - except TypeError as exc: - yield Issue( - rule=self.name, - severity=Severity.BLOCKING, - message=( - f"in-place __class__ promotion to {Warp.__name__} failed:" - f" {exc}. The warp variant likely diverged in slots/layout" - f" from the stable one; rebuild cfg explicitly with the warp" - f" SceneEntityCfg." - ), - location=".".join(path), - detail={"warp_class": Warp.__name__, "stable_class": Stable.__name__}, - ) - return - value.joint_mask = None - value.joint_ids_wp = None - value.body_ids_wp = None - promoted += 1 - if promoted: - yield Change( - rule=self.name, - description=f"promoted {promoted} SceneEntityCfg instance(s) to warp variant", - ) - - -class SwapMdpFunctions(CompatRule): - """Replace ``term.func`` with the same-named callable from warp ``mdp`` modules. - - Walks observation / event / reward / termination / command / curriculum - groups. A re-export from stable code is rejected by checking - ``__module__``. Missing twins are dropped (term set to ``None``) and - surfaced as :attr:`Severity.WARNING`; in strict mode they're escalated - to :attr:`Severity.BLOCKING`. - """ - - name = "swap_mdp_functions" - - GROUPS: ClassVar[tuple[tuple[str, ...], ...]] = ( - ("observations", "policy"), - ("events",), - ("rewards",), - ("terminations",), - ("commands",), - ("curriculum",), - ) - - def __init__( - self, - stable_root: str = "isaaclab_tasks", - warp_root: str = "isaaclab_tasks_experimental", - fallback_mdp: str = "isaaclab_experimental.envs.mdp", - ): - self.stable_root = stable_root - self.warp_root = warp_root - self.fallback_mdp = fallback_mdp - - def applies_to(self, ctx: ResolveContext) -> bool: - return ctx.task.workflow == Workflow.MANAGER_BASED - - def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: - modules = self._mdp_modules(ctx.task) - logger.info("WarpFrontend: warp mdp modules → %s", [m.__name__ for m in modules]) - searched = tuple(m.__name__ for m in modules) - for path in self.GROUPS: - group = walk_attrs(cfg, path) - if group is None: - continue - location_prefix = ".".join(path) - for name, term in iter_term_attrs(group): - if not hasattr(term, "func"): - continue - stable = term.func - if not callable(stable): - continue - # Idempotency: if the term already references a warp-native - # callable (e.g. running the bridge against an already-warp - # task), don't try to swap. Without this guard the rule would - # ``resolve_warp_twin`` against modules that might not contain - # the same symbol and silently drop the term. - origin = getattr(stable, "__module__", "") or "" - if origin.startswith(WARP_ROOT_PREFIXES): - continue - twin = resolve_warp_twin(stable.__name__, modules) - if twin is not None: - term.func = twin - continue - # No twin found: drop the term and report. - setattr(group, name, None) - yield Issue( - rule=self.name, - severity=Severity.BLOCKING if ctx.strict else Severity.WARNING, - message=(f"no warp twin for stable func={stable.__name__!r}; term dropped."), - location=f"{location_prefix}.{name}", - detail={"expected": stable.__name__, "searched": list(searched)}, - ) - - def _mdp_modules(self, task: TaskMeta) -> list[ModuleType]: - """Locate warp mdp modules that mirror the task's stable mdp.""" - modules: list[ModuleType] = [] - entry = task.env_cfg_entry_point - # Use a trailing dot when matching the stable root so a task registered - # under ``isaaclab_tasks_experimental.*`` (the warp side) doesn't match - # ``isaaclab_tasks`` and end up with a double-replaced - # ``isaaclab_tasks_experimental_experimental.*`` import path. - stable_prefix = self.stable_root + "." - if isinstance(entry, str) and entry.startswith(stable_prefix): - warp_pkg = entry.rsplit(".", 1)[0].replace(self.stable_root, self.warp_root, 1) - parts = warp_pkg.split(".") - for depth in range(len(parts), 0, -1): - target = ".".join(parts[:depth] + ["mdp"]) - try: - modules.append(importlib.import_module(target)) - break - except ModuleNotFoundError as exc: - # Only swallow "this candidate module does not exist". - # An ImportError raised *inside* a module that does exist - # is a real bug we should surface, not paper over by - # falling back. - if exc.name == target: - continue - raise - try: - modules.append(importlib.import_module(self.fallback_mdp)) - except ModuleNotFoundError as exc: - if exc.name != self.fallback_mdp: - raise - logger.warning("WarpFrontend: fallback mdp module %r not importable", self.fallback_mdp) - return modules - - -class SwapActionClassType(CompatRule): - """Swap ``cfg.actions..class_type`` with the same-named warp class. - - Always strict at the term level: an action with no warp twin can't run - on the warp runtime, so a missing twin is :attr:`Severity.BLOCKING` - regardless of ``ctx.strict``. - """ - - name = "swap_action_class_type" - - def __init__(self, warp_actions_module: str = "isaaclab_experimental.envs.mdp.actions.joint_actions"): - self.warp_actions_module = warp_actions_module - - def applies_to(self, ctx: ResolveContext) -> bool: - return ctx.task.workflow == Workflow.MANAGER_BASED - - def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: - actions = getattr(cfg, "actions", None) - if actions is None: - return - try: - module = importlib.import_module(self.warp_actions_module) - except ImportError as exc: - yield Issue( - rule=self.name, - severity=Severity.BLOCKING, - message=(f"warp action module {self.warp_actions_module!r} not importable: {exc}"), - location="actions", - ) - return - searched = (module.__name__,) - for name, term in iter_term_attrs(actions): - if not hasattr(term, "class_type"): - continue - stable = term.class_type - if not callable(stable): - continue - twin = resolve_warp_twin(stable.__name__, [module]) - if twin is not None: - term.class_type = twin - continue - yield Issue( - rule=self.name, - severity=Severity.BLOCKING, - message=f"no warp twin for stable class_type={stable.__name__!r}", - location=f"actions.{name}", - detail={"expected": stable.__name__, "searched": list(searched)}, - ) - - -class VerifyDirectIsWarp(CompatRule): - """For direct envs, block if the registered entry-point isn't a warp class. - - Direct envs encode their runtime in the env class itself (each warp - direct env is a separate class with its own kernels), so the warp - frontend can't *adapt* a stable direct cfg — it can only refuse it - with a clear message. A direct task whose entry-point already lives - under :data:`WARP_ROOT_PREFIXES` is a pass-through. - """ - - name = "verify_direct_is_warp" - - def applies_to(self, ctx: ResolveContext) -> bool: - return ctx.task.workflow == Workflow.DIRECT - - def run(self, cfg: Any, ctx: ResolveContext) -> Iterable[Issue | Change]: - if ctx.task.runtime == Runtime.WARP: - return - ep = ctx.task.entry_point - yield Issue( - rule=self.name, - severity=Severity.BLOCKING, - message=( - f"direct env entry_point {ep!r} is not a warp implementation." - f" --frontend=warp on a direct task requires the registered class to live under" - f" {list(WARP_ROOT_PREFIXES)} (e.g. *-Direct-Warp-v0)." - ), - location="task.entry_point", - detail={"entry_point": ep}, - ) - - -# --------------------------------------------------------------------------- -# WarpFrontend -# --------------------------------------------------------------------------- - - -class WarpFrontend(Frontend): - """Run a stable env cfg on the warp runtime. - - Manager-based cfgs get the full rule pipeline and are constructed on - :class:`ManagerBasedRLEnvWarp`. Direct cfgs run :class:`VerifyDirectIsWarp` - only and dispatch through :func:`gym.make`. - """ - - name = "warp" - rules = ( - CheckPhysicsIsNewton, - ResolvePhysicsPreset, - DropUnsupportedSensors, - PromoteSceneEntityCfg, - SwapMdpFunctions, - SwapActionClassType, - VerifyDirectIsWarp, - ) - - # -- construction --------------------------------------------------------- - - def construct(self, cfg: Any, meta: TaskMeta, **kwargs: Any) -> Any: - if meta.workflow == Workflow.MANAGER_BASED: - from isaaclab_experimental.envs import ManagerBasedRLEnvWarp - - return ManagerBasedRLEnvWarp(cfg=cfg, **kwargs) - if meta.workflow == Workflow.DIRECT: - import gymnasium as gym - - return gym.make(meta.task_id, cfg=cfg, **kwargs) - raise FrontendIncompatibleError( - f"WarpFrontend does not support workflow {meta.workflow.value!r} (task {meta.task_id!r})." - ) - - # -- CLI hook ------------------------------------------------------------- - - def preprocess_hydra_args(self, task_id: str, args: list[str]) -> list[str]: - """Inject ``presets=newton`` for stable manager-based tasks. - - Hydra resolves ``presets=...`` *before* the frontend runs, so the - injection has to happen at the CLI layer. We only inject for tasks - whose ``env_cfg_entry_point`` is under ``isaaclab_tasks.manager_based``; - other tasks (direct, pre-warp ``*-Warp-v0``) don't carry a preset - system and would error out. - - If the user already passed an explicit ``presets=`` we don't override, - but we warn when their preset isn't ``newton``. - """ - import gymnasium as gym - - try: - spec = gym.spec(task_id) - except gym.error.NameNotFound: - return args - cfg_ep = spec.kwargs.get("env_cfg_entry_point") - if not isinstance(cfg_ep, str) or not cfg_ep.startswith("isaaclab_tasks.manager_based"): - return args - - # Match both ``presets=foo`` and ``--presets=foo`` — Hydra accepts either. - def _is_preset_arg(arg: str) -> bool: - return arg.lstrip("-").startswith("presets=") - - explicit = next((a for a in args if _is_preset_arg(a)), None) - if explicit is None: - return [*args, "presets=newton"] - if explicit.lstrip("-") != "presets=newton": - logger.warning( - "--frontend=warp on %r expects presets=newton; got %r — adapter may fail.", - task_id, - explicit, - ) - return args diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py b/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py index 1930f6ce1cb7..d92fe54b1198 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py @@ -34,6 +34,30 @@ class SceneEntityCfg(_SceneEntityCfg): """Integer indices of selected bodies — used for subset-sized body gathers.""" body_ids_wp: wp.array | None = None + @classmethod + def from_stable(cls, stable: _SceneEntityCfg) -> SceneEntityCfg: + """Build a warp scene-entity cfg from a stable one. + + Copies every selection field declared on the stable cfg (name, joint + names/ids, body names/ids, fixed-tendon names/ids, object-collection + names/ids, ``preserve_order``); warp-specific fields stay ``None`` + and are filled by :meth:`resolve` at scene build time. Used by the + warp frontend to promote :class:`SceneEntityCfg` instances embedded + in term ``params`` without resorting to a ``__class__`` reassign. + """ + return cls( + name=stable.name, + joint_names=stable.joint_names, + joint_ids=stable.joint_ids, + fixed_tendon_names=stable.fixed_tendon_names, + fixed_tendon_ids=stable.fixed_tendon_ids, + body_names=stable.body_names, + body_ids=stable.body_ids, + object_collection_names=stable.object_collection_names, + object_collection_ids=stable.object_collection_ids, + preserve_order=stable.preserve_order, + ) + def resolve(self, scene: InteractiveScene): # run the stable resolution first (fills joint_ids/body_ids from names/regex) super().resolve(scene) diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py new file mode 100644 index 000000000000..a35c21dd07ed --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -0,0 +1,417 @@ +# 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 + +"""Tests for :mod:`isaaclab_experimental.envs.frontend`. + +Pure-Python unit tests; no app launch. Covers: + +* :class:`Frontend` / :class:`Workflow` enum surface. +* :func:`SceneEntityCfg.from_stable` field copy. +* :func:`_require_newton_physics` hard-check. +* :func:`_promote_scene_entity_cfgs` walks ``term.params`` dicts. +* :func:`_swap_mdp` swaps ``func`` *and* ``class_type``; raises with a path + list when twins are missing. +* :func:`_resolve_warp_twin` rejects stable-origin re-exports. +* :func:`_require_direct_is_warp_task` accepts warp-rooted entry points + and rejects stable ones. +""" + +from __future__ import annotations + +import types +import unittest +from typing import Any + +import gymnasium as gym +from isaaclab_experimental.envs.frontend import ( + Frontend, + FrontendIncompatibleError, + Workflow, + _is_swap_candidate, + _iter_term_attrs, + _promote_scene_entity_cfgs, + _require_direct_is_warp_task, + _require_newton_physics, + _resolve_warp_twin, + _swap_mdp, + _walk_attrs, +) +from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as WarpSceneEntityCfg +from isaaclab_newton.physics import NewtonCfg +from isaaclab_physx.physics import PhysxCfg + +from isaaclab.managers.scene_entity_cfg import SceneEntityCfg as StableSceneEntityCfg + +# ====================================================================== +# Enums +# ====================================================================== + + +class TestEnums(unittest.TestCase): + def test_frontend_values(self): + self.assertEqual(Frontend.TORCH, "torch") + self.assertEqual(Frontend.WARP, "warp") + + def test_frontend_coercion(self): + self.assertIs(Frontend("torch"), Frontend.TORCH) + self.assertIs(Frontend("warp"), Frontend.WARP) + with self.assertRaises(ValueError): + Frontend("kit") + + def test_workflow_values(self): + self.assertEqual(Workflow.MANAGER_BASED, "manager_based") + self.assertEqual(Workflow.DIRECT, "direct") + + +# ====================================================================== +# SceneEntityCfg.from_stable +# ====================================================================== + + +class TestFromStable(unittest.TestCase): + def test_copies_minimum_fields(self): + stable = StableSceneEntityCfg(name="robot") + warp = WarpSceneEntityCfg.from_stable(stable) + self.assertIsInstance(warp, WarpSceneEntityCfg) + self.assertEqual(warp.name, "robot") + # Warp-only fields stay None until :meth:`resolve` runs. + self.assertIsNone(warp.joint_mask) + self.assertIsNone(warp.joint_ids_wp) + self.assertIsNone(warp.body_ids_wp) + + def test_copies_all_selection_fields(self): + stable = StableSceneEntityCfg( + name="robot", + joint_names=["lf_hip"], + joint_ids=[0, 1, 2], + fixed_tendon_names=["tendon_a"], + fixed_tendon_ids=[5], + body_names=["base", "lf_foot"], + body_ids=[0, 4], + object_collection_names=["objs"], + object_collection_ids=[7], + preserve_order=True, + ) + warp = WarpSceneEntityCfg.from_stable(stable) + for field in ( + "name", + "joint_names", + "joint_ids", + "fixed_tendon_names", + "fixed_tendon_ids", + "body_names", + "body_ids", + "object_collection_names", + "object_collection_ids", + "preserve_order", + ): + self.assertEqual(getattr(warp, field), getattr(stable, field), msg=f"field {field!r} mismatch") + + +# ====================================================================== +# _require_newton_physics +# ====================================================================== + + +class TestRequireNewtonPhysics(unittest.TestCase): + def _cfg_with(self, physics: Any) -> Any: + cfg = types.SimpleNamespace() + cfg.sim = types.SimpleNamespace(physics=physics) + return cfg + + def test_passes_for_newton(self): + cfg = self._cfg_with(NewtonCfg()) + _require_newton_physics(cfg, "Isaac-Test-v0") # no raise + + def test_rejects_physx(self): + cfg = self._cfg_with(PhysxCfg()) + with self.assertRaises(FrontendIncompatibleError) as exc: + _require_newton_physics(cfg, "Isaac-Test-v0") + self.assertIn("presets=newton", str(exc.exception)) + self.assertIn("PhysxCfg", str(exc.exception)) + + def test_rejects_none(self): + cfg = self._cfg_with(None) + with self.assertRaises(FrontendIncompatibleError): + _require_newton_physics(cfg, "Isaac-Test-v0") + + +# ====================================================================== +# _promote_scene_entity_cfgs +# ====================================================================== + + +class _Term: + """Minimal stand-in for a manager term cfg.""" + + def __init__(self, params: dict | None = None, func=None, class_type=None): + self.params = params if params is not None else {} + if func is not None: + self.func = func + if class_type is not None: + self.class_type = class_type + + +class _Group: + """Minimal stand-in for a manager-cfg group (e.g. RewardsCfg).""" + + def __init__(self, **terms: Any): + for name, term in terms.items(): + setattr(self, name, term) + + +class _ObsCfg: + def __init__(self, policy: _Group): + self.policy = policy + + +def _make_cfg(**groups: Any) -> Any: + cfg = types.SimpleNamespace() + for name, value in groups.items(): + setattr(cfg, name, value) + return cfg + + +class TestPromoteSceneEntityCfgs(unittest.TestCase): + def test_promotes_in_params(self): + stable = StableSceneEntityCfg(name="robot", joint_names=["lf_hip"]) + term = _Term(params={"asset_cfg": stable, "scale": 1.0}) + cfg = _make_cfg(rewards=_Group(track=term)) + _promote_scene_entity_cfgs(cfg) + promoted = term.params["asset_cfg"] + self.assertIsInstance(promoted, WarpSceneEntityCfg) + self.assertEqual(promoted.name, "robot") + self.assertEqual(promoted.joint_names, ["lf_hip"]) + # Non-SceneEntityCfg params are untouched. + self.assertEqual(term.params["scale"], 1.0) + + def test_skips_already_warp(self): + warp = WarpSceneEntityCfg(name="robot") + term = _Term(params={"asset_cfg": warp}) + cfg = _make_cfg(rewards=_Group(t=term)) + _promote_scene_entity_cfgs(cfg) + self.assertIs(term.params["asset_cfg"], warp) # unchanged identity + + def test_walks_all_term_paths(self): + # Drop a stable cfg in every supported group; all should promote. + groups = {} + terms: list[_Term] = [] + for group_name in ("rewards", "events", "terminations", "commands", "curriculum"): + t = _Term(params={"asset_cfg": StableSceneEntityCfg(name=group_name)}) + terms.append(t) + groups[group_name] = _Group(t=t) + # observations is nested as observations.policy + obs_term = _Term(params={"asset_cfg": StableSceneEntityCfg(name="obs")}) + terms.append(obs_term) + groups["observations"] = _ObsCfg(policy=_Group(t=obs_term)) + # actions group also walked + act_term = _Term(params={"asset_cfg": StableSceneEntityCfg(name="act")}) + terms.append(act_term) + groups["actions"] = _Group(t=act_term) + + cfg = _make_cfg(**groups) + _promote_scene_entity_cfgs(cfg) + for t in terms: + self.assertIsInstance(t.params["asset_cfg"], WarpSceneEntityCfg) + + +# ====================================================================== +# _swap_mdp +# ====================================================================== + + +# A fake "stable" mdp module the test fixture stands in for +# ``isaaclab_tasks..mdp`` — symbols here pretend to live there. +def _stable_func(env, **params): + return None + + +class _StableActionCls: + pass + + +_stable_func.__module__ = "isaaclab_tasks.fake_task.mdp" +_StableActionCls.__module__ = "isaaclab_tasks.fake_task.mdp" + + +# And a fake warp twin module the test installs into sys.modules under the +# expected name resolution path. Symbols here pretend to live under +# ``isaaclab_experimental.envs.mdp`` (the fallback). +def _warp_twin_func(env, out, **params): + return None + + +class _WarpActionCls: + pass + + +_warp_twin_func.__module__ = "isaaclab_experimental.envs.mdp" +_WarpActionCls.__module__ = "isaaclab_experimental.envs.mdp" + + +class _FakeMdpModule: + """Stand-in for a warp mdp module containing twins.""" + + __name__ = "test_fake_warp_mdp" + + +class TestSwapMdp(unittest.TestCase): + def _patch_modules(self, name_to_symbol: dict[str, Any]) -> _FakeMdpModule: + m = _FakeMdpModule() + for name, sym in name_to_symbol.items(): + setattr(m, name, sym) + return m + + def _patched_warp_mdp_modules(self, modules: list[Any]): + # Patch _warp_mdp_modules at the frontend module level so _swap_mdp + # uses our fake module list instead of doing real importlib lookups. + import isaaclab_experimental.envs.frontend as fe + + self._orig = fe._warp_mdp_modules + fe._warp_mdp_modules = lambda task_id: modules # type: ignore[assignment] + + def setUp(self) -> None: + self._orig = None + + def tearDown(self) -> None: + if self._orig is not None: + import isaaclab_experimental.envs.frontend as fe + + fe._warp_mdp_modules = self._orig + + def test_swaps_func_and_class_type(self): + fake = self._patch_modules({"_stable_func": _warp_twin_func, "_StableActionCls": _WarpActionCls}) + self._patched_warp_mdp_modules([fake]) + term_reward = _Term(func=_stable_func) + term_action = _Term(class_type=_StableActionCls) + cfg = _make_cfg(rewards=_Group(r=term_reward), actions=_Group(a=term_action)) + _swap_mdp(cfg, "Isaac-Test-v0") + self.assertIs(term_reward.func, _warp_twin_func) + self.assertIs(term_action.class_type, _WarpActionCls) + + def test_missing_twin_raises_with_path_list(self): + # No twins available — every term should be reported as missing. + fake = self._patch_modules({}) + self._patched_warp_mdp_modules([fake]) + term = _Term(func=_stable_func) + cfg = _make_cfg(rewards=_Group(track=term)) + with self.assertRaises(FrontendIncompatibleError) as exc: + _swap_mdp(cfg, "Isaac-Test-v0") + msg = str(exc.exception) + self.assertIn("rewards.track.func", msg) + self.assertIn("_stable_func", msg) + # The cfg term wasn't mutated (we hard-failed before partial application). + self.assertIs(term.func, _stable_func) + + def test_skips_already_warp(self): + fake = self._patch_modules({}) + self._patched_warp_mdp_modules([fake]) + term = _Term(func=_warp_twin_func) # already warp-origin + cfg = _make_cfg(rewards=_Group(r=term)) + _swap_mdp(cfg, "Isaac-Test-v0") # no raise + self.assertIs(term.func, _warp_twin_func) + + +# ====================================================================== +# Twin resolution +# ====================================================================== + + +class TestResolveWarpTwin(unittest.TestCase): + def test_accepts_warp_origin(self): + m = types.SimpleNamespace() + m.foo = _warp_twin_func + result = _resolve_warp_twin("foo", [m]) + self.assertIs(result, _warp_twin_func) + + def test_rejects_stable_origin(self): + m = types.SimpleNamespace() + m.foo = _stable_func # same name, stable origin + self.assertIsNone(_resolve_warp_twin("foo", [m])) + + def test_returns_none_when_absent(self): + m = types.SimpleNamespace() + self.assertIsNone(_resolve_warp_twin("missing", [m])) + + +# ====================================================================== +# Swap candidate heuristic +# ====================================================================== + + +class TestIsSwapCandidate(unittest.TestCase): + def test_stable_callable_is_candidate(self): + self.assertTrue(_is_swap_candidate(_stable_func)) + + def test_warp_callable_is_not(self): + self.assertFalse(_is_swap_candidate(_warp_twin_func)) + + def test_non_callable_is_not(self): + self.assertFalse(_is_swap_candidate(42)) + self.assertFalse(_is_swap_candidate("string")) + + +# ====================================================================== +# Direct workflow guard +# ====================================================================== + + +_DIRECT_TEST_TASKS = { + "_Frontend-Test-Warp-Direct-v0": ("isaaclab_tasks_experimental.fake:DirectEnv", True), + "_Frontend-Test-Stable-Direct-v0": ("isaaclab_tasks.fake:DirectEnv", False), +} + + +class TestRequireDirectIsWarpTask(unittest.TestCase): + @classmethod + def setUpClass(cls): + # Register stub tasks so ``gym.spec`` resolves them. Entry-point strings + # are never invoked (the guard only inspects them), so a fake import + # path is fine. + cls._registered: list[str] = [] + for task_id, (ep, _) in _DIRECT_TEST_TASKS.items(): + try: + gym.register(id=task_id, entry_point=ep, disable_env_checker=True) + cls._registered.append(task_id) + except gym.error.Error: + # Already registered from a previous run — that's fine. + pass + + def test_accepts_warp_rooted(self): + _require_direct_is_warp_task("_Frontend-Test-Warp-Direct-v0") # no raise + + def test_rejects_stable_rooted(self): + with self.assertRaises(FrontendIncompatibleError) as exc: + _require_direct_is_warp_task("_Frontend-Test-Stable-Direct-v0") + self.assertIn("isaaclab_experimental", str(exc.exception)) + + def test_rejects_unknown_task(self): + with self.assertRaises(FrontendIncompatibleError): + _require_direct_is_warp_task("Frontend-Test-NotRegistered-v0") + + +# ====================================================================== +# Walk helpers +# ====================================================================== + + +class TestWalkHelpers(unittest.TestCase): + def test_walk_attrs_hit(self): + root = types.SimpleNamespace(a=types.SimpleNamespace(b=types.SimpleNamespace(c=42))) + self.assertEqual(_walk_attrs(root, ("a", "b", "c")), 42) + + def test_walk_attrs_miss(self): + root = types.SimpleNamespace(a=types.SimpleNamespace()) + self.assertIsNone(_walk_attrs(root, ("a", "missing"))) + self.assertIsNone(_walk_attrs(root, ("missing", "b"))) + + def test_iter_term_attrs_skips_dunders_and_none(self): + g = types.SimpleNamespace(t1=_Term(), t2=None, _internal=_Term()) + names = sorted(n for n, _ in _iter_term_attrs(g)) + self.assertEqual(names, ["t1"]) + + +if __name__ == "__main__": + unittest.main() From bd27af41f81d175e40373924517d81c330b7b22c Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 26 May 2026 06:51:58 +0000 Subject: [PATCH 11/58] Simplify frontend cfg walker and unify term-cfg type contract Frontend (envs/frontend.py) - Replace _TERM_PATHS hardcode (with its policy-only observation limitation) with _walk_terms: a recursive ManagerTermBaseCfg discovery that descends any configclass and yields each term with its path. New cfg layouts and observation sub-groups (perception, critic, camera_images, ...) are picked up automatically with no framework change. - _promote_scene_entity_cfgs and _swap_mdp consume the new walker; delete _walk_attrs / _iter_term_attrs / _TERM_PATHS. - Per-promotion log now lists the actual SceneEntityCfg paths. - Rename _require_direct_is_warp_task to _assert_direct_warp_registration. Update _detect_workflow with a note for adding new cfg roots. Manager term cfg contract - Relax stable ObservationTermCfg / RewardTermCfg / TerminationTermCfg func annotation to Callable[..., torch.Tensor | None] so the warp func(env, out) -> None contract type-checks alongside the existing torch return-value contract. - Drop the now-redundant warp-side manager_term_cfg.py shim and redirect the 9 relative imports onto isaaclab.managers.manager_term_cfg. Tests - New TestWalkTerms cases verify the recursive type-driven discovery. - TestPromoteSceneEntityCfgs / TestSwapMdp use real term cfgs and configclass fixtures so they exercise the actual walker contract. --- .../reinforcement_learning/rsl_rl/train.py | 8 +- .../isaaclab/managers/manager_term_cfg.py | 6 +- .../isaaclab_experimental/envs/frontend.py | 156 +++++---- .../managers/__init__.py | 1 - .../managers/action_manager.py | 2 +- .../managers/command_manager.py | 3 +- .../managers/event_manager.py | 3 +- .../managers/manager_base.py | 2 +- .../managers/manager_term_cfg.py | 25 -- .../managers/observation_manager.py | 2 +- .../managers/reward_manager.py | 3 +- .../managers/termination_manager.py | 3 +- .../test/envs/test_frontend.py | 303 ++++++++++-------- .../classic/humanoid/mdp/rewards.py | 3 +- 14 files changed, 278 insertions(+), 242 deletions(-) delete mode 100644 source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py diff --git a/scripts/reinforcement_learning/rsl_rl/train.py b/scripts/reinforcement_learning/rsl_rl/train.py index 23bff4218629..1f30de33eb12 100644 --- a/scripts/reinforcement_learning/rsl_rl/train.py +++ b/scripts/reinforcement_learning/rsl_rl/train.py @@ -218,12 +218,8 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen # set the log directory for the environment (works for all environment types) env_cfg.log_dir = log_dir - # Build the env via the selected frontend. The warp frontend mutates env_cfg - # in place (SceneEntityCfg promotion + MDP func/class_type swap) and then - # constructs ``ManagerBasedRLEnvWarp``; missing warp twins are a hard failure. - # When ``isaaclab_experimental`` isn't installed, the argparse step above - # has already constrained --frontend to ``torch``, so the plain ``gym.make`` - # fallback is safe. + # Build the env via the selected frontend. Warp adapts env_cfg in place; + # missing warp twins are a hard failure. render_mode = "rgb_array" if args_cli.video else None if _frontend_build is not None: env = _frontend_build(args_cli.frontend, env_cfg, args_cli.task, render_mode=render_mode) diff --git a/source/isaaclab/isaaclab/managers/manager_term_cfg.py b/source/isaaclab/isaaclab/managers/manager_term_cfg.py index aa118f8e9784..9372fbd64662 100644 --- a/source/isaaclab/isaaclab/managers/manager_term_cfg.py +++ b/source/isaaclab/isaaclab/managers/manager_term_cfg.py @@ -152,7 +152,7 @@ class CurriculumTermCfg(ManagerTermBaseCfg): class ObservationTermCfg(ManagerTermBaseCfg): """Configuration for an observation term.""" - func: Callable[..., torch.Tensor] = MISSING + func: Callable[..., torch.Tensor | None] = MISSING """The name of the function to be called. This function should take the environment object and any other parameters @@ -316,7 +316,7 @@ class EventTermCfg(ManagerTermBaseCfg): class RewardTermCfg(ManagerTermBaseCfg): """Configuration for a reward term.""" - func: Callable[..., torch.Tensor] = MISSING + func: Callable[..., torch.Tensor | None] = MISSING """The name of the function to be called. This function should take the environment object and any other parameters @@ -344,7 +344,7 @@ class RewardTermCfg(ManagerTermBaseCfg): class TerminationTermCfg(ManagerTermBaseCfg): """Configuration for a termination term.""" - func: Callable[..., torch.Tensor] = MISSING + func: Callable[..., torch.Tensor | None] = MISSING """The name of the function to be called. This function should take the environment object and any other parameters diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index 4dfc312d6481..f8b4acd2284a 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -75,18 +75,6 @@ class FrontendIncompatibleError(RuntimeError): # the warp packages. Used by the direct-workflow guard. _WARP_ROOT_PREFIXES: tuple[str, ...] = ("isaaclab_experimental", "isaaclab_tasks_experimental") -# Term containers walked by the cfg-adaptation step. Actions are included so -# the swap logic handles ``class_type`` and ``func`` in one pass. -_TERM_PATHS: tuple[tuple[str, ...], ...] = ( - ("observations", "policy"), - ("events",), - ("rewards",), - ("terminations",), - ("commands",), - ("curriculum",), - ("actions",), -) - def build( frontend: Frontend | str, @@ -119,7 +107,7 @@ def build( if workflow is Workflow.DIRECT: # Direct workflows aren't adapted — they must already be registered # under the warp packages (e.g. ``Isaac-Cartpole-Direct-Warp-v0``). - _require_direct_is_warp_task(task_id) + _assert_direct_warp_registration(task_id) return gym.make(task_id, cfg=env_cfg, **construct_kwargs) _adapt_cfg_for_warp(env_cfg, task_id) @@ -148,11 +136,12 @@ def _adapt_cfg_for_warp(cfg: Any, task_id: str) -> None: :class:`~isaaclab.managers.SceneEntityCfg` instances under each term's ``params`` with the warp variant (which adds warp-cached ``joint_mask``, ``joint_ids_wp``, ``body_ids_wp`` fields). - 3. :func:`_swap_mdp` — for every MDP term in every group (observations, - events, rewards, terminations, commands, curriculum, actions), replace - any stable ``func`` *or* ``class_type`` with its same-named warp twin. - A missing twin raises :class:`FrontendIncompatibleError` — partial - coverage is unsafe under the warp managers' kernel-only signature. + 3. :func:`_swap_mdp` — for every MDP term found anywhere in the cfg tree + (discovered by :func:`_walk_terms` via :class:`ManagerTermBaseCfg` + subclassing, not by hard-coded attribute names), replace any stable + ``func`` *or* ``class_type`` with its same-named warp twin. A missing + twin raises :class:`FrontendIncompatibleError` — partial coverage is + unsafe under the warp managers' kernel-only signature. """ _require_newton_physics(cfg, task_id) _promote_scene_entity_cfgs(cfg) @@ -183,8 +172,8 @@ def _require_newton_physics(cfg: Any, task_id: str) -> None: def _promote_scene_entity_cfgs(cfg: Any) -> None: """Replace stable :class:`SceneEntityCfg` instances with the warp variant. - Walks every ``term.params: dict`` under each term group and rebuilds any - stable :class:`SceneEntityCfg` value via + Iterates every term cfg in the tree (via :func:`_walk_terms`) and rebuilds + any stable :class:`SceneEntityCfg` value under ``term.params`` through :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable`. The warp variant subclasses the stable one, so type checks elsewhere stay valid; the new fields (``joint_mask`` / ``joint_ids_wp`` / ``body_ids_wp``) @@ -192,35 +181,42 @@ def _promote_scene_entity_cfgs(cfg: Any) -> None: """ from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as _WarpSceneEntityCfg - promoted = 0 - for path in _TERM_PATHS: - group = _walk_attrs(cfg, path) - if group is None: + promoted: list[str] = [] + for path, term in _walk_terms(cfg): + params = getattr(term, "params", None) + if not isinstance(params, dict): continue - for _name, term in _iter_term_attrs(group): - params = getattr(term, "params", None) - if not isinstance(params, dict): + for key, value in list(params.items()): + if isinstance(value, _WarpSceneEntityCfg) or not isinstance(value, _StableSceneEntityCfg): continue - for key, value in list(params.items()): - if isinstance(value, _WarpSceneEntityCfg) or not isinstance(value, _StableSceneEntityCfg): - continue - params[key] = _WarpSceneEntityCfg.from_stable(value) - promoted += 1 + params[key] = _WarpSceneEntityCfg.from_stable(value) + promoted.append(f"{'.'.join(path)}.params[{key!r}] ({value.name!r})") if promoted: - logger.info("frontend.warp: promoted %d SceneEntityCfg instance(s) to warp variant", promoted) + logger.info( + "frontend.warp: promoted %d SceneEntityCfg instance(s) to warp variant:\n %s", + len(promoted), + "\n ".join(promoted), + ) def _swap_mdp(cfg: Any, task_id: str) -> None: """Replace ``term.func`` and ``term.class_type`` with their warp twins. - Walks the same term paths used by :func:`_promote_scene_entity_cfgs` and - on each term swaps whichever of ``func`` / ``class_type`` is set to a - stable-origin symbol. Twin lookup is name-based against the task's - matching ``isaaclab_tasks_experimental.<...>.mdp`` module and, as a - fallback, :mod:`isaaclab_experimental.envs.mdp`. Any missing twin raises + Iterates every term cfg in the tree (via :func:`_walk_terms`) and on each + term swaps whichever of ``func`` / ``class_type`` is a stable-origin + callable. Twin lookup is name-based against the task's matching + ``isaaclab_tasks_experimental.<...>.mdp`` module and, as a fallback, + :mod:`isaaclab_experimental.envs.mdp`. Any missing twin raises :class:`FrontendIncompatibleError` listing every affected term — partial swaps would leave torch-style callables in the cfg and the warp managers would call them with the wrong signature. + + The warp-side side declarations (``out_dim``, ``axes``, ``observation_type``) + that the warp managers need at init are *not* supplied by this swap; they + travel with the warp twin function itself via its own + ``@generic_io_descriptor_warp(out_dim=…)`` decorator. This function only + substitutes the callable; the manager reads the new func's annotations + when it parses the term cfg. """ modules = _warp_mdp_modules(task_id) searched = tuple(m.__name__ for m in modules) @@ -228,22 +224,18 @@ def _swap_mdp(cfg: Any, task_id: str) -> None: swapped = 0 missing: list[tuple[str, str, str]] = [] # (location, attr, symbol) - for path in _TERM_PATHS: - group = _walk_attrs(cfg, path) - if group is None: - continue - location_prefix = ".".join(path) - for name, term in _iter_term_attrs(group): - for attr in ("func", "class_type"): - stable = getattr(term, attr, None) - if stable is None or not _is_swap_candidate(stable): - continue - twin = _resolve_warp_twin(stable.__name__, modules) - if twin is None: - missing.append((f"{location_prefix}.{name}", attr, stable.__name__)) - continue - setattr(term, attr, twin) - swapped += 1 + for path, term in _walk_terms(cfg): + location = ".".join(path) + for attr in ("func", "class_type"): + stable = getattr(term, attr, None) + if stable is None or not _is_swap_candidate(stable): + continue + twin = _resolve_warp_twin(stable.__name__, modules) + if twin is None: + missing.append((location, attr, stable.__name__)) + continue + setattr(term, attr, twin) + swapped += 1 if missing: lines = "\n ".join(f"{loc}.{attr}: no warp twin for {sym!r}" for loc, attr, sym in missing) @@ -260,7 +252,13 @@ def _swap_mdp(cfg: Any, task_id: str) -> None: def _detect_workflow(cfg: Any) -> Workflow: - """Classify the env cfg into manager-based or direct (used to pick build path).""" + """Classify the env cfg into manager-based or direct (used to pick build path). + + Note: + The four env cfg roots (ManagerBasedEnvCfg, ManagerBasedRLEnvCfg, + DirectRLEnvCfg, DirectMARLEnvCfg) do not share a common base class. + When a new cfg root is added, extend the isinstance tuples below. + """ if isinstance(cfg, ManagerBasedRLEnvCfg): return Workflow.MANAGER_BASED if isinstance(cfg, (DirectRLEnvCfg, DirectMARLEnvCfg)): @@ -271,7 +269,7 @@ def _detect_workflow(cfg: Any) -> Workflow: ) -def _require_direct_is_warp_task(task_id: str) -> None: +def _assert_direct_warp_registration(task_id: str) -> None: """For direct workflows, the task must be pre-registered under the warp packages.""" try: spec = gym.spec(task_id) @@ -352,22 +350,38 @@ def _is_swap_candidate(value: Any) -> bool: return True -def _walk_attrs(root: Any, path: tuple[str, ...]) -> Any: - """Walk ``root..…``; return ``None`` on any miss.""" - node = root - for attr in path: - node = getattr(node, attr, None) - if node is None: - return None - return node +def _walk_terms(node: Any, path: tuple[str, ...] = ()) -> Iterator[tuple[tuple[str, ...], Any]]: + """Yield ``(path, term)`` for every :class:`ManagerTermBaseCfg` in the cfg tree. + + Behavior at each node: + + * Match (:class:`ManagerTermBaseCfg` instance): yield ``(path, node)`` and stop — + do not descend into ``term.params`` / ``term.func``. + * Configclass: don't yield; recurse into every non-underscore attribute, + extending the path. ``observations``, ``rewards``, ``events``, sub-groups + like ``observations.policy`` / ``observations.perception``, and anything + nested deeper are reached transparently. + * Anything else (plain Python data, callables, non-configclass objects): + stop. No yield, no recursion. + Driven entirely by type — no attribute names are hardcoded — so future + cfg layouts (extra observation groups, new nesting, etc.) are picked up + automatically as long as their terms subclass :class:`ManagerTermBaseCfg`. + """ + from isaaclab.managers.manager_term_cfg import ManagerTermBaseCfg -def _iter_term_attrs(group: Any) -> Iterator[tuple[str, Any]]: - """Yield ``(name, term)`` pairs from a manager-cfg group, skipping dunders/Nones.""" - if group is None: + if isinstance(node, ManagerTermBaseCfg): + yield path, node + return + if not hasattr(node, "__dataclass_fields__"): return - for name in [n for n in dir(group) if not n.startswith("_")]: - term = getattr(group, name, None) - if term is None: + for name in dir(node): + if name.startswith("_"): + continue + try: + value = getattr(node, name, None) + except Exception: # noqa: BLE001 — defensive; some descriptors can raise on attribute access + continue + if value is None: continue - yield name, term + yield from _walk_terms(value, path + (name,)) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/managers/__init__.py index b4521b98434a..757da5def93f 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/__init__.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/__init__.py @@ -17,7 +17,6 @@ # Override the stable implementation with the experimental fork. from .manager_base import ManagerTermBase # noqa: F401 -from .manager_term_cfg import ObservationTermCfg, RewardTermCfg, TerminationTermCfg # noqa: F401 from .observation_manager import ObservationManager # noqa: F401 from .reward_manager import RewardManager # noqa: F401 from .scene_entity_cfg import SceneEntityCfg # noqa: F401 diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py index 672af82bc8ab..902c74d6d095 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py @@ -20,9 +20,9 @@ from isaaclab.assets import AssetBase from isaaclab.envs.utils.io_descriptors import GenericActionIODescriptor +from isaaclab.managers.manager_term_cfg import ActionTermCfg from .manager_base import ManagerBase, ManagerTermBase -from .manager_term_cfg import ActionTermCfg if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py index 8b4e8f83dbe7..dfafeef5e206 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py @@ -17,10 +17,11 @@ import warp as wp from prettytable import PrettyTable +from isaaclab.managers.manager_term_cfg import CommandTermCfg + from isaaclab_experimental.utils.warp.kernels import compute_reset_scale, count_masked from .manager_base import ManagerBase, ManagerTermBase -from .manager_term_cfg import CommandTermCfg # import omni.kit.app diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py index 3b36613a2ac9..cc5b77222674 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py @@ -29,8 +29,9 @@ import warp as wp from prettytable import PrettyTable +from isaaclab.managers.manager_term_cfg import EventTermCfg + from .manager_base import ManagerBase -from .manager_term_cfg import EventTermCfg logger = logging.getLogger(__name__) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py index 6caea3eebcbb..699e27170e02 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py @@ -26,11 +26,11 @@ import warp as wp import isaaclab.utils.string as string_utils +from isaaclab.managers.manager_term_cfg import ManagerTermBaseCfg from isaaclab.utils import class_to_dict, string_to_callable from isaaclab_experimental.utils.warp import is_warp_capturable -from .manager_term_cfg import ManagerTermBaseCfg from .scene_entity_cfg import SceneEntityCfg # import omni.timeline diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py deleted file mode 100644 index 604b138003d4..000000000000 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py +++ /dev/null @@ -1,25 +0,0 @@ -# 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 - -"""Configuration terms for different managers (experimental, Warp-first). - -The warp manager classes accept the same term cfg shapes as the stable -managers, so this module simply re-exports the stable term cfg classes -to keep ``isinstance(stable_term, warp.TermCfg)`` true. This is what -lets the :class:`isaaclab_experimental.envs.warp_frontend.WarpFrontend` -adapter run a stable cfg through the warp runtime without rewrapping -every term. - -The ``func`` callable on each term is still expected to follow the -warp-first ``func(env, out, **params) -> None`` signature when run on -the warp runtime; only the *type* is shared with stable. -""" - -from __future__ import annotations - -# Re-export stable manager term cfg classes verbatim. -# `from … import *` carries `ObservationTermCfg`, `RewardTermCfg`, -# `TerminationTermCfg`, `ManagerTermBaseCfg`, etc. -from isaaclab.managers.manager_term_cfg import * # noqa: F401,F403 diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index 74d0a9a955be..059ae2afdce1 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -55,6 +55,7 @@ import warp as wp from prettytable import PrettyTable +from isaaclab.managers.manager_term_cfg import ObservationGroupCfg, ObservationTermCfg from isaaclab.utils import class_to_dict from isaaclab_experimental.utils import modifiers, noise @@ -62,7 +63,6 @@ from isaaclab_experimental.utils.torch_utils import clone_obs_buffer from .manager_base import ManagerBase, ManagerTermBase -from .manager_term_cfg import ObservationGroupCfg, ObservationTermCfg if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py index 67dcbc055c2f..99e3d9af979c 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py @@ -18,10 +18,11 @@ import warp as wp from prettytable import PrettyTable +from isaaclab.managers.manager_term_cfg import RewardTermCfg + from isaaclab_experimental.utils.warp.kernels import compute_reset_scale, count_masked from .manager_base import ManagerBase, ManagerTermBase -from .manager_term_cfg import RewardTermCfg if TYPE_CHECKING: from isaaclab.envs import ManagerBasedRLEnv diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py index 8a768651b292..732edf9f2fe2 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py @@ -22,8 +22,9 @@ import warp as wp from prettytable import PrettyTable +from isaaclab.managers.manager_term_cfg import TerminationTermCfg + from .manager_base import ManagerBase, ManagerTermBase -from .manager_term_cfg import TerminationTermCfg if TYPE_CHECKING: from isaaclab.envs import ManagerBasedRLEnv diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py index a35c21dd07ed..7f884674f58e 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend.py +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -10,12 +10,13 @@ * :class:`Frontend` / :class:`Workflow` enum surface. * :func:`SceneEntityCfg.from_stable` field copy. * :func:`_require_newton_physics` hard-check. +* :func:`_walk_terms` recursive ManagerTermBaseCfg discovery. * :func:`_promote_scene_entity_cfgs` walks ``term.params`` dicts. * :func:`_swap_mdp` swaps ``func`` *and* ``class_type``; raises with a path list when twins are missing. * :func:`_resolve_warp_twin` rejects stable-origin re-exports. -* :func:`_require_direct_is_warp_task` accepts warp-rooted entry points - and rejects stable ones. +* :func:`_assert_direct_warp_registration` accepts warp-rooted entry + points and rejects stable ones. """ from __future__ import annotations @@ -29,20 +30,21 @@ Frontend, FrontendIncompatibleError, Workflow, + _assert_direct_warp_registration, _is_swap_candidate, - _iter_term_attrs, _promote_scene_entity_cfgs, - _require_direct_is_warp_task, _require_newton_physics, _resolve_warp_twin, _swap_mdp, - _walk_attrs, + _walk_terms, ) from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as WarpSceneEntityCfg from isaaclab_newton.physics import NewtonCfg from isaaclab_physx.physics import PhysxCfg +from isaaclab.managers.manager_term_cfg import EventTermCfg, ObservationTermCfg, RewardTermCfg from isaaclab.managers.scene_entity_cfg import SceneEntityCfg as StableSceneEntityCfg +from isaaclab.utils.configclass import configclass # ====================================================================== # Enums @@ -139,116 +141,189 @@ def test_rejects_none(self): # ====================================================================== -# _promote_scene_entity_cfgs +# Configclass fixtures for the walker / swap tests. # ====================================================================== +# +# These mirror the real cfg shape so :func:`_walk_terms` descends into them +# (it only descends into objects with ``__dataclass_fields__``). Term cfgs +# use the real :class:`EventTermCfg`/:class:`RewardTermCfg`/:class:`ObservationTermCfg` +# so the walker's ``isinstance(ManagerTermBaseCfg)`` discriminator yields them. -class _Term: - """Minimal stand-in for a manager term cfg.""" +def _stable_func(env, **params): + return None - def __init__(self, params: dict | None = None, func=None, class_type=None): - self.params = params if params is not None else {} - if func is not None: - self.func = func - if class_type is not None: - self.class_type = class_type +class _StableActionCls: + pass -class _Group: - """Minimal stand-in for a manager-cfg group (e.g. RewardsCfg).""" - def __init__(self, **terms: Any): - for name, term in terms.items(): - setattr(self, name, term) +_stable_func.__module__ = "isaaclab_tasks.fake_task.mdp" +_StableActionCls.__module__ = "isaaclab_tasks.fake_task.mdp" -class _ObsCfg: - def __init__(self, policy: _Group): - self.policy = policy +def _warp_twin_func(env, out, **params): + return None -def _make_cfg(**groups: Any) -> Any: - cfg = types.SimpleNamespace() - for name, value in groups.items(): - setattr(cfg, name, value) - return cfg +class _WarpActionCls: + pass -class TestPromoteSceneEntityCfgs(unittest.TestCase): - def test_promotes_in_params(self): - stable = StableSceneEntityCfg(name="robot", joint_names=["lf_hip"]) - term = _Term(params={"asset_cfg": stable, "scale": 1.0}) - cfg = _make_cfg(rewards=_Group(track=term)) - _promote_scene_entity_cfgs(cfg) - promoted = term.params["asset_cfg"] - self.assertIsInstance(promoted, WarpSceneEntityCfg) - self.assertEqual(promoted.name, "robot") - self.assertEqual(promoted.joint_names, ["lf_hip"]) - # Non-SceneEntityCfg params are untouched. - self.assertEqual(term.params["scale"], 1.0) +_warp_twin_func.__module__ = "isaaclab_experimental.envs.mdp" +_WarpActionCls.__module__ = "isaaclab_experimental.envs.mdp" + + +@configclass +class _PolicyObsGroup: + """Stand-in for a per-task ObservationsCfg sub-group (e.g. PolicyCfg).""" + + o1: ObservationTermCfg | None = None + o2: ObservationTermCfg | None = None + + +@configclass +class _ExtraObsGroup: + """A second obs group (named arbitrarily) to exercise multi-group walks.""" + + o3: ObservationTermCfg | None = None - def test_skips_already_warp(self): - warp = WarpSceneEntityCfg(name="robot") - term = _Term(params={"asset_cfg": warp}) - cfg = _make_cfg(rewards=_Group(t=term)) - _promote_scene_entity_cfgs(cfg) - self.assertIs(term.params["asset_cfg"], warp) # unchanged identity - - def test_walks_all_term_paths(self): - # Drop a stable cfg in every supported group; all should promote. - groups = {} - terms: list[_Term] = [] - for group_name in ("rewards", "events", "terminations", "commands", "curriculum"): - t = _Term(params={"asset_cfg": StableSceneEntityCfg(name=group_name)}) - terms.append(t) - groups[group_name] = _Group(t=t) - # observations is nested as observations.policy - obs_term = _Term(params={"asset_cfg": StableSceneEntityCfg(name="obs")}) - terms.append(obs_term) - groups["observations"] = _ObsCfg(policy=_Group(t=obs_term)) - # actions group also walked - act_term = _Term(params={"asset_cfg": StableSceneEntityCfg(name="act")}) - terms.append(act_term) - groups["actions"] = _Group(t=act_term) - - cfg = _make_cfg(**groups) - _promote_scene_entity_cfgs(cfg) - for t in terms: - self.assertIsInstance(t.params["asset_cfg"], WarpSceneEntityCfg) + +@configclass +class _ObservationsCfg: + policy: _PolicyObsGroup | None = None + perception: _ExtraObsGroup | None = None + + +@configclass +class _RewardsCfg: + r1: RewardTermCfg | None = None + r2: RewardTermCfg | None = None + + +@configclass +class _EventsCfg: + e1: EventTermCfg | None = None + + +@configclass +class _CfgFixture: + observations: _ObservationsCfg | None = None + rewards: _RewardsCfg | None = None + events: _EventsCfg | None = None + + +def _term(func=None, params: dict | None = None) -> RewardTermCfg: + """Cheap RewardTermCfg builder for tests; the cfg class is irrelevant for swap/walk logic.""" + return RewardTermCfg(func=func or _stable_func, weight=1.0, params=params or {}) # ====================================================================== -# _swap_mdp +# _walk_terms # ====================================================================== -# A fake "stable" mdp module the test fixture stands in for -# ``isaaclab_tasks..mdp`` — symbols here pretend to live there. -def _stable_func(env, **params): - return None +class TestWalkTerms(unittest.TestCase): + def test_yields_each_term_with_its_path(self): + cfg = _CfgFixture( + rewards=_RewardsCfg(r1=_term(), r2=_term()), + events=_EventsCfg(e1=EventTermCfg(func=_stable_func, mode="reset")), + ) + # Configclass instances aren't hashable, so collect paths only. + paths = {".".join(p) for p, _ in _walk_terms(cfg)} + self.assertEqual(paths, {"rewards.r1", "rewards.r2", "events.e1"}) + + def test_descends_into_obs_subgroups(self): + cfg = _CfgFixture( + observations=_ObservationsCfg( + policy=_PolicyObsGroup(o1=ObservationTermCfg(func=_stable_func)), + perception=_ExtraObsGroup(o3=ObservationTermCfg(func=_stable_func)), + ), + ) + paths = {".".join(p) for p, _ in _walk_terms(cfg)} + # Discovery is purely type-driven; no obs group name is hardcoded. + self.assertEqual(paths, {"observations.policy.o1", "observations.perception.o3"}) + def test_stops_at_terms(self): + # The walker must not descend into term.params / term.func — yields the term itself. + nested_se_cfg = StableSceneEntityCfg(name="robot") + cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(params={"asset_cfg": nested_se_cfg}))) + terms = list(_walk_terms(cfg)) + self.assertEqual(len(terms), 1) + _, term = terms[0] + self.assertIsInstance(term, RewardTermCfg) -class _StableActionCls: - pass + def test_skips_non_configclass_attrs(self): + # A namespace without __dataclass_fields__ is not descended into. + cfg = types.SimpleNamespace(some_plain_attr="hello") + self.assertEqual(list(_walk_terms(cfg)), []) + def test_skips_none_subtrees(self): + cfg = _CfgFixture(rewards=None, events=None) + self.assertEqual(list(_walk_terms(cfg)), []) -_stable_func.__module__ = "isaaclab_tasks.fake_task.mdp" -_StableActionCls.__module__ = "isaaclab_tasks.fake_task.mdp" +# ====================================================================== +# _promote_scene_entity_cfgs +# ====================================================================== -# And a fake warp twin module the test installs into sys.modules under the -# expected name resolution path. Symbols here pretend to live under -# ``isaaclab_experimental.envs.mdp`` (the fallback). -def _warp_twin_func(env, out, **params): - return None +class TestPromoteSceneEntityCfgs(unittest.TestCase): + def test_promotes_in_params(self): + cfg = _CfgFixture( + rewards=_RewardsCfg( + r1=_term(params={"asset_cfg": StableSceneEntityCfg(name="robot", joint_names=["lf_hip"]), "scale": 1.0}) + ) + ) + _promote_scene_entity_cfgs(cfg) + promoted = cfg.rewards.r1.params["asset_cfg"] + self.assertIsInstance(promoted, WarpSceneEntityCfg) + self.assertEqual(promoted.name, "robot") + self.assertEqual(promoted.joint_names, ["lf_hip"]) + # Non-SceneEntityCfg params are untouched. + self.assertEqual(cfg.rewards.r1.params["scale"], 1.0) -class _WarpActionCls: - pass + def test_skips_already_warp(self): + # configclass init deep-copies params, so identity won't hold across + # construction; what we actually want to assert is "no re-promotion": + # the asset_cfg remains a WarpSceneEntityCfg (i.e., wasn't passed + # back through `from_stable`). + warp = WarpSceneEntityCfg(name="robot", joint_names=["lf_hip"]) + cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(params={"asset_cfg": warp}))) + before = cfg.rewards.r1.params["asset_cfg"] + _promote_scene_entity_cfgs(cfg) + after = cfg.rewards.r1.params["asset_cfg"] + self.assertIsInstance(after, WarpSceneEntityCfg) + # The asset_cfg object was not replaced by another from_stable call. + self.assertIs(after, before) + + def test_walks_all_term_groups(self): + cfg = _CfgFixture( + rewards=_RewardsCfg(r1=_term(params={"asset_cfg": StableSceneEntityCfg(name="r")})), + events=_EventsCfg( + e1=EventTermCfg(func=_stable_func, mode="reset", params={"asset_cfg": StableSceneEntityCfg(name="e")}) + ), + observations=_ObservationsCfg( + policy=_PolicyObsGroup( + o1=ObservationTermCfg(func=_stable_func, params={"asset_cfg": StableSceneEntityCfg(name="o-pol")}) + ), + perception=_ExtraObsGroup( + o3=ObservationTermCfg(func=_stable_func, params={"asset_cfg": StableSceneEntityCfg(name="o-per")}) + ), + ), + ) + _promote_scene_entity_cfgs(cfg) + self.assertIsInstance(cfg.rewards.r1.params["asset_cfg"], WarpSceneEntityCfg) + self.assertIsInstance(cfg.events.e1.params["asset_cfg"], WarpSceneEntityCfg) + self.assertIsInstance(cfg.observations.policy.o1.params["asset_cfg"], WarpSceneEntityCfg) + # The perception sub-group is reached even though its attribute name is + # not hardcoded in the framework. + self.assertIsInstance(cfg.observations.perception.o3.params["asset_cfg"], WarpSceneEntityCfg) -_warp_twin_func.__module__ = "isaaclab_experimental.envs.mdp" -_WarpActionCls.__module__ = "isaaclab_experimental.envs.mdp" +# ====================================================================== +# _swap_mdp +# ====================================================================== class _FakeMdpModule: @@ -265,8 +340,6 @@ def _patch_modules(self, name_to_symbol: dict[str, Any]) -> _FakeMdpModule: return m def _patched_warp_mdp_modules(self, modules: list[Any]): - # Patch _warp_mdp_modules at the frontend module level so _swap_mdp - # uses our fake module list instead of doing real importlib lookups. import isaaclab_experimental.envs.frontend as fe self._orig = fe._warp_mdp_modules @@ -284,34 +357,32 @@ def tearDown(self) -> None: def test_swaps_func_and_class_type(self): fake = self._patch_modules({"_stable_func": _warp_twin_func, "_StableActionCls": _WarpActionCls}) self._patched_warp_mdp_modules([fake]) - term_reward = _Term(func=_stable_func) - term_action = _Term(class_type=_StableActionCls) - cfg = _make_cfg(rewards=_Group(r=term_reward), actions=_Group(a=term_action)) + term_reward = _term(func=_stable_func) + term_action = _term() + term_action.class_type = _StableActionCls # set attr to exercise class_type swap + cfg = _CfgFixture(rewards=_RewardsCfg(r1=term_reward, r2=term_action)) _swap_mdp(cfg, "Isaac-Test-v0") - self.assertIs(term_reward.func, _warp_twin_func) - self.assertIs(term_action.class_type, _WarpActionCls) + self.assertIs(cfg.rewards.r1.func, _warp_twin_func) + self.assertIs(cfg.rewards.r2.class_type, _WarpActionCls) def test_missing_twin_raises_with_path_list(self): - # No twins available — every term should be reported as missing. fake = self._patch_modules({}) self._patched_warp_mdp_modules([fake]) - term = _Term(func=_stable_func) - cfg = _make_cfg(rewards=_Group(track=term)) + cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(func=_stable_func))) with self.assertRaises(FrontendIncompatibleError) as exc: _swap_mdp(cfg, "Isaac-Test-v0") msg = str(exc.exception) - self.assertIn("rewards.track.func", msg) + self.assertIn("rewards.r1.func", msg) self.assertIn("_stable_func", msg) - # The cfg term wasn't mutated (we hard-failed before partial application). - self.assertIs(term.func, _stable_func) + # The cfg term wasn't mutated for the missing twin. + self.assertIs(cfg.rewards.r1.func, _stable_func) def test_skips_already_warp(self): fake = self._patch_modules({}) self._patched_warp_mdp_modules([fake]) - term = _Term(func=_warp_twin_func) # already warp-origin - cfg = _make_cfg(rewards=_Group(r=term)) + cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(func=_warp_twin_func))) _swap_mdp(cfg, "Isaac-Test-v0") # no raise - self.assertIs(term.func, _warp_twin_func) + self.assertIs(cfg.rewards.r1.func, _warp_twin_func) # ====================================================================== @@ -364,7 +435,7 @@ def test_non_callable_is_not(self): } -class TestRequireDirectIsWarpTask(unittest.TestCase): +class TestAssertDirectWarpRegistration(unittest.TestCase): @classmethod def setUpClass(cls): # Register stub tasks so ``gym.spec`` resolves them. Entry-point strings @@ -376,41 +447,19 @@ def setUpClass(cls): gym.register(id=task_id, entry_point=ep, disable_env_checker=True) cls._registered.append(task_id) except gym.error.Error: - # Already registered from a previous run — that's fine. pass def test_accepts_warp_rooted(self): - _require_direct_is_warp_task("_Frontend-Test-Warp-Direct-v0") # no raise + _assert_direct_warp_registration("_Frontend-Test-Warp-Direct-v0") # no raise def test_rejects_stable_rooted(self): with self.assertRaises(FrontendIncompatibleError) as exc: - _require_direct_is_warp_task("_Frontend-Test-Stable-Direct-v0") + _assert_direct_warp_registration("_Frontend-Test-Stable-Direct-v0") self.assertIn("isaaclab_experimental", str(exc.exception)) def test_rejects_unknown_task(self): with self.assertRaises(FrontendIncompatibleError): - _require_direct_is_warp_task("Frontend-Test-NotRegistered-v0") - - -# ====================================================================== -# Walk helpers -# ====================================================================== - - -class TestWalkHelpers(unittest.TestCase): - def test_walk_attrs_hit(self): - root = types.SimpleNamespace(a=types.SimpleNamespace(b=types.SimpleNamespace(c=42))) - self.assertEqual(_walk_attrs(root, ("a", "b", "c")), 42) - - def test_walk_attrs_miss(self): - root = types.SimpleNamespace(a=types.SimpleNamespace()) - self.assertIsNone(_walk_attrs(root, ("a", "missing"))) - self.assertIsNone(_walk_attrs(root, ("missing", "b"))) - - def test_iter_term_attrs_skips_dunders_and_none(self): - g = types.SimpleNamespace(t1=_Term(), t2=None, _internal=_Term()) - names = sorted(n for n, _ in _iter_term_attrs(g)) - self.assertEqual(names, ["t1"]) + _assert_direct_warp_registration("Frontend-Test-NotRegistered-v0") if __name__ == "__main__": diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py index f0e47dfbbf05..beecd526d245 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py @@ -23,9 +23,8 @@ from isaaclab.assets import Articulation if TYPE_CHECKING: - from isaaclab_experimental.managers.manager_term_cfg import RewardTermCfg - from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.managers.manager_term_cfg import RewardTermCfg # --------------------------------------------------------------------------- From 21e8bba0d68db95e8a436982d107436a32974295 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 26 May 2026 07:07:27 +0000 Subject: [PATCH 12/58] Compress changelog fragment + add per-package fragments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The isaaclab_experimental fragment shrinks from 47 lines of nested narrative to 12 lines per the existing terse-fragment rule. Adds the matching fragments for the two other touched packages: isaaclab (the term-cfg type-hint relax) and isaaclab_tasks_experimental (.skip — the humanoid import redirect is internal). --- .../changelog.d/warp-manager-bridge.rst | 8 ++++ .../changelog.d/warp-manager-bridge.rst | 48 ++++--------------- .../changelog.d/warp-manager-bridge.skip | 1 + 3 files changed, 17 insertions(+), 40 deletions(-) create mode 100644 source/isaaclab/changelog.d/warp-manager-bridge.rst create mode 100644 source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.skip diff --git a/source/isaaclab/changelog.d/warp-manager-bridge.rst b/source/isaaclab/changelog.d/warp-manager-bridge.rst new file mode 100644 index 000000000000..8aa615a5294a --- /dev/null +++ b/source/isaaclab/changelog.d/warp-manager-bridge.rst @@ -0,0 +1,8 @@ +Changed +^^^^^^^ + +* Relaxed ``func`` annotation on :class:`~isaaclab.managers.ObservationTermCfg`, + :class:`~isaaclab.managers.RewardTermCfg`, and + :class:`~isaaclab.managers.TerminationTermCfg` to ``Callable[..., torch.Tensor | None]`` + so kernel-style ``func(env, out) -> None`` terms type-check alongside the + existing torch return-tensor form. diff --git a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst index 92dfd4f06f7b..3e198ce1db37 100644 --- a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst +++ b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst @@ -1,47 +1,15 @@ Added ^^^^^ -* Added :mod:`isaaclab_experimental.envs.frontend`, a small runtime selector - used by ``--frontend {torch,warp}`` to choose how a task is constructed. - - * ``torch`` (default) dispatches via :func:`gym.make` unchanged. - * ``warp`` adapts a stable manager-based cfg in place and constructs - :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp`. The adapter - does three things, all hard-failing on incompatibility: - - 1. Requires :class:`~isaaclab_newton.physics.NewtonCfg` as the active - physics — the user is responsible for selecting it via - ``presets=newton``; no Hydra-arg mutation happens behind the scenes. - 2. Promotes every stable :class:`~isaaclab.managers.SceneEntityCfg` - embedded under term ``params`` to - :class:`isaaclab_experimental.managers.SceneEntityCfg` via the new - :meth:`~isaaclab_experimental.managers.SceneEntityCfg.from_stable` - classmethod (no ``__class__`` reassignment). - 3. Swaps every stable ``term.func`` *and* ``term.class_type`` (one pass, - handles observations / events / rewards / terminations / commands / - curriculum / actions) with its same-named warp twin from - ``isaaclab_tasks_experimental..mdp`` or - :mod:`isaaclab_experimental.envs.mdp`. Any missing twin raises - :class:`FrontendIncompatibleError` listing the affected term paths. - - Direct workflows aren't adapted; ``--frontend=warp`` on a direct task - requires the task to be pre-registered under ``isaaclab_experimental`` / - ``isaaclab_tasks_experimental`` (e.g. ``*-Direct-Warp-v0``). - -* Added :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable`, - a classmethod that builds a warp scene-entity cfg from a stable one by - copying every selection field through ``__init__``. - -* Added a ``--frontend {torch,warp}`` flag to ``rsl_rl/train.py``. The flag - selects the runtime via :func:`isaaclab_experimental.envs.frontend.build`; - ``isaaclab_experimental`` is treated as optional and ``--frontend=torch`` - falls back to ``gym.make`` when it isn't installed. +* Added ``--frontend {torch,warp}`` to ``rsl_rl/train.py`` for selecting the env + runtime; default ``torch`` is unchanged. +* Added :mod:`isaaclab_experimental.envs.frontend` runtime selector and + :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable` used by + ``--frontend=warp`` to adapt stable cfgs onto the warp runtime. Fixed ^^^^^ -* Fixed a regression in :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` - introduced when the ``SimulationContext.get_setting`` API was reshaped: - the warp env now mirrors the stable env and probes - :meth:`~isaaclab.sim.SimulationContext.has_active_visualizers` instead of - splitting a string setting that no longer exists. +* Fixed :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` to probe + :meth:`~isaaclab.sim.SimulationContext.has_active_visualizers` after the + ``get_setting`` API was reshaped. diff --git a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.skip b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.skip new file mode 100644 index 000000000000..e55772b9bda3 --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.skip @@ -0,0 +1 @@ +internal: redirect humanoid reward import to isaaclab.managers.manager_term_cfg From b932897dd2894af615b7c829da21cc6b0dbedc41 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 26 May 2026 07:20:12 +0000 Subject: [PATCH 13/58] Sync warp envs to stable visualizer / has_gui API Two missed-sync fixes carried alongside the manager_based viewport sync that already landed earlier in this branch: 1. DirectRLEnvWarp viewport controller now mirrors stable DirectRLEnv: '(has_gui or has_active_visualizers()) and has_kit()' gate, instantiate ViewportCameraController accordingly. Replaces the stale get_setting('/isaaclab/has_gui'/'render/offscreen') check and removes the commented-out instantiation that left viewport_camera_controller unconditionally None. 2. render(mode='rgb_array') in ManagerBasedRLEnvWarp and DirectRLEnvWarp now read SimulationContext.has_gui and has_offscreen_render properties instead of get_setting on keys that were removed when the simulation manager was refactored. Without this the gate silently fell through on None values. --- .../envs/direct_rl_env_warp.py | 30 ++++++++----------- .../envs/manager_based_rl_env_warp.py | 8 ++--- 2 files changed, 16 insertions(+), 22 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py index 20c82e582690..979826e07345 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py @@ -27,15 +27,15 @@ from isaaclab.envs.common import VecEnvObs, VecEnvStepReturn from isaaclab.envs.direct_rl_env import DirectRLEnv from isaaclab.envs.direct_rl_env_cfg import DirectRLEnvCfg +from isaaclab.envs.ui import ViewportCameraController from isaaclab.envs.utils.spaces import sample_space, spec_to_gym_space - -# from isaaclab.envs.ui import ViewportCameraController from isaaclab.managers import EventManager from isaaclab.sim import SimulationContext from isaaclab.sim.utils import use_stage from isaaclab.utils.noise import NoiseModel from isaaclab.utils.seed import configure_seed from isaaclab.utils.timer import Timer +from isaaclab.utils.version import has_kit from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache @@ -168,15 +168,12 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs # attach_stage_to_usd_context() print("[INFO]: Scene manager: ", self.scene) - # set up camera viewport controller - # viewport is not available in other rendering modes so the function will throw a warning - # FIXME: This needs to be fixed in the future when we unify the UI functionalities even for - # non-rendering modes. - has_gui = bool(self.sim.get_setting("/isaaclab/has_gui")) - offscreen_render = bool(self.sim.get_setting("/isaaclab/render/offscreen")) - if has_gui or offscreen_render: - # self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) - self.viewport_camera_controller = None + # Initialize when a Kit viewport exists. ViewportCameraController uses omni.kit + # (renderer camera); skip in kitless Newton-only runs where no Kit app is running. + # Mirrors stable :class:`isaaclab.envs.DirectRLEnv` (PR #5103 changed the visualizer API). + has_visualizers = self.sim.has_active_visualizers() + if (self.sim.has_gui or has_visualizers) and has_kit(): + self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) else: self.viewport_camera_controller = None @@ -571,13 +568,10 @@ def render(self, recompute: bool = False) -> np.ndarray | None: if self.render_mode == "human" or self.render_mode is None: return None elif self.render_mode == "rgb_array": - # check that if any render could have happened - has_gui = bool(self.sim.get_setting("/isaaclab/has_gui")) - offscreen_render = bool(self.sim.get_setting("/isaaclab/render/offscreen")) - # Rendering is possible if we have GUI or offscreen rendering enabled - can_render = has_gui or offscreen_render - - if not can_render: + # Rendering is possible if we have GUI or offscreen rendering enabled — mirror + # stable (PR #4646 replaced /isaaclab/has_gui and /isaaclab/render/offscreen + # settings with cached SimulationContext properties). + if not (self.sim.has_gui or self.sim.has_offscreen_render): render_mode_name = "NO_GUI_OR_RENDERING" raise RuntimeError( f"Cannot render '{self.render_mode}' when the simulation render mode is" diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index 2e1894e54003..06ba2db190ae 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -408,10 +408,10 @@ def render(self, recompute: bool = False) -> np.ndarray | None: if self.render_mode == "human" or self.render_mode is None: return None elif self.render_mode == "rgb_array": - # check that if any render could have happened - has_gui = bool(self.sim.get_setting("/isaaclab/has_gui")) - offscreen_render = bool(self.sim.get_setting("/isaaclab/render/offscreen")) - if not (has_gui or offscreen_render): + # check that if any render could have happened — mirror stable + # (PR #4646 replaced /isaaclab/has_gui and /isaaclab/render/offscreen settings + # with cached SimulationContext properties). + if not (self.sim.has_gui or self.sim.has_offscreen_render): raise RuntimeError( f"Cannot render '{self.render_mode}' when the simulation render mode does not support" " rendering. Please set the simulation render mode to 'PARTIAL_RENDERING' or" From 0f3447541b8860cdd7b0efcd1b7a3b337994019d Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 27 May 2026 04:48:54 +0000 Subject: [PATCH 14/58] Rename CLI preset hint to newton_mjwarp The legacy bare 'newton'/'kamino' preset names were renamed on develop to 'newton_mjwarp'/'newton_kamino' to disambiguate from the Newton backend label. Update the warp-frontend hard-check error message, the `_adapt_cfg_for_warp`/`_require_newton_physics` docstrings, the rsl_rl/train.py comment, and the matching test assertion to spell the canonical preset name so users get the right CLI token in errors. --- scripts/reinforcement_learning/rsl_rl/train.py | 2 +- .../isaaclab_experimental/envs/frontend.py | 6 +++--- source/isaaclab_experimental/test/envs/test_frontend.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/scripts/reinforcement_learning/rsl_rl/train.py b/scripts/reinforcement_learning/rsl_rl/train.py index 1f30de33eb12..65a4e03623b1 100644 --- a/scripts/reinforcement_learning/rsl_rl/train.py +++ b/scripts/reinforcement_learning/rsl_rl/train.py @@ -121,7 +121,7 @@ # Build path for ``--frontend=warp`` lives in ``isaaclab_experimental``, which is an # optional package — torch (the default) must still work without it, so import lazily. -# Warp callers are required to pass ``presets=newton`` themselves; the frontend hard-checks +# Warp callers are required to pass ``presets=newton_mjwarp`` themselves; the frontend hard-checks # it at build time rather than mutating Hydra args here. _frontend_build = None try: diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index f8b4acd2284a..518cebb45e3a 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -131,7 +131,7 @@ def _adapt_cfg_for_warp(cfg: Any, task_id: str) -> None: 1. :func:`_require_newton_physics` — hard check that ``cfg.sim.physics`` is :class:`~isaaclab_newton.physics.NewtonCfg`. The user is responsible for selecting the Newton variant of the task's :class:`PresetCfg` via - ``presets=newton``; we don't auto-inject. + ``presets=newton_mjwarp``; we don't auto-inject. 2. :func:`_promote_scene_entity_cfgs` — replace stable :class:`~isaaclab.managers.SceneEntityCfg` instances under each term's ``params`` with the warp variant (which adds warp-cached ``joint_mask``, @@ -153,7 +153,7 @@ def _require_newton_physics(cfg: Any, task_id: str) -> None: The warp managers' assets read state through :class:`NewtonManager`; a :class:`PhysxCfg` (or unresolved :class:`PresetCfg`) is a hard - incompatibility. The fix is to pass ``presets=newton`` on the CLI so + incompatibility. The fix is to pass ``presets=newton_mjwarp`` on the CLI so Hydra resolves the task's :class:`PresetCfg` wrapper to the Newton field before construction. """ @@ -164,7 +164,7 @@ def _require_newton_physics(cfg: Any, task_id: str) -> None: return raise FrontendIncompatibleError( f"--frontend=warp on {task_id!r}: expected cfg.sim.physics to be NewtonCfg," - f" got {type(physics).__name__!r}. Pass `presets=newton` on the CLI so" + f" got {type(physics).__name__!r}. Pass `presets=newton_mjwarp` on the CLI so" f" Hydra resolves the task's PresetCfg wrapper to the Newton variant." ) diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py index 7f884674f58e..28666ed78e34 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend.py +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -131,7 +131,7 @@ def test_rejects_physx(self): cfg = self._cfg_with(PhysxCfg()) with self.assertRaises(FrontendIncompatibleError) as exc: _require_newton_physics(cfg, "Isaac-Test-v0") - self.assertIn("presets=newton", str(exc.exception)) + self.assertIn("presets=newton_mjwarp", str(exc.exception)) self.assertIn("PhysxCfg", str(exc.exception)) def test_rejects_none(self): From 1a66c3d449a7281cd1816dff65f838acace35247 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 4 Jun 2026 23:34:55 +0000 Subject: [PATCH 15/58] Self-adapt stable cfg inside warp env construction Move the stable->warp cfg adaptation out of the train.py --frontend path and into ManagerBasedEnvWarp.__init__, so a task registered directly as *-Warp-v0 (with env_cfg_entry_point pointing at the stable cfg) goes through the exact same Newton-physics check, SceneEntityCfg promotion and MDP twin swap. Twin lookup is now keyed off the stable cfg class module (type(cfg).__module__) instead of the gym task id, so it no longer needs a registration to resolve. Adaptation is idempotent. This lets the warp side reuse stable env cfgs instead of carrying parallel copies that drift. --- .../isaaclab_experimental/envs/frontend.py | 53 +++++++++++-------- .../envs/manager_based_env_warp.py | 8 +++ 2 files changed, 40 insertions(+), 21 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index 518cebb45e3a..2f186d3617d8 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -39,6 +39,7 @@ "Frontend", "FrontendIncompatibleError", "Workflow", + "adapt_cfg_for_warp", "build", ] @@ -110,9 +111,11 @@ def build( _assert_direct_warp_registration(task_id) return gym.make(task_id, cfg=env_cfg, **construct_kwargs) - _adapt_cfg_for_warp(env_cfg, task_id) # Imported lazily so that ``--frontend=torch`` callers don't pay the - # ``isaaclab_experimental.envs`` import cost. + # ``isaaclab_experimental.envs`` import cost. The warp env adapts the stable + # cfg in its own ``__init__`` (see :func:`adapt_cfg_for_warp`), so a task + # registered directly as ``*-Warp-v0`` (constructed straight from the stable + # cfg) goes through the exact same adaptation as ``--frontend=warp``. from isaaclab_experimental.envs import ManagerBasedRLEnvWarp return ManagerBasedRLEnvWarp(cfg=env_cfg, **construct_kwargs) @@ -123,8 +126,14 @@ def build( # --------------------------------------------------------------------------- -def _adapt_cfg_for_warp(cfg: Any, task_id: str) -> None: - """Mutate ``cfg`` in place so warp managers can consume it. +def adapt_cfg_for_warp(cfg: Any) -> None: + """Mutate a stable manager-based ``cfg`` in place so warp managers can consume it. + + Called from :meth:`ManagerBasedEnvWarp.__init__`, so it runs identically + whether the env is built via ``--frontend=warp`` on a stable task id or by + constructing a ``*-Warp-v0`` registration straight from the stable cfg. + Idempotent: re-running on an already-adapted cfg is a no-op (the steps below + skip warp-origin symbols / already-promoted entities). Three steps, each independently testable: @@ -143,12 +152,13 @@ def _adapt_cfg_for_warp(cfg: Any, task_id: str) -> None: twin raises :class:`FrontendIncompatibleError` — partial coverage is unsafe under the warp managers' kernel-only signature. """ - _require_newton_physics(cfg, task_id) + label = type(cfg).__name__ + _require_newton_physics(cfg, label) _promote_scene_entity_cfgs(cfg) - _swap_mdp(cfg, task_id) + _swap_mdp(cfg, label) -def _require_newton_physics(cfg: Any, task_id: str) -> None: +def _require_newton_physics(cfg: Any, label: str) -> None: """Block unless ``cfg.sim.physics`` is :class:`NewtonCfg`. The warp managers' assets read state through :class:`NewtonManager`; @@ -163,7 +173,7 @@ def _require_newton_physics(cfg: Any, task_id: str) -> None: if isinstance(physics, NewtonCfg): return raise FrontendIncompatibleError( - f"--frontend=warp on {task_id!r}: expected cfg.sim.physics to be NewtonCfg," + f"warp env {label!r}: expected cfg.sim.physics to be NewtonCfg," f" got {type(physics).__name__!r}. Pass `presets=newton_mjwarp` on the CLI so" f" Hydra resolves the task's PresetCfg wrapper to the Newton variant." ) @@ -199,7 +209,7 @@ def _promote_scene_entity_cfgs(cfg: Any) -> None: ) -def _swap_mdp(cfg: Any, task_id: str) -> None: +def _swap_mdp(cfg: Any, label: str) -> None: """Replace ``term.func`` and ``term.class_type`` with their warp twins. Iterates every term cfg in the tree (via :func:`_walk_terms`) and on each @@ -218,7 +228,7 @@ def _swap_mdp(cfg: Any, task_id: str) -> None: substitutes the callable; the manager reads the new func's annotations when it parses the term cfg. """ - modules = _warp_mdp_modules(task_id) + modules = _warp_mdp_modules(type(cfg).__module__) searched = tuple(m.__name__ for m in modules) logger.info("frontend.warp: searching warp mdp modules %s", list(searched)) @@ -240,7 +250,7 @@ def _swap_mdp(cfg: Any, task_id: str) -> None: if missing: lines = "\n ".join(f"{loc}.{attr}: no warp twin for {sym!r}" for loc, attr, sym in missing) raise FrontendIncompatibleError( - f"--frontend=warp on {task_id!r}: missing warp MDP twins (searched {list(searched)}):\n {lines}" + f"warp env {label!r}: missing warp MDP twins (searched {list(searched)}):\n {lines}" ) logger.info("frontend.warp: swapped %d MDP symbol(s) to warp twins", swapped) @@ -285,25 +295,26 @@ def _assert_direct_warp_registration(task_id: str) -> None: ) -def _warp_mdp_modules(task_id: str) -> list[ModuleType]: +def _warp_mdp_modules(cfg_module: str) -> list[ModuleType]: """Locate warp MDP modules to consult for twin lookups. + The lookup is keyed off the *stable* cfg class's module + (``type(cfg).__module__``) rather than the gym task id, so it resolves the + same twins whether the env is built via ``--frontend=warp`` or constructed + directly from a ``*-Warp-v0`` registration. + Order of preference: 1. The task-specific module, derived by replacing ``isaaclab_tasks`` with - ``isaaclab_tasks_experimental`` in the task's ``env_cfg_entry_point`` - package and walking up to the first existing ``.mdp`` submodule. + ``isaaclab_tasks_experimental`` in ``cfg_module`` and walking up to the + first existing ``.mdp`` submodule. This relies on the experimental + package mirroring the stable ``core/`` layout. 2. The shared :mod:`isaaclab_experimental.envs.mdp` fallback (where generic warp twins live). """ modules: list[ModuleType] = [] - try: - spec = gym.spec(task_id) - except gym.error.NameNotFound: - spec = None - entry = spec.kwargs.get("env_cfg_entry_point") if spec is not None else None - if isinstance(entry, str) and entry.startswith("isaaclab_tasks."): - warp_pkg = entry.rsplit(".", 1)[0].replace("isaaclab_tasks", "isaaclab_tasks_experimental", 1) + if isinstance(cfg_module, str) and cfg_module.startswith("isaaclab_tasks."): + warp_pkg = cfg_module.replace("isaaclab_tasks", "isaaclab_tasks_experimental", 1) parts = warp_pkg.split(".") for depth in range(len(parts), 0, -1): target = ".".join(parts[:depth] + ["mdp"]) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py index 9c05dc207700..707a3978ee05 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py @@ -73,6 +73,14 @@ def __init__(self, cfg: ManagerBasedEnvCfg): RuntimeError: If a simulation context already exists. The environment must always create one since it configures the simulation context and controls the simulation. """ + # Adapt a stable manager-based cfg to the warp runtime in place: assert + # Newton physics, promote SceneEntityCfg to the warp variant, and swap + # every MDP func/class to its warp twin. Idempotent, so an already-warp + # cfg passes through untouched. This lets a ``*-Warp-v0`` task point its + # env_cfg_entry_point at the stable cfg rather than carry a drifting copy. + from isaaclab_experimental.envs.frontend import adapt_cfg_for_warp + + adapt_cfg_for_warp(cfg) # check that the config is valid cfg.validate() # store inputs to class From 40db07a58fb9329eb5d0f1f4aee8dac92b4b3e9f Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 4 Jun 2026 23:54:05 +0000 Subject: [PATCH 16/58] Mirror core/ layout in tasks_experimental; reuse stable cfgs Flatten isaaclab_tasks_experimental from the legacy manager_based/ | direct/ | classic/ trees into core// mirroring isaaclab_tasks/core, so the frontend's stable->warp module-path rewrite resolves twins again after the stable migration. Manager-based warp tasks now reuse the stable env cfg (env_cfg_entry_point points at isaaclab_tasks.core...; physics + MDP twins are swapped at construction) instead of carrying parallel copies that drift; the drift cfgs are deleted. Direct warp tasks keep their own env class + cfg. Twin lookup keys off each stable symbol's own module, so ant reusing manager_humanoid's mdp resolves with no per-task shim (mirrors core, which gives manager_ant no mdp/ of its own). --- .../isaaclab_experimental/envs/frontend.py | 66 ++-- .../{direct => core}/allegro_hand/__init__.py | 2 +- .../allegro_hand/allegro_hand_warp_env_cfg.py | 0 .../core/cartpole/__init__.py | 53 ++++ .../cartpole/cartpole_warp_env.py | 0 .../cartpole/cartpole_warp_env_cfg.py | 0 .../classic => core}/cartpole/mdp/__init__.py | 0 .../cartpole/mdp/__init__.pyi | 0 .../classic => core}/cartpole/mdp/rewards.py | 0 .../ant => core/direct_ant}/__init__.py | 0 .../ant => core/direct_ant}/ant_env_warp.py | 2 +- .../direct_ant}/ant_env_warp_cfg.py | 0 .../direct_humanoid}/__init__.py | 0 .../direct_humanoid}/humanoid_warp_env.py | 2 +- .../direct_humanoid}/humanoid_warp_env_cfg.py | 0 .../direct_locomotion}/__init__.py | 0 .../direct_locomotion}/locomotion_env_warp.py | 0 .../inhand_manipulation}/__init__.py | 0 .../inhand_manipulation_warp_env.py | 2 +- .../ant => core/manager_ant}/__init__.py | 6 +- .../manager_humanoid}/__init__.py | 6 +- .../manager_humanoid}/mdp/__init__.py | 0 .../manager_humanoid}/mdp/__init__.pyi | 0 .../manager_humanoid}/mdp/observations.py | 0 .../manager_humanoid}/mdp/rewards.py | 0 .../manipulation => core}/reach/__init__.py | 0 .../reach/config/__init__.py | 0 .../reach/config/franka/__init__.py | 8 +- .../reach/config/ur_10/__init__.py | 8 +- .../reach/mdp/__init__.py | 0 .../reach/mdp/__init__.pyi | 0 .../reach/mdp/rewards.py | 0 .../locomotion => core}/velocity/__init__.py | 0 .../velocity/config/__init__.py | 0 .../velocity/config/a1/__init__.py | 8 +- .../velocity/config/anymal_b/__init__.py | 8 +- .../velocity/config/anymal_c/__init__.py | 8 +- .../velocity/config/anymal_d/__init__.py | 12 +- .../velocity/config/cassie/__init__.py | 8 +- .../velocity/config/g1/__init__.py | 12 +- .../velocity/config/go1/__init__.py | 8 +- .../velocity/config/go2/__init__.py | 8 +- .../velocity/config/h1/__init__.py | 12 +- .../velocity/mdp/__init__.py | 0 .../velocity/mdp/__init__.pyi | 0 .../velocity/mdp/curriculums.py | 0 .../velocity/mdp/rewards.py | 0 .../velocity/mdp/terminations.py | 0 .../direct/__init__.py | 10 - .../direct/cartpole/__init__.py | 29 -- .../manager_based/__init__.py | 10 - .../manager_based/classic/__init__.py | 12 - .../manager_based/classic/ant/ant_env_cfg.py | 197 ------------ .../classic/cartpole/__init__.py | 30 -- .../classic/cartpole/cartpole_env_cfg.py | 199 ------------ .../classic/humanoid/humanoid_env_cfg.py | 232 -------------- .../manager_based/locomotion/__init__.py | 6 - .../velocity/config/a1/flat_env_cfg.py | 60 ---- .../velocity/config/a1/rough_env_cfg.py | 91 ------ .../velocity/config/anymal_b/flat_env_cfg.py | 60 ---- .../velocity/config/anymal_b/rough_env_cfg.py | 34 -- .../velocity/config/anymal_c/flat_env_cfg.py | 52 --- .../velocity/config/anymal_c/rough_env_cfg.py | 40 --- .../velocity/config/anymal_d/flat_env_cfg.py | 46 --- .../velocity/config/anymal_d/rough_env_cfg.py | 37 --- .../velocity/config/cassie/flat_env_cfg.py | 45 --- .../velocity/config/cassie/rough_env_cfg.py | 96 ------ .../velocity/config/g1/flat_env_cfg.py | 58 ---- .../velocity/config/g1/rough_env_cfg.py | 178 ----------- .../velocity/config/go1/flat_env_cfg.py | 46 --- .../velocity/config/go1/rough_env_cfg.py | 61 ---- .../velocity/config/go2/flat_env_cfg.py | 46 --- .../velocity/config/go2/rough_env_cfg.py | 60 ---- .../velocity/config/h1/flat_env_cfg.py | 46 --- .../velocity/config/h1/rough_env_cfg.py | 131 -------- .../locomotion/velocity/velocity_env_cfg.py | 296 ------------------ .../manager_based/manipulation/__init__.py | 6 - .../reach/config/franka/joint_pos_env_cfg.py | 74 ----- .../reach/config/ur_10/joint_pos_env_cfg.py | 74 ----- .../manipulation/reach/reach_env_cfg.py | 206 ------------ 80 files changed, 177 insertions(+), 2630 deletions(-) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct => core}/allegro_hand/__init__.py (92%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct => core}/allegro_hand/allegro_hand_warp_env_cfg.py (100%) create mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct => core}/cartpole/cartpole_warp_env.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct => core}/cartpole/cartpole_warp_env_cfg.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic => core}/cartpole/mdp/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic => core}/cartpole/mdp/__init__.pyi (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic => core}/cartpole/mdp/rewards.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct/ant => core/direct_ant}/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct/ant => core/direct_ant}/ant_env_warp.py (82%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct/ant => core/direct_ant}/ant_env_warp_cfg.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct/humanoid => core/direct_humanoid}/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct/humanoid => core/direct_humanoid}/humanoid_warp_env.py (83%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct/humanoid => core/direct_humanoid}/humanoid_warp_env_cfg.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct/inhand_manipulation => core/direct_locomotion}/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct/locomotion => core/direct_locomotion}/locomotion_env_warp.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct/locomotion => core/inhand_manipulation}/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct => core}/inhand_manipulation/inhand_manipulation_warp_env.py (99%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic/ant => core/manager_ant}/__init__.py (77%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic/humanoid => core/manager_humanoid}/__init__.py (76%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic/humanoid => core/manager_humanoid}/mdp/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic/humanoid => core/manager_humanoid}/mdp/__init__.pyi (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic/humanoid => core/manager_humanoid}/mdp/observations.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic/humanoid => core/manager_humanoid}/mdp/rewards.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/manipulation => core}/reach/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/manipulation => core}/reach/config/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/manipulation => core}/reach/config/franka/__init__.py (76%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/manipulation => core}/reach/config/ur_10/__init__.py (79%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/manipulation => core}/reach/mdp/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/manipulation => core}/reach/mdp/__init__.pyi (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/manipulation => core}/reach/mdp/rewards.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/a1/__init__.py (77%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/anymal_b/__init__.py (78%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/anymal_c/__init__.py (80%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/anymal_d/__init__.py (79%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/cassie/__init__.py (74%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/g1/__init__.py (77%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/go1/__init__.py (74%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/go2/__init__.py (74%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/h1/__init__.py (77%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/mdp/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/mdp/__init__.pyi (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/mdp/curriculums.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/mdp/rewards.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/mdp/terminations.py (100%) delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/cartpole_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/joint_pos_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/joint_pos_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index 2f186d3617d8..2b85e474da8f 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -214,12 +214,15 @@ def _swap_mdp(cfg: Any, label: str) -> None: Iterates every term cfg in the tree (via :func:`_walk_terms`) and on each term swaps whichever of ``func`` / ``class_type`` is a stable-origin - callable. Twin lookup is name-based against the task's matching - ``isaaclab_tasks_experimental.<...>.mdp`` module and, as a fallback, - :mod:`isaaclab_experimental.envs.mdp`. Any missing twin raises - :class:`FrontendIncompatibleError` listing every affected term — partial - swaps would leave torch-style callables in the cfg and the warp managers - would call them with the wrong signature. + callable. Twin lookup is name-based against the warp mirror of the *stable + symbol's own module* (:func:`_warp_mdp_modules`) plus the + :mod:`isaaclab_experimental.envs.mdp` fallback. Keying off the symbol's + module — not the cfg's — means a task that borrows another task's MDP (e.g. + ``manager_ant`` reuses ``manager_humanoid.mdp``) resolves to the right warp + twins without a per-task shim, mirroring the stable ``core/`` layout. Any + missing twin raises :class:`FrontendIncompatibleError` listing every + affected term — partial swaps would leave torch-style callables in the cfg + and the warp managers would call them with the wrong signature. The warp-side side declarations (``out_dim``, ``axes``, ``observation_type``) that the warp managers need at init are *not* supplied by this swap; they @@ -228,9 +231,8 @@ def _swap_mdp(cfg: Any, label: str) -> None: substitutes the callable; the manager reads the new func's annotations when it parses the term cfg. """ - modules = _warp_mdp_modules(type(cfg).__module__) - searched = tuple(m.__name__ for m in modules) - logger.info("frontend.warp: searching warp mdp modules %s", list(searched)) + module_cache: dict[str, list[ModuleType]] = {} + searched: set[str] = set() swapped = 0 missing: list[tuple[str, str, str]] = [] # (location, attr, symbol) @@ -240,7 +242,11 @@ def _swap_mdp(cfg: Any, label: str) -> None: stable = getattr(term, attr, None) if stable is None or not _is_swap_candidate(stable): continue - twin = _resolve_warp_twin(stable.__name__, modules) + origin = getattr(stable, "__module__", "") or "" + if origin not in module_cache: + module_cache[origin] = _warp_mdp_modules(origin) + searched.update(m.__name__ for m in module_cache[origin]) + twin = _resolve_warp_twin(stable.__name__, module_cache[origin]) if twin is None: missing.append((location, attr, stable.__name__)) continue @@ -250,7 +256,7 @@ def _swap_mdp(cfg: Any, label: str) -> None: if missing: lines = "\n ".join(f"{loc}.{attr}: no warp twin for {sym!r}" for loc, attr, sym in missing) raise FrontendIncompatibleError( - f"warp env {label!r}: missing warp MDP twins (searched {list(searched)}):\n {lines}" + f"warp env {label!r}: missing warp MDP twins (searched {sorted(searched)}):\n {lines}" ) logger.info("frontend.warp: swapped %d MDP symbol(s) to warp twins", swapped) @@ -295,36 +301,40 @@ def _assert_direct_warp_registration(task_id: str) -> None: ) -def _warp_mdp_modules(cfg_module: str) -> list[ModuleType]: - """Locate warp MDP modules to consult for twin lookups. +def _warp_mdp_modules(symbol_module: str) -> list[ModuleType]: + """Locate warp MDP modules to consult for a stable symbol's twin. - The lookup is keyed off the *stable* cfg class's module - (``type(cfg).__module__``) rather than the gym task id, so it resolves the - same twins whether the env is built via ``--frontend=warp`` or constructed - directly from a ``*-Warp-v0`` registration. + The lookup is keyed off the *stable symbol's own module* + (``func.__module__`` / ``class_type.__module__``), e.g. + ``isaaclab_tasks.core.manager_humanoid.mdp.rewards``. We mirror that path + into the experimental package and walk up to the nearest importable module, + then append the shared fallback. Keying off the symbol's module — not the + cfg's — resolves the right twins even when a task borrows another task's MDP + (``manager_ant`` reuses ``manager_humanoid.mdp``), so the experimental + package can mirror the stable ``core/`` layout one-to-one without per-task + re-export shims. Order of preference: - 1. The task-specific module, derived by replacing ``isaaclab_tasks`` with - ``isaaclab_tasks_experimental`` in ``cfg_module`` and walking up to the - first existing ``.mdp`` submodule. This relies on the experimental - package mirroring the stable ``core/`` layout. + 1. The warp mirror of ``symbol_module`` (or its nearest importable ancestor, + e.g. the ``...mdp`` package when the exact submodule isn't mirrored). 2. The shared :mod:`isaaclab_experimental.envs.mdp` fallback (where generic warp twins live). """ modules: list[ModuleType] = [] - if isinstance(cfg_module, str) and cfg_module.startswith("isaaclab_tasks."): - warp_pkg = cfg_module.replace("isaaclab_tasks", "isaaclab_tasks_experimental", 1) - parts = warp_pkg.split(".") + if isinstance(symbol_module, str) and symbol_module.startswith("isaaclab_tasks."): + warp_mod = symbol_module.replace("isaaclab_tasks", "isaaclab_tasks_experimental", 1) + parts = warp_mod.split(".") for depth in range(len(parts), 0, -1): - target = ".".join(parts[:depth] + ["mdp"]) + target = ".".join(parts[:depth]) try: modules.append(importlib.import_module(target)) break except ModuleNotFoundError as exc: - # Only swallow "this candidate doesn't exist" — a real - # import error inside an existing module is a bug. - if exc.name == target: + # Keep walking up while the missing module is part of the path + # we're probing; a genuine import error inside an existing + # module (missing third-party dep, etc.) must surface. + if exc.name and warp_mod.startswith(exc.name): continue raise # Generic fallback. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/allegro_hand/__init__.py similarity index 92% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/allegro_hand/__init__.py index 95e9aafb7800..2bd91e74598b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/allegro_hand/__init__.py @@ -13,7 +13,7 @@ # Register Gym environments. ## -inhand_task_entry = "isaaclab_tasks_experimental.direct.inhand_manipulation" +inhand_task_entry = "isaaclab_tasks_experimental.core.inhand_manipulation" stable_agents = "isaaclab_tasks.core.allegro_hand.agents" gym.register( diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/allegro_hand_warp_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/allegro_hand/allegro_hand_warp_env_cfg.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/allegro_hand_warp_env_cfg.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/allegro_hand/allegro_hand_warp_env_cfg.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py new file mode 100644 index 000000000000..d5d93b3dba48 --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py @@ -0,0 +1,53 @@ +# 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 + +""" +Cartpole balancing environment (warp entry points). + +Mirrors stable ``isaaclab_tasks.core.cartpole``: the manager-based warp task +reuses the stable cfg (no parallel copy), while the direct warp task keeps its +own env class + cfg here since direct envs encode behaviour in the class. +""" + +import gymnasium as gym + +# Reuse agent configs from the stable task package. +from isaaclab_tasks.core.cartpole import agents + +# Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) +# and MDP twins are swapped at construction (see adapt_cfg_for_warp). +_stable_pkg = agents.__name__.rsplit(".", 1)[0] + +## +# Register Gym environments. +## + +# Manager-based: stable cfg adapted to warp at construction. +gym.register( + id="Isaac-Cartpole-Warp-v0", + entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{_stable_pkg}.cartpole_manager_env_cfg:CartpoleEnvCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_manager_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_manager_ppo_cfg:CartpolePPORunnerCfg", + "skrl_cfg_entry_point": f"{agents.__name__}:skrl_manager_ppo_cfg.yaml", + "sb3_cfg_entry_point": f"{agents.__name__}:sb3_manager_ppo_cfg.yaml", + }, +) + +# Direct: dedicated warp env class + local cfg (not adaptable from a cfg). +gym.register( + id="Isaac-Cartpole-Direct-Warp-v0", + entry_point=f"{__name__}.cartpole_warp_env:CartpoleWarpEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.cartpole_warp_env_cfg:CartpoleWarpEnvCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_direct_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_direct_ppo_cfg:CartpolePPORunnerCfg", + "skrl_cfg_entry_point": f"{agents.__name__}:skrl_direct_ppo_cfg.yaml", + "sb3_cfg_entry_point": f"{agents.__name__}:sb3_direct_ppo_cfg.yaml", + }, +) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env_cfg.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env_cfg.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env_cfg.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.pyi similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/__init__.pyi rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.pyi diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_ant/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_ant/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_ant/ant_env_warp.py similarity index 82% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_ant/ant_env_warp.py index 2bd7390039b7..d9eba6dfeba1 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_ant/ant_env_warp.py @@ -5,7 +5,7 @@ from __future__ import annotations -from isaaclab_tasks_experimental.direct.locomotion.locomotion_env_warp import LocomotionWarpEnv +from isaaclab_tasks_experimental.core.direct_locomotion.locomotion_env_warp import LocomotionWarpEnv from .ant_env_warp_cfg import AntWarpEnvCfg diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_ant/ant_env_warp_cfg.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp_cfg.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_ant/ant_env_warp_cfg.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_humanoid/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_humanoid/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_humanoid/humanoid_warp_env.py similarity index 83% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_humanoid/humanoid_warp_env.py index 2fd0ac8433e2..f704a6a8859a 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_humanoid/humanoid_warp_env.py @@ -5,7 +5,7 @@ from __future__ import annotations -from isaaclab_tasks_experimental.direct.locomotion.locomotion_env_warp import LocomotionWarpEnv +from isaaclab_tasks_experimental.core.direct_locomotion.locomotion_env_warp import LocomotionWarpEnv from .humanoid_warp_env_cfg import HumanoidWarpEnvCfg diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_humanoid/humanoid_warp_env_cfg.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env_cfg.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_humanoid/humanoid_warp_env_cfg.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_locomotion/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_locomotion/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_locomotion/locomotion_env_warp.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_locomotion/locomotion_env_warp.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/inhand_manipulation/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/inhand_manipulation/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/inhand_manipulation/inhand_manipulation_warp_env.py similarity index 99% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/inhand_manipulation/inhand_manipulation_warp_env.py index d5bc6a2ad529..b3f1f4abb789 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/inhand_manipulation/inhand_manipulation_warp_env.py @@ -19,7 +19,7 @@ from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane if TYPE_CHECKING: - from isaaclab_tasks_experimental.direct.allegro_hand.allegro_hand_warp_env_cfg import AllegroHandWarpEnvCfg + from isaaclab_tasks_experimental.core.allegro_hand.allegro_hand_warp_env_cfg import AllegroHandWarpEnvCfg @wp.kernel diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_ant/__init__.py similarity index 77% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_ant/__init__.py index 643996b7f730..26996b4881f9 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_ant/__init__.py @@ -12,6 +12,10 @@ # Reuse agent configs from the stable task package. from isaaclab_tasks.core.manager_ant import agents +# Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) +# and MDP twins are swapped at construction (see adapt_cfg_for_warp). +_stable_pkg = agents.__name__.rsplit(".", 1)[0] + ## # Register Gym environments. ## @@ -21,7 +25,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.ant_env_cfg:AntEnvCfg", + "env_cfg_entry_point": f"{_stable_pkg}.ant_env_cfg:AntEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AntPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/__init__.py similarity index 76% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/__init__.py index c1d30e76a9e2..183b42d2aaeb 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/__init__.py @@ -12,6 +12,10 @@ # Reuse agent configs from the stable task package. from isaaclab_tasks.core.manager_humanoid import agents +# Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) +# and MDP twins are swapped at construction (see adapt_cfg_for_warp). +_stable_pkg = agents.__name__.rsplit(".", 1)[0] + ## # Register Gym environments. ## @@ -21,7 +25,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.humanoid_env_cfg:HumanoidEnvCfg", + "env_cfg_entry_point": f"{_stable_pkg}.humanoid_env_cfg:HumanoidEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:HumanoidPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/mdp/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/mdp/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/mdp/__init__.pyi similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/__init__.pyi rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/mdp/__init__.pyi diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/observations.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/mdp/observations.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/observations.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/mdp/observations.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/mdp/rewards.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/mdp/rewards.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/franka/__init__.py similarity index 76% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/franka/__init__.py index 9f01c7b2e729..a82f977588fd 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/franka/__init__.py @@ -8,6 +8,10 @@ # Reuse agent configs from the stable task package. from isaaclab_tasks.core.reach.config.franka import agents +# Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) +# and MDP twins are swapped at construction (see adapt_cfg_for_warp). +_stable_pkg = agents.__name__.rsplit(".", 1)[0] + ## # Register Gym environments. ## @@ -21,7 +25,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaReachEnvCfg", + "env_cfg_entry_point": f"{_stable_pkg}.joint_pos_env_cfg:FrankaReachEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", @@ -33,7 +37,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaReachEnvCfg_PLAY", + "env_cfg_entry_point": f"{_stable_pkg}.joint_pos_env_cfg:FrankaReachEnvCfg_PLAY", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/ur_10/__init__.py similarity index 79% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/ur_10/__init__.py index cdff71bf6523..03843620488d 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/ur_10/__init__.py @@ -11,12 +11,16 @@ # import gymnasium as gym # from isaaclab_tasks.core.reach.config.ur_10 import agents +# Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) +# and MDP twins are swapped at construction (see adapt_cfg_for_warp). +_stable_pkg = agents.__name__.rsplit(".", 1)[0] + # gym.register( # id="Isaac-Reach-UR10-Warp-v0", # entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", # disable_env_checker=True, # kwargs={ -# "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:UR10ReachEnvCfg", +# "env_cfg_entry_point": f"{_stable_pkg}.joint_pos_env_cfg:UR10ReachEnvCfg", # "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", # "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UR10ReachPPORunnerCfg", # "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", @@ -28,7 +32,7 @@ # entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", # disable_env_checker=True, # kwargs={ -# "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:UR10ReachEnvCfg_PLAY", +# "env_cfg_entry_point": f"{_stable_pkg}.joint_pos_env_cfg:UR10ReachEnvCfg_PLAY", # "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", # "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UR10ReachPPORunnerCfg", # "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.pyi similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.pyi rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.pyi diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/__init__.py similarity index 77% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/__init__.py index 598d37c28c73..92ca6742257f 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/__init__.py @@ -8,6 +8,10 @@ # Reuse agent configs from the stable task package. from isaaclab_tasks.contrib.velocity.config.a1 import agents +# Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) +# and MDP twins are swapped at construction (see adapt_cfg_for_warp). +_stable_pkg = agents.__name__.rsplit(".", 1)[0] + ## # Register Gym environments. ## @@ -17,7 +21,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeA1FlatEnvCfg", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:UnitreeA1FlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeA1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml", @@ -29,7 +33,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeA1FlatEnvCfg_PLAY", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:UnitreeA1FlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeA1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/__init__.py similarity index 78% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/__init__.py index e33364fd7cff..fcb1393f5092 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/__init__.py @@ -8,6 +8,10 @@ # Reuse agent configs from the stable task package. from isaaclab_tasks.contrib.velocity.config.anymal_b import agents +# Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) +# and MDP twins are swapped at construction (see adapt_cfg_for_warp). +_stable_pkg = agents.__name__.rsplit(".", 1)[0] + ## # Register Gym environments. ## @@ -17,7 +21,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalBFlatEnvCfg", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:AnymalBFlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBFlatPPORunnerCfg", "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBFlatPPORunnerWithSymmetryCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", @@ -29,7 +33,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalBFlatEnvCfg_PLAY", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:AnymalBFlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBFlatPPORunnerCfg", "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBFlatPPORunnerWithSymmetryCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/__init__.py similarity index 80% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/__init__.py index 58f71adc0d06..7ee48552d0d2 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/__init__.py @@ -8,6 +8,10 @@ # Reuse agent configs from the stable task package. from isaaclab_tasks.contrib.velocity.config.anymal_c import agents +# Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) +# and MDP twins are swapped at construction (see adapt_cfg_for_warp). +_stable_pkg = agents.__name__.rsplit(".", 1)[0] + ## # Register Gym environments. ## @@ -17,7 +21,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalCFlatEnvCfg", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:AnymalCFlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerCfg", "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerWithSymmetryCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_flat_ppo_cfg.yaml", @@ -30,7 +34,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalCFlatEnvCfg_PLAY", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:AnymalCFlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerCfg", "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerWithSymmetryCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_flat_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/__init__.py similarity index 79% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/__init__.py index 20dda722e09c..611ec20f9540 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/__init__.py @@ -8,6 +8,10 @@ # Reuse agent configs from the stable task package. from isaaclab_tasks.core.velocity.config.anymal_d import agents +# Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) +# and MDP twins are swapped at construction (see adapt_cfg_for_warp). +_stable_pkg = agents.__name__.rsplit(".", 1)[0] + ## # Register Gym environments. ## @@ -20,7 +24,7 @@ # entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", # disable_env_checker=True, # kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:AnymalDRoughEnvCfg", +# "env_cfg_entry_point": f"{_stable_pkg}.rough_env_cfg:AnymalDRoughEnvCfg", # "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDRoughPPORunnerCfg", # "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", # }, @@ -31,7 +35,7 @@ # entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", # disable_env_checker=True, # kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:AnymalDRoughEnvCfg_PLAY", +# "env_cfg_entry_point": f"{_stable_pkg}.rough_env_cfg:AnymalDRoughEnvCfg_PLAY", # "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDRoughPPORunnerCfg", # "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", # }, @@ -42,7 +46,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalDFlatEnvCfg", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:AnymalDFlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -53,7 +57,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalDFlatEnvCfg_PLAY", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:AnymalDFlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/__init__.py similarity index 74% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/__init__.py index e3af41334541..db4c9858b764 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/__init__.py @@ -8,6 +8,10 @@ # Reuse agent configs from the stable task package. from isaaclab_tasks.core.velocity.config.cassie import agents +# Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) +# and MDP twins are swapped at construction (see adapt_cfg_for_warp). +_stable_pkg = agents.__name__.rsplit(".", 1)[0] + ## # Register Gym environments. ## @@ -17,7 +21,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:CassieFlatEnvCfg", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:CassieFlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CassieFlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -28,7 +32,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:CassieFlatEnvCfg_PLAY", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:CassieFlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CassieFlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/__init__.py similarity index 77% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/__init__.py index f49e103f431a..5a9e6bfa9e9c 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/__init__.py @@ -8,6 +8,10 @@ # Reuse agent configs from the stable task package. from isaaclab_tasks.core.velocity.config.g1 import agents +# Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) +# and MDP twins are swapped at construction (see adapt_cfg_for_warp). +_stable_pkg = agents.__name__.rsplit(".", 1)[0] + ## # Register Gym environments. ## @@ -17,7 +21,7 @@ # entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", # disable_env_checker=True, # kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:G1RoughEnvCfg", +# "env_cfg_entry_point": f"{_stable_pkg}.rough_env_cfg:G1RoughEnvCfg", # "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1RoughPPORunnerCfg", # "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", # }, @@ -29,7 +33,7 @@ # entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", # disable_env_checker=True, # kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:G1RoughEnvCfg_PLAY", +# "env_cfg_entry_point": f"{_stable_pkg}.rough_env_cfg:G1RoughEnvCfg_PLAY", # "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1RoughPPORunnerCfg", # "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", # }, @@ -40,7 +44,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:G1FlatEnvCfg", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:G1FlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -51,7 +55,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:G1FlatEnvCfg_PLAY", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:G1FlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/__init__.py similarity index 74% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/__init__.py index 3206b985ea0e..bbdb15b7da88 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/__init__.py @@ -8,6 +8,10 @@ # Reuse agent configs from the stable task package. from isaaclab_tasks.contrib.velocity.config.go1 import agents +# Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) +# and MDP twins are swapped at construction (see adapt_cfg_for_warp). +_stable_pkg = agents.__name__.rsplit(".", 1)[0] + ## # Register Gym environments. ## @@ -17,7 +21,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo1FlatEnvCfg", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:UnitreeGo1FlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -28,7 +32,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo1FlatEnvCfg_PLAY", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:UnitreeGo1FlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/__init__.py similarity index 74% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/__init__.py index 041579201f35..0e0d42e604ba 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/__init__.py @@ -8,6 +8,10 @@ # Reuse agent configs from the stable task package. from isaaclab_tasks.core.velocity.config.go2 import agents +# Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) +# and MDP twins are swapped at construction (see adapt_cfg_for_warp). +_stable_pkg = agents.__name__.rsplit(".", 1)[0] + ## # Register Gym environments. ## @@ -17,7 +21,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo2FlatEnvCfg", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:UnitreeGo2FlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo2FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -28,7 +32,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo2FlatEnvCfg_PLAY", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:UnitreeGo2FlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo2FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/__init__.py similarity index 77% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/__init__.py index c731af7ca986..2c1de23e1679 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/__init__.py @@ -8,6 +8,10 @@ # Reuse agent configs from the stable task package. from isaaclab_tasks.core.velocity.config.h1 import agents +# Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) +# and MDP twins are swapped at construction (see adapt_cfg_for_warp). +_stable_pkg = agents.__name__.rsplit(".", 1)[0] + ## # Register Gym environments. ## @@ -17,7 +21,7 @@ # entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", # disable_env_checker=True, # kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:H1RoughEnvCfg", +# "env_cfg_entry_point": f"{_stable_pkg}.rough_env_cfg:H1RoughEnvCfg", # "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1RoughPPORunnerCfg", # "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", # }, @@ -28,7 +32,7 @@ # entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", # disable_env_checker=True, # kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:H1RoughEnvCfg_PLAY", +# "env_cfg_entry_point": f"{_stable_pkg}.rough_env_cfg:H1RoughEnvCfg_PLAY", # "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1RoughPPORunnerCfg", # "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", # }, @@ -39,7 +43,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:H1FlatEnvCfg", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:H1FlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -50,7 +54,7 @@ entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:H1FlatEnvCfg_PLAY", + "env_cfg_entry_point": f"{_stable_pkg}.flat_env_cfg:H1FlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.pyi rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/terminations.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/terminations.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/__init__.py deleted file mode 100644 index 3e2b7945ebde..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# 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 - -""" -Direct workflow environments. -""" - -import gymnasium as gym diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py deleted file mode 100644 index 48021a84f4d5..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -# 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 - -""" -Cartpole balancing environment. -""" - -import gymnasium as gym - -## -# Register Gym environments. -## - -stable_agents = "isaaclab_tasks.core.cartpole.agents" - -gym.register( - id="Isaac-Cartpole-Direct-Warp-v0", - entry_point=f"{__name__}.cartpole_warp_env:CartpoleWarpEnv", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.cartpole_warp_env_cfg:CartpoleWarpEnvCfg", - "rl_games_cfg_entry_point": f"{stable_agents}:rl_games_direct_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{stable_agents}.rsl_rl_direct_ppo_cfg:CartpolePPORunnerCfg", - "skrl_cfg_entry_point": f"{stable_agents}:skrl_direct_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{stable_agents}:sb3_direct_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/__init__.py deleted file mode 100644 index 7f23883e6332..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# 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 - -"""Experimental registrations for manager-based tasks. - -We intentionally only register new Gym IDs pointing at experimental entry points. -Task definitions (configs/mdp) remain in `isaaclab_tasks` to avoid duplication. -""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/__init__.py deleted file mode 100644 index 79c13e2aa8f4..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# 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 - -"""Classic environments for control. - -These environments are based on the MuJoCo environments provided by OpenAI. - -Reference: - https://github.com/openai/gym/tree/master/gym/envs/mujoco -""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py deleted file mode 100644 index 66a9131011f3..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py +++ /dev/null @@ -1,197 +0,0 @@ -# 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 - -# Ant reuses humanoid's experimental MDP (mirrors stable pattern). -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.classic.humanoid.mdp as mdp - -## -# Pre-defined configs -## -from isaaclab_assets.robots.ant import ANT_CFG # isort: skip - - -@configclass -class MySceneCfg(InteractiveSceneCfg): - """Configuration for the terrain scene with an ant robot.""" - - # terrain - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="plane", - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg( - friction_combine_mode="average", - restitution_combine_mode="average", - static_friction=1.0, - dynamic_friction=1.0, - restitution=0.0, - ), - debug_vis=False, - ) - - # robot - robot = ANT_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - - # lights - light = AssetBaseCfg( - prim_path="/World/light", - spawn=sim_utils.DistantLightCfg(color=(0.75, 0.75, 0.75), intensity=3000.0), - ) - - -## -# MDP settings -## - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - joint_effort = mdp.JointEffortActionCfg(asset_name="robot", joint_names=[".*"], scale=7.5) - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - @configclass - class PolicyCfg(ObsGroup): - """Observations for the policy.""" - - base_height = ObsTerm(func=mdp.base_pos_z) - base_lin_vel = ObsTerm(func=mdp.base_lin_vel) - base_ang_vel = ObsTerm(func=mdp.base_ang_vel) - base_yaw_roll = ObsTerm(func=mdp.base_yaw_roll) - base_angle_to_target = ObsTerm(func=mdp.base_angle_to_target, params={"target_pos": (1000.0, 0.0, 0.0)}) - base_up_proj = ObsTerm(func=mdp.base_up_proj) - base_heading_proj = ObsTerm(func=mdp.base_heading_proj, params={"target_pos": (1000.0, 0.0, 0.0)}) - joint_pos_norm = ObsTerm(func=mdp.joint_pos_limit_normalized) - joint_vel_rel = ObsTerm(func=mdp.joint_vel_rel, scale=0.2) - actions = ObsTerm(func=mdp.last_action) - - def __post_init__(self): - self.enable_corruption = False - self.concatenate_terms = True - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - reset_base = EventTerm( - func=mdp.reset_root_state_uniform, - mode="reset", - params={"pose_range": {}, "velocity_range": {}}, - ) - - reset_robot_joints = EventTerm( - func=mdp.reset_joints_by_offset, - mode="reset", - params={ - "position_range": (-0.2, 0.2), - "velocity_range": (-0.1, 0.1), - }, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # (1) Reward for moving forward - progress = RewTerm(func=mdp.progress_reward, weight=1.0, params={"target_pos": (1000.0, 0.0, 0.0)}) - # (2) Stay alive bonus - alive = RewTerm(func=mdp.is_alive, weight=0.5) - # (3) Reward for non-upright posture - upright = RewTerm(func=mdp.upright_posture_bonus, weight=0.1, params={"threshold": 0.93}) - # (4) Reward for moving in the right direction - move_to_target = RewTerm( - func=mdp.move_to_target_bonus, weight=0.5, params={"threshold": 0.8, "target_pos": (1000.0, 0.0, 0.0)} - ) - # (5) Penalty for large action commands - action_l2 = RewTerm(func=mdp.action_l2, weight=-0.005) - # (6) Penalty for energy consumption - energy = RewTerm(func=mdp.power_consumption, weight=-0.05, params={"gear_ratio": {".*": 15.0}}) - # (7) Penalty for reaching close to joint limits - joint_pos_limits = RewTerm( - func=mdp.joint_pos_limits_penalty_ratio, weight=-0.1, params={"threshold": 0.99, "gear_ratio": {".*": 15.0}} - ) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - # (1) Terminate if the episode length is exceeded - time_out = DoneTerm(func=mdp.time_out, time_out=True) - # (2) Terminate if the robot falls - torso_height = DoneTerm(func=mdp.root_height_below_minimum, params={"minimum_height": 0.31}) - - -## -# Environment configuration -## - - -@configclass -class AntEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the MuJoCo-style Ant walking environment.""" - - # Simulation settings - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=38, - nconmax=15, - ls_iterations=10, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - # Scene settings - scene: MySceneCfg = MySceneCfg(num_envs=4096, env_spacing=5.0, clone_in_fabric=True) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() - - def __post_init__(self): - """Post initialization.""" - # general settings - self.decimation = 2 - self.episode_length_s = 16.0 - # simulation settings - self.sim.dt = 1 / 120.0 - self.sim.render_interval = self.decimation - # default friction material - self.sim.physics_material.static_friction = 1.0 - self.sim.physics_material.dynamic_friction = 1.0 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py deleted file mode 100644 index 9f46551de255..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py +++ /dev/null @@ -1,30 +0,0 @@ -# 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 - -""" -Cartpole balancing environment (experimental manager-based entry point). -""" - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.cartpole import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Cartpole-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.cartpole_env_cfg:CartpoleEnvCfg", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_manager_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_manager_ppo_cfg:CartpolePPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_manager_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{agents.__name__}:sb3_manager_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/cartpole_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/cartpole_env_cfg.py deleted file mode 100644 index ae187a593a16..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/cartpole_env_cfg.py +++ /dev/null @@ -1,199 +0,0 @@ -# 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 - -import math - -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg, AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.classic.cartpole.mdp as mdp - -## -# Pre-defined configs -## -from isaaclab_assets.robots.cartpole import CARTPOLE_CFG # isort:skip - - -## -# Scene definition -## - - -@configclass -class CartpoleSceneCfg(InteractiveSceneCfg): - """Configuration for a cart-pole scene.""" - - # ground plane - # ground = AssetBaseCfg( - # prim_path="/World/ground", - # spawn=sim_utils.GroundPlaneCfg(size=(100.0, 100.0)), - # ) - - # cartpole - robot: ArticulationCfg = CARTPOLE_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - - # lights - dome_light = AssetBaseCfg( - prim_path="/World/DomeLight", - spawn=sim_utils.DomeLightCfg(color=(0.9, 0.9, 0.9), intensity=500.0), - ) - - -## -# MDP settings -## - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - joint_effort = mdp.JointEffortActionCfg(asset_name="robot", joint_names=["slider_to_cart"], scale=100.0) - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - @configclass - class PolicyCfg(ObsGroup): - """Observations for policy group.""" - - # observation terms (order preserved) - joint_pos_rel = ObsTerm(func=mdp.joint_pos_rel) - joint_vel_rel = ObsTerm(func=mdp.joint_vel_rel) - - def __post_init__(self) -> None: - self.enable_corruption = False - self.concatenate_terms = True - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - # reset - reset_cart_position = EventTerm( - func=mdp.reset_joints_by_offset, - mode="reset", - params={ - "asset_cfg": SceneEntityCfg("robot", joint_names=["slider_to_cart"]), - "position_range": (-1.0, 1.0), - "velocity_range": (-0.5, 0.5), - }, - ) - - reset_pole_position = EventTerm( - func=mdp.reset_joints_by_offset, - mode="reset", - params={ - "asset_cfg": SceneEntityCfg("robot", joint_names=["cart_to_pole"]), - "position_range": (-0.25 * math.pi, 0.25 * math.pi), - "velocity_range": (-0.25 * math.pi, 0.25 * math.pi), - }, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # (1) Constant running reward - alive = RewTerm(func=mdp.is_alive, weight=1.0) - # (2) Failure penalty - terminating = RewTerm(func=mdp.is_terminated, weight=-2.0) - # (3) Primary task: keep pole upright - pole_pos = RewTerm( - func=mdp.joint_pos_target_l2, - weight=-1.0, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["cart_to_pole"]), "target": 0.0}, - ) - # (4) Shaping tasks: lower cart velocity - cart_vel = RewTerm( - func=mdp.joint_vel_l1, - weight=-0.01, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["slider_to_cart"])}, - ) - # (5) Shaping tasks: lower pole angular velocity - pole_vel = RewTerm( - func=mdp.joint_vel_l1, - weight=-0.005, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["cart_to_pole"])}, - ) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - # (1) Time out - time_out = DoneTerm(func=mdp.time_out, time_out=True) - # (2) Cart out of bounds - cart_out_of_bounds = DoneTerm( - func=mdp.joint_pos_out_of_manual_limit, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["slider_to_cart"]), "bounds": (-3.0, 3.0)}, - ) - - -## -# Environment configuration -## - - -@configclass -class CartpoleEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the cartpole environment.""" - - # Scene settings - scene: CartpoleSceneCfg = CartpoleSceneCfg(num_envs=4096, env_spacing=4.0, clone_in_fabric=True) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - events: EventCfg = EventCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - # Simulation settings - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=5, - nconmax=3, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - use_cuda_graph=True, - ) - ) - - # Post initialization - def __post_init__(self) -> None: - """Post initialization.""" - # general settings - self.decimation = 2 - self.episode_length_s = 5 - # viewer settings - self.viewer.eye = (8.0, 0.0, 5.0) - # simulation settings - self.sim.dt = 1 / 120 - self.sim.render_interval = self.decimation diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py deleted file mode 100644 index 9de689ffb98e..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py +++ /dev/null @@ -1,232 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.classic.humanoid.mdp as mdp - -from isaaclab_assets.robots.humanoid import HUMANOID_CFG # isort:skip - - -## -# Scene definition -## - - -@configclass -class MySceneCfg(InteractiveSceneCfg): - """Configuration for the terrain scene with a humanoid robot.""" - - # terrain - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="plane", - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg(static_friction=1.0, dynamic_friction=1.0, restitution=0.0), - debug_vis=False, - ) - - # robot - robot = HUMANOID_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - - # lights - light = AssetBaseCfg( - prim_path="/World/light", - spawn=sim_utils.DistantLightCfg(color=(0.75, 0.75, 0.75), intensity=3000.0), - ) - - -## -# MDP settings -## - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - joint_effort = mdp.JointEffortActionCfg( - asset_name="robot", - joint_names=[".*"], - scale={ - ".*_waist.*": 67.5, - ".*_upper_arm.*": 67.5, - "pelvis": 67.5, - ".*_lower_arm": 45.0, - ".*_thigh:0": 45.0, - ".*_thigh:1": 135.0, - ".*_thigh:2": 45.0, - ".*_shin": 90.0, - ".*_foot.*": 22.5, - }, - ) - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - @configclass - class PolicyCfg(ObsGroup): - """Observations for the policy.""" - - base_height = ObsTerm(func=mdp.base_pos_z) - base_lin_vel = ObsTerm(func=mdp.base_lin_vel) - base_ang_vel = ObsTerm(func=mdp.base_ang_vel, scale=0.25) - base_yaw_roll = ObsTerm(func=mdp.base_yaw_roll) - base_angle_to_target = ObsTerm(func=mdp.base_angle_to_target, params={"target_pos": (1000.0, 0.0, 0.0)}) - base_up_proj = ObsTerm(func=mdp.base_up_proj) - base_heading_proj = ObsTerm(func=mdp.base_heading_proj, params={"target_pos": (1000.0, 0.0, 0.0)}) - joint_pos_norm = ObsTerm(func=mdp.joint_pos_limit_normalized) - joint_vel_rel = ObsTerm(func=mdp.joint_vel_rel, scale=0.1) - actions = ObsTerm(func=mdp.last_action) - - def __post_init__(self): - self.enable_corruption = False - self.concatenate_terms = True - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - reset_base = EventTerm( - func=mdp.reset_root_state_uniform, - mode="reset", - params={"pose_range": {}, "velocity_range": {}}, - ) - - reset_robot_joints = EventTerm( - func=mdp.reset_joints_by_offset, - mode="reset", - params={ - "position_range": (-0.2, 0.2), - "velocity_range": (-0.1, 0.1), - }, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # (1) Reward for moving forward - progress = RewTerm(func=mdp.progress_reward, weight=1.0, params={"target_pos": (1000.0, 0.0, 0.0)}) - # (2) Stay alive bonus - alive = RewTerm(func=mdp.is_alive, weight=2.0) - # (3) Reward for non-upright posture - upright = RewTerm(func=mdp.upright_posture_bonus, weight=0.1, params={"threshold": 0.93}) - # (4) Reward for moving in the right direction - move_to_target = RewTerm( - func=mdp.move_to_target_bonus, weight=0.5, params={"threshold": 0.8, "target_pos": (1000.0, 0.0, 0.0)} - ) - # (5) Penalty for large action commands - action_l2 = RewTerm(func=mdp.action_l2, weight=-0.01) - # (6) Penalty for energy consumption - energy = RewTerm( - func=mdp.power_consumption, - weight=-0.005, - params={ - "gear_ratio": { - ".*_waist.*": 67.5, - ".*_upper_arm.*": 67.5, - "pelvis": 67.5, - ".*_lower_arm": 45.0, - ".*_thigh:0": 45.0, - ".*_thigh:1": 135.0, - ".*_thigh:2": 45.0, - ".*_shin": 90.0, - ".*_foot.*": 22.5, - } - }, - ) - # (7) Penalty for reaching close to joint limits - joint_pos_limits = RewTerm( - func=mdp.joint_pos_limits_penalty_ratio, - weight=-0.25, - params={ - "threshold": 0.98, - "gear_ratio": { - ".*_waist.*": 67.5, - ".*_upper_arm.*": 67.5, - "pelvis": 67.5, - ".*_lower_arm": 45.0, - ".*_thigh:0": 45.0, - ".*_thigh:1": 135.0, - ".*_thigh:2": 45.0, - ".*_shin": 90.0, - ".*_foot.*": 22.5, - }, - }, - ) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - # (1) Terminate if the episode length is exceeded - time_out = DoneTerm(func=mdp.time_out, time_out=True) - # (2) Terminate if the robot falls - torso_height = DoneTerm(func=mdp.root_height_below_minimum, params={"minimum_height": 0.8}) - - -@configclass -class HumanoidEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the MuJoCo-style Humanoid walking environment.""" - - # Scene settings - scene: MySceneCfg = MySceneCfg(num_envs=4096, env_spacing=5.0, clone_in_fabric=True) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() - - def __post_init__(self): - """Post initialization.""" - # general settings - self.decimation = 2 - self.episode_length_s = 16.0 - # simulation settings - self.sim: SimulationCfg = SimulationCfg( - dt=1 / 120.0, - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=80, - nconmax=25, - ls_iterations=15, - cone="pyramidal", - update_data_interval=2, - impratio=1, - integrator="implicitfast", - ), - num_substeps=2, - debug_mode=False, - ), - ) - # self.sim.dt = 1 / 120.0 - self.sim.render_interval = self.decimation - # default friction material - self.sim.physics_material.static_friction = 1.0 - self.sim.physics_material.dynamic_friction = 1.0 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/__init__.py deleted file mode 100644 index 0660d38f0658..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# 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 - -"""Locomotion experimental task registrations (manager-based).""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py deleted file mode 100644 index f213286d0cfb..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py +++ /dev/null @@ -1,60 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import UnitreeA1RoughEnvCfg - - -@configclass -class UnitreeA1FlatEnvCfg(UnitreeA1RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=60, - nconmax=30, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - # post init of parent - super().__post_init__() - - # override rewards - self.rewards.flat_orientation_l2.weight = -2.5 - self.rewards.feet_air_time.weight = 0.25 - - # change terrain to flat - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - # no height scan - # self.scene.height_scanner = None - # self.observations.policy.height_scan = None - # no terrain curriculum - self.curriculum.terrain_levels = None - - -class UnitreeA1FlatEnvCfg_PLAY(UnitreeA1FlatEnvCfg): - def __post_init__(self) -> None: - # post init of parent - super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing event - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/rough_env_cfg.py deleted file mode 100644 index ee0f236043a2..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/rough_env_cfg.py +++ /dev/null @@ -1,91 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm - -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import ( - LocomotionVelocityRoughEnvCfg, - TerminationsCfg, -) - -from isaaclab_assets.robots.unitree import UNITREE_A1_CFG # isort: skip - - -class TerminationsCfg_A1(TerminationsCfg): - base_too_low = DoneTerm(func=mdp.root_height_below_minimum, params={"minimum_height": 0.2}) - - -@configclass -class UnitreeA1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - terminations: TerminationsCfg_A1 = TerminationsCfg_A1() - - def __post_init__(self): - # post init of parent - super().__post_init__() - - self.scene.robot = UNITREE_A1_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01 - - # reduce action scale - self.actions.joint_pos.scale = 0.25 - - # event - self.events.push_robot = None - # TODO: TEMPORARILY DISABLED - adding this causes NaNs in the simulation - # self.events.add_base_mass.params["mass_distribution_params"] = (-1.0, 3.0) - # self.events.add_base_mass.params["asset_cfg"].body_names = "trunk" - self.events.base_external_force_torque.params["asset_cfg"].body_names = "trunk" - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - - # rewards - self.rewards.feet_air_time.params["sensor_cfg"].body_names = ".*_foot" - self.rewards.feet_air_time.weight = 0.01 - self.rewards.undesired_contacts.params["sensor_cfg"].body_names = ".*thigh" - self.rewards.dof_torques_l2.weight = -0.0002 - self.rewards.track_lin_vel_xy_exp.weight = 1.5 - self.rewards.track_ang_vel_z_exp.weight = 0.75 - self.rewards.dof_acc_l2.weight = -2.5e-7 - self.terminations.base_contact.params["sensor_cfg"].body_names = "trunk" - - -@configclass -class UnitreeA1RoughEnvCfg_PLAY(UnitreeA1RoughEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # spawn the robot randomly in the grid (instead of their terrain levels) - self.scene.terrain.max_init_terrain_level = None - # reduce the number of terrains to save memory - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing event - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py deleted file mode 100644 index a5c825bf2e98..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py +++ /dev/null @@ -1,60 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import AnymalBRoughEnvCfg - - -@configclass -class AnymalBFlatEnvCfg(AnymalBRoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=75, - nconmax=15, - cone="elliptic", - impratio=100, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - # post init of parent - super().__post_init__() - - # override rewards - self.rewards.flat_orientation_l2.weight = -5.0 - self.rewards.dof_torques_l2.weight = -2.5e-5 - self.rewards.feet_air_time.weight = 0.5 - # change terrain to flat - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - # no height scan - self.scene.height_scanner = None - self.observations.policy.height_scan = None - # no terrain curriculum - self.curriculum.terrain_levels = None - - -class AnymalBFlatEnvCfg_PLAY(AnymalBFlatEnvCfg): - def __post_init__(self) -> None: - # post init of parent - super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing event - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py deleted file mode 100644 index 3a7d80af470f..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py +++ /dev/null @@ -1,34 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -from isaaclab_assets import ANYMAL_B_CFG # isort: skip - - -@configclass -class AnymalBRoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.robot = ANYMAL_B_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.manager_call_max_mode = {"Scene": 1} - - -@configclass -class AnymalBRoughEnvCfg_PLAY(AnymalBRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py deleted file mode 100644 index 3cc348a97675..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py +++ /dev/null @@ -1,52 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import AnymalCRoughEnvCfg - - -@configclass -class AnymalCFlatEnvCfg(AnymalCRoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=120, - nconmax=15, - cone="elliptic", - impratio=100, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - # post init of parent - super().__post_init__() - - # override rewards - self.rewards.flat_orientation_l2.weight = -5.0 - self.rewards.dof_torques_l2.weight = -2.5e-5 - self.rewards.feet_air_time.weight = 0.5 - # change terrain to flat - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - # no terrain curriculum - self.curriculum.terrain_levels = None - - -class AnymalCFlatEnvCfg_PLAY(AnymalCFlatEnvCfg): - def __post_init__(self) -> None: - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py deleted file mode 100644 index 1d2f2676c64a..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py +++ /dev/null @@ -1,40 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -## -# Pre-defined configs -## -from isaaclab_assets.robots.anymal import ANYMAL_C_CFG # isort: skip - - -@configclass -class AnymalCRoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - # switch robot to anymal-c - self.scene.robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.scene.robot.actuators["legs"].armature = 0.01 - self.manager_call_max_mode = {"Scene": 1} - - -@configclass -class AnymalCRoughEnvCfg_PLAY(AnymalCRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py deleted file mode 100644 index 702443e815e2..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py +++ /dev/null @@ -1,46 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import AnymalDRoughEnvCfg - - -@configclass -class AnymalDFlatEnvCfg(AnymalDRoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=60, - nconmax=25, - cone="elliptic", - impratio=100.0, - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - super().__post_init__() - self.rewards.flat_orientation_l2.weight = -5.0 - self.rewards.dof_torques_l2.weight = -2.5e-5 - self.rewards.feet_air_time.weight = 0.5 - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None - - -class AnymalDFlatEnvCfg_PLAY(AnymalDFlatEnvCfg): - def __post_init__(self) -> None: - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py deleted file mode 100644 index 71f49edd2e68..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py +++ /dev/null @@ -1,37 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -## -# Pre-defined configs -## -from isaaclab_assets.robots.anymal import ANYMAL_D_CFG # isort: skip - - -@configclass -class AnymalDRoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.robot = ANYMAL_D_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.manager_call_max_mode = {"Scene": 1} - - -@configclass -class AnymalDRoughEnvCfg_PLAY(AnymalDRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/flat_env_cfg.py deleted file mode 100644 index 818c7aa30a1e..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/flat_env_cfg.py +++ /dev/null @@ -1,45 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import CassieRoughEnvCfg - - -@configclass -class CassieFlatEnvCfg(CassieRoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=52, - nconmax=15, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - super().__post_init__() - self.rewards.flat_orientation_l2.weight = -2.5 - self.rewards.feet_air_time.weight = 5.0 - self.rewards.joint_deviation_hip.params["asset_cfg"].joint_names = ["hip_rotation_.*"] - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None - - -class CassieFlatEnvCfg_PLAY(CassieFlatEnvCfg): - def __post_init__(self) -> None: - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/rough_env_cfg.py deleted file mode 100644 index 2fc44ea36cdc..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/rough_env_cfg.py +++ /dev/null @@ -1,96 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg - -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import ( - LocomotionVelocityRoughEnvCfg, - RewardsCfg, -) - -from isaaclab_assets.robots.cassie import CASSIE_CFG # isort: skip - - -@configclass -class CassieRewardsCfg(RewardsCfg): - termination_penalty = RewTerm(func=mdp.is_terminated, weight=-200.0) - feet_air_time = RewTerm( - func=mdp.feet_air_time_positive_biped, - weight=2.5, - params={ - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*toe"), - "command_name": "base_velocity", - "threshold": 0.3, - }, - ) - joint_deviation_hip = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["hip_abduction_.*", "hip_rotation_.*"])}, - ) - joint_deviation_toes = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["toe_joint_.*"])}, - ) - dof_pos_limits = RewTerm( - func=mdp.joint_pos_limits, - weight=-1.0, - params={"asset_cfg": SceneEntityCfg("robot", joint_names="toe_joint_.*")}, - ) - - -@configclass -class CassieRoughEnvCfg(LocomotionVelocityRoughEnvCfg): - rewards: CassieRewardsCfg = CassieRewardsCfg() - - def __post_init__(self): - super().__post_init__() - self.scene.robot = CASSIE_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.actions.joint_pos.scale = 0.5 - self.events.push_robot = None - # TODO: TEMPORARILY DISABLED - adding this causes NaNs in the simulation - # self.events.add_base_mass = None - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.base_external_force_torque.params["asset_cfg"].body_names = [".*pelvis"] - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - self.terminations.base_contact.params["sensor_cfg"].body_names = [".*pelvis"] - self.rewards.undesired_contacts = None - self.rewards.dof_torques_l2.weight = -5.0e-6 - self.rewards.track_lin_vel_xy_exp.weight = 2.0 - self.rewards.track_ang_vel_z_exp.weight = 1.0 - self.rewards.action_rate_l2.weight *= 1.5 - self.rewards.dof_acc_l2.weight *= 1.5 - - -@configclass -class CassieRoughEnvCfg_PLAY(CassieRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py deleted file mode 100644 index e99dd95ff82f..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py +++ /dev/null @@ -1,58 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import G1RoughEnvCfg - - -@configclass -class G1FlatEnvCfg(G1RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=95, - nconmax=10, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - super().__post_init__() - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None - self.rewards.track_ang_vel_z_exp.weight = 1.0 - self.rewards.lin_vel_z_l2.weight = -0.2 - self.rewards.action_rate_l2.weight = -0.005 - self.rewards.dof_acc_l2.weight = -1.0e-7 - self.rewards.feet_air_time.weight = 0.75 - self.rewards.feet_air_time.params["threshold"] = 0.4 - self.rewards.dof_torques_l2.weight = -2.0e-6 - self.rewards.dof_torques_l2.params["asset_cfg"] = SceneEntityCfg( - "robot", joint_names=[".*_hip_.*", ".*_knee_joint"] - ) - self.commands.base_velocity.ranges.lin_vel_x = (0.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (-0.5, 0.5) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - - -class G1FlatEnvCfg_PLAY(G1FlatEnvCfg): - def __post_init__(self) -> None: - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py deleted file mode 100644 index 7f9d12d0f201..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py +++ /dev/null @@ -1,178 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg - -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import ( - LocomotionVelocityRoughEnvCfg, - RewardsCfg, -) - -## -# Pre-defined configs -## -from isaaclab_assets import G1_MINIMAL_CFG # isort: skip - - -@configclass -class G1Rewards(RewardsCfg): - """Reward terms for the MDP.""" - - termination_penalty = RewTerm(func=mdp.is_terminated, weight=-200.0) - track_lin_vel_xy_exp = RewTerm( - func=mdp.track_lin_vel_xy_yaw_frame_exp, - weight=1.0, - params={"command_name": "base_velocity", "std": 0.5}, - ) - track_ang_vel_z_exp = RewTerm( - func=mdp.track_ang_vel_z_world_exp, weight=2.0, params={"command_name": "base_velocity", "std": 0.5} - ) - feet_air_time = RewTerm( - func=mdp.feet_air_time_positive_biped, - weight=0.25, - params={ - "command_name": "base_velocity", - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*_ankle_roll_link"), - "threshold": 0.4, - }, - ) - feet_slide = RewTerm( - func=mdp.feet_slide, - weight=-0.1, - params={ - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*_ankle_roll_link"), - "asset_cfg": SceneEntityCfg("robot", body_names=".*_ankle_roll_link"), - }, - ) - dof_pos_limits = RewTerm( - func=mdp.joint_pos_limits, - weight=-1.0, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_ankle_pitch_joint", ".*_ankle_roll_joint"])}, - ) - joint_deviation_hip = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.1, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_hip_yaw_joint", ".*_hip_roll_joint"])}, - ) - joint_deviation_arms = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.1, - params={ - "asset_cfg": SceneEntityCfg( - "robot", - joint_names=[ - ".*_shoulder_pitch_joint", - ".*_shoulder_roll_joint", - ".*_shoulder_yaw_joint", - ".*_elbow_pitch_joint", - ".*_elbow_roll_joint", - ], - ) - }, - ) - joint_deviation_fingers = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.05, - params={ - "asset_cfg": SceneEntityCfg( - "robot", - joint_names=[ - ".*_five_joint", - ".*_three_joint", - ".*_six_joint", - ".*_four_joint", - ".*_zero_joint", - ".*_one_joint", - ".*_two_joint", - ], - ) - }, - ) - joint_deviation_torso = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.1, - params={"asset_cfg": SceneEntityCfg("robot", joint_names="torso_joint")}, - ) - - -@configclass -class G1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - rewards: G1Rewards = G1Rewards() - - def __post_init__(self): - super().__post_init__() - self.scene.robot = G1_MINIMAL_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.events.push_robot = None - # TODO: TEMPORARILY DISABLED - adding this causes NaNs in the simulation - # self.events.add_base_mass = None - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.base_external_force_torque.params["asset_cfg"].body_names = ["torso_link"] - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - - # Rewards - self.rewards.lin_vel_z_l2.weight = 0.0 - self.rewards.undesired_contacts = None - self.rewards.flat_orientation_l2.weight = -1.0 - self.rewards.action_rate_l2.weight = -0.005 - self.rewards.dof_acc_l2.weight = -1.25e-7 - self.rewards.dof_acc_l2.params["asset_cfg"] = SceneEntityCfg( - "robot", joint_names=[".*_hip_.*", ".*_knee_joint"] - ) - self.rewards.dof_torques_l2.weight = -1.5e-7 - self.rewards.dof_torques_l2.params["asset_cfg"] = SceneEntityCfg( - "robot", joint_names=[".*_hip_.*", ".*_knee_joint", ".*_ankle_.*"] - ) - - # Commands - self.commands.base_velocity.ranges.lin_vel_x = (0.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (-0.0, 0.0) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - - # terminations - self.terminations.base_contact.params["sensor_cfg"].body_names = "torso_link" - - -@configclass -class G1RoughEnvCfg_PLAY(G1RoughEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.episode_length_s = 40.0 - # spawn the robot randomly in the grid (instead of their terrain levels) - self.scene.terrain.max_init_terrain_level = None - # reduce the number of terrains to save memory - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - - self.commands.base_velocity.ranges.lin_vel_x = (1.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - self.commands.base_velocity.ranges.heading = (0.0, 0.0) - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py deleted file mode 100644 index dd4012882075..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py +++ /dev/null @@ -1,46 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import UnitreeGo1RoughEnvCfg - - -@configclass -class UnitreeGo1FlatEnvCfg(UnitreeGo1RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=60, - nconmax=25, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - super().__post_init__() - self.rewards.flat_orientation_l2.weight = -2.5 - self.rewards.feet_air_time.weight = 0.25 - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None - - -class UnitreeGo1FlatEnvCfg_PLAY(UnitreeGo1FlatEnvCfg): - def __post_init__(self) -> None: - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py deleted file mode 100644 index cb3a602394a6..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py +++ /dev/null @@ -1,61 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -from isaaclab_assets.robots.unitree import UNITREE_GO1_CFG # isort: skip - - -@configclass -class UnitreeGo1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.robot = UNITREE_GO1_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.manager_call_max_mode = {"Scene": 1} - self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01 - self.actions.joint_pos.scale = 0.25 - self.events.push_robot = None - self.events.base_external_force_torque.params["asset_cfg"].body_names = "trunk" - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - self.rewards.feet_air_time.params["sensor_cfg"].body_names = ".*_foot" - self.rewards.feet_air_time.weight = 0.01 - self.rewards.undesired_contacts = None - self.rewards.dof_torques_l2.weight = -0.0002 - self.rewards.track_lin_vel_xy_exp.weight = 1.5 - self.rewards.track_ang_vel_z_exp.weight = 0.75 - self.rewards.dof_acc_l2.weight = -2.5e-7 - self.terminations.base_contact.params["sensor_cfg"].body_names = "trunk" - - -@configclass -class UnitreeGo1RoughEnvCfg_PLAY(UnitreeGo1RoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py deleted file mode 100644 index 349cc450ca9f..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py +++ /dev/null @@ -1,46 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import UnitreeGo2RoughEnvCfg - - -@configclass -class UnitreeGo2FlatEnvCfg(UnitreeGo2RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=65, - nconmax=35, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - super().__post_init__() - self.rewards.flat_orientation_l2.weight = -2.5 - self.rewards.feet_air_time.weight = 0.25 - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None - - -class UnitreeGo2FlatEnvCfg_PLAY(UnitreeGo2FlatEnvCfg): - def __post_init__(self) -> None: - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/rough_env_cfg.py deleted file mode 100644 index a25748da9885..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/rough_env_cfg.py +++ /dev/null @@ -1,60 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -from isaaclab_assets.robots.unitree import UNITREE_GO2_CFG # isort: skip - - -@configclass -class UnitreeGo2RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.robot = UNITREE_GO2_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01 - self.actions.joint_pos.scale = 0.25 - self.events.push_robot = None - self.events.base_external_force_torque.params["asset_cfg"].body_names = "base" - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - self.rewards.feet_air_time.params["sensor_cfg"].body_names = ".*_foot" - self.rewards.feet_air_time.weight = 0.01 - self.rewards.undesired_contacts = None - self.rewards.dof_torques_l2.weight = -0.0002 - self.rewards.track_lin_vel_xy_exp.weight = 1.5 - self.rewards.track_ang_vel_z_exp.weight = 0.75 - self.rewards.dof_acc_l2.weight = -2.5e-7 - self.terminations.base_contact.params["sensor_cfg"].body_names = "base" - - -@configclass -class UnitreeGo2RoughEnvCfg_PLAY(UnitreeGo2RoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py deleted file mode 100644 index a7021343464a..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py +++ /dev/null @@ -1,46 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import H1RoughEnvCfg - - -@configclass -class H1FlatEnvCfg(H1RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=65, - nconmax=15, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - super().__post_init__() - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None - self.rewards.feet_air_time.weight = 1.0 - self.rewards.feet_air_time.params["threshold"] = 0.6 - - -class H1FlatEnvCfg_PLAY(H1FlatEnvCfg): - def __post_init__(self) -> None: - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py deleted file mode 100644 index 1b65e018c9e4..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py +++ /dev/null @@ -1,131 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg - -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import ( - LocomotionVelocityRoughEnvCfg, - RewardsCfg, -) - -## -# Pre-defined configs -## -from isaaclab_assets import H1_MINIMAL_CFG # isort: skip - - -@configclass -class H1Rewards(RewardsCfg): - """Reward terms for the MDP.""" - - termination_penalty = RewTerm(func=mdp.is_terminated, weight=-200.0) - lin_vel_z_l2 = None - track_lin_vel_xy_exp = RewTerm( - func=mdp.track_lin_vel_xy_yaw_frame_exp, - weight=1.0, - params={"command_name": "base_velocity", "std": 0.5}, - ) - track_ang_vel_z_exp = RewTerm( - func=mdp.track_ang_vel_z_world_exp, weight=1.0, params={"command_name": "base_velocity", "std": 0.5} - ) - feet_air_time = RewTerm( - func=mdp.feet_air_time_positive_biped, - weight=0.25, - params={ - "command_name": "base_velocity", - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*ankle_link"), - "threshold": 0.4, - }, - ) - feet_slide = RewTerm( - func=mdp.feet_slide, - weight=-0.25, - params={ - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*ankle_link"), - "asset_cfg": SceneEntityCfg("robot", body_names=".*ankle_link"), - }, - ) - dof_pos_limits = RewTerm( - func=mdp.joint_pos_limits, weight=-1.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*_ankle")} - ) - joint_deviation_hip = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_hip_yaw", ".*_hip_roll"])}, - ) - joint_deviation_arms = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_shoulder_.*", ".*_elbow"])}, - ) - joint_deviation_torso = RewTerm( - func=mdp.joint_deviation_l1, weight=-0.1, params={"asset_cfg": SceneEntityCfg("robot", joint_names="torso")} - ) - - -@configclass -class H1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - rewards: H1Rewards = H1Rewards() - - def __post_init__(self): - super().__post_init__() - self.scene.robot = H1_MINIMAL_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.events.push_robot = None - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.base_external_force_torque.params["asset_cfg"].body_names = [".*torso_link"] - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - self.rewards.undesired_contacts = None - self.rewards.flat_orientation_l2.weight = -1.0 - self.rewards.dof_torques_l2.weight = 0.0 - self.rewards.action_rate_l2.weight = -0.005 - self.rewards.dof_acc_l2.weight = -1.25e-7 - self.commands.base_velocity.ranges.lin_vel_x = (0.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - self.terminations.base_contact.params["sensor_cfg"].body_names = ".*torso_link" - - -@configclass -class H1RoughEnvCfg_PLAY(H1RoughEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.episode_length_s = 40.0 - # spawn the robot randomly in the grid (instead of their terrain levels) - self.scene.terrain.max_init_terrain_level = None - # reduce the number of terrains to save memory - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - - self.commands.base_velocity.ranges.lin_vel_x = (1.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - self.commands.base_velocity.ranges.heading = (0.0, 0.0) - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py deleted file mode 100644 index 52d53c3db268..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py +++ /dev/null @@ -1,296 +0,0 @@ -# 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 - -import math -from dataclasses import MISSING - -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg, AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import CurriculumTermCfg as CurrTerm -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sensors import ContactSensorCfg -from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR -from isaaclab.utils.configclass import configclass -from isaaclab.utils.noise import UniformNoiseCfg as Unoise - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp - -## -# Pre-defined configs -## -from isaaclab.terrains.config.rough import ROUGH_TERRAINS_CFG # isort: skip - - -## -# Scene definition -## - - -@configclass -class MySceneCfg(InteractiveSceneCfg): - """Configuration for the terrain scene with a legged robot.""" - - # ground terrain - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="generator", - terrain_generator=ROUGH_TERRAINS_CFG, - max_init_terrain_level=5, - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg( - friction_combine_mode="multiply", - restitution_combine_mode="multiply", - static_friction=1.0, - dynamic_friction=1.0, - ), - visual_material=sim_utils.MdlFileCfg( - mdl_path=f"{ISAACLAB_NUCLEUS_DIR}/Materials/TilesMarbleSpiderWhiteBrickBondHoned/TilesMarbleSpiderWhiteBrickBondHoned.mdl", - project_uvw=True, - texture_scale=(0.25, 0.25), - ), - debug_vis=False, - ) - # robots - robot: ArticulationCfg = MISSING - # sensors - contact_forces = ContactSensorCfg( - prim_path="{ENV_REGEX_NS}/Robot/.*", - filter_prim_paths_expr=[], - history_length=3, - track_air_time=True, - ) - # lights - sky_light = AssetBaseCfg( - prim_path="/World/skyLight", - spawn=sim_utils.DomeLightCfg( - intensity=750.0, - texture_file=f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr", - ), - ) - - -## -# MDP settings -## - - -@configclass -class CommandsCfg: - """Command specifications for the MDP.""" - - base_velocity = mdp.UniformVelocityCommandCfg( - asset_name="robot", - resampling_time_range=(10.0, 10.0), - rel_standing_envs=0.02, - rel_heading_envs=1.0, - heading_command=True, - heading_control_stiffness=0.5, - debug_vis=True, - ranges=mdp.UniformVelocityCommandCfg.Ranges( - lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.0, 1.0), heading=(-math.pi, math.pi) - ), - ) - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - joint_pos = mdp.JointPositionActionCfg(asset_name="robot", joint_names=[".*"], scale=0.5, use_default_offset=True) - - -@configclass -class PolicyCfg(ObsGroup): - """Observations for policy group.""" - - # observation terms (order preserved) - base_lin_vel = ObsTerm(func=mdp.base_lin_vel, noise=Unoise(n_min=-0.1, n_max=0.1)) - base_ang_vel = ObsTerm(func=mdp.base_ang_vel, noise=Unoise(n_min=-0.2, n_max=0.2)) - projected_gravity = ObsTerm( - func=mdp.projected_gravity, - noise=Unoise(n_min=-0.05, n_max=0.05), - ) - velocity_commands = ObsTerm(func=mdp.generated_commands, params={"command_name": "base_velocity"}) - joint_pos = ObsTerm(func=mdp.joint_pos_rel, noise=Unoise(n_min=-0.01, n_max=0.01)) - joint_vel = ObsTerm(func=mdp.joint_vel_rel, noise=Unoise(n_min=-1.5, n_max=1.5)) - actions = ObsTerm(func=mdp.last_action) - - def __post_init__(self): - self.enable_corruption = True - self.concatenate_terms = True - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - # FIXME(warp-migration): COM randomization in exp manager-based locomotion currently causes - # NaNs and is temporarily disabled. - # base_com = EventTerm( - # func=mdp.randomize_rigid_body_com, - # mode="startup", - # params={ - # "asset_cfg": SceneEntityCfg("robot", body_names="base"), - # "com_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.01, 0.01)}, - # }, - # ) - base_com = None - - # reset - base_external_force_torque = EventTerm( - func=mdp.apply_external_force_torque, - mode="reset", - params={ - "asset_cfg": SceneEntityCfg("robot", body_names="base"), - "force_range": (0.0, 0.0), - "torque_range": (-0.0, 0.0), - }, - ) - - reset_base = EventTerm( - func=mdp.reset_root_state_uniform, - mode="reset", - params={ - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (-0.5, 0.5), - "y": (-0.5, 0.5), - "z": (-0.5, 0.5), - "roll": (-0.5, 0.5), - "pitch": (-0.5, 0.5), - "yaw": (-0.5, 0.5), - }, - }, - ) - - reset_robot_joints = EventTerm( - func=mdp.reset_joints_by_scale, - mode="reset", - params={ - "position_range": (0.5, 1.5), - "velocity_range": (0.0, 0.0), - }, - ) - - # interval - push_robot = EventTerm( - func=mdp.push_by_setting_velocity, - mode="interval", - interval_range_s=(10.0, 15.0), - params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}}, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # -- task - track_lin_vel_xy_exp = RewTerm( - func=mdp.track_lin_vel_xy_exp, weight=1.0, params={"command_name": "base_velocity", "std": math.sqrt(0.25)} - ) - track_ang_vel_z_exp = RewTerm( - func=mdp.track_ang_vel_z_exp, weight=0.5, params={"command_name": "base_velocity", "std": math.sqrt(0.25)} - ) - # -- penalties - lin_vel_z_l2 = RewTerm(func=mdp.lin_vel_z_l2, weight=-2.0) - ang_vel_xy_l2 = RewTerm(func=mdp.ang_vel_xy_l2, weight=-0.05) - dof_torques_l2 = RewTerm(func=mdp.joint_torques_l2, weight=-1.0e-5) - dof_acc_l2 = RewTerm(func=mdp.joint_acc_l2, weight=-2.5e-7) - action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=-0.01) - feet_air_time = RewTerm( - func=mdp.feet_air_time, - weight=0.125, - params={ - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*FOOT"), - "command_name": "base_velocity", - "threshold": 0.5, - }, - ) - undesired_contacts = RewTerm( - func=mdp.undesired_contacts, - weight=-1.0, - params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*THIGH"), "threshold": 1.0}, - ) - # -- optional penalties - flat_orientation_l2 = RewTerm(func=mdp.flat_orientation_l2, weight=0.0) - dof_pos_limits = RewTerm(func=mdp.joint_pos_limits, weight=0.0) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - time_out = DoneTerm(func=mdp.time_out, time_out=True) - base_contact = DoneTerm( - func=mdp.illegal_contact, - params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names="base"), "threshold": 1.0}, - ) - - -@configclass -class CurriculumCfg: - """Curriculum terms for the MDP.""" - - terrain_levels = CurrTerm(func=mdp.terrain_levels_vel) - - -## -# Environment configuration -## - - -@configclass -class LocomotionVelocityRoughEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the locomotion velocity-tracking environment.""" - - # Scene settings - scene: MySceneCfg = MySceneCfg(num_envs=4096, env_spacing=2.5, replicate_physics=True) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - commands: CommandsCfg = CommandsCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() - curriculum: CurriculumCfg = CurriculumCfg() - - def __post_init__(self): - """Post initialization.""" - # general settings - self.decimation = 4 - self.episode_length_s = 20.0 - # simulation settings - self.sim.dt = 1.0 / 200.0 - self.sim.render_interval = self.decimation - self.sim.physics_material = self.scene.terrain.physics_material - # update sensor update periods - if self.scene.contact_forces is not None: - self.scene.contact_forces.update_period = self.sim.dt - # check if terrain levels curriculum is enabled - if getattr(self.curriculum, "terrain_levels", None) is not None: - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.curriculum = True - else: - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.curriculum = False diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/__init__.py deleted file mode 100644 index 6cd56351b6e5..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# 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 - -from .reach import * diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/joint_pos_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/joint_pos_env_cfg.py deleted file mode 100644 index d6d36afd15f5..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/joint_pos_env_cfg.py +++ /dev/null @@ -1,74 +0,0 @@ -# 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 - -import math - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.manipulation.reach.mdp as mdp -from isaaclab_tasks_experimental.manager_based.manipulation.reach.reach_env_cfg import ReachEnvCfg - -## -# Pre-defined configs -## -from isaaclab_assets import FRANKA_PANDA_CFG # isort: skip - - -## -# Environment configuration -## - - -@configclass -class FrankaReachEnvCfg(ReachEnvCfg): - sim: SimulationCfg = SimulationCfg( - dt=1 / 60, - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=50, - nconmax=20, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ), - ) - - def __post_init__(self): - # post init of parent - super().__post_init__() - - # switch robot to franka - self.scene.robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - # override rewards - self.rewards.end_effector_position_tracking.params["asset_cfg"].body_names = ["panda_hand"] - self.rewards.end_effector_position_tracking_fine_grained.params["asset_cfg"].body_names = ["panda_hand"] - self.rewards.end_effector_orientation_tracking.params["asset_cfg"].body_names = ["panda_hand"] - - # override actions - self.actions.arm_action = mdp.JointPositionActionCfg( - asset_name="robot", joint_names=["panda_joint.*"], scale=0.5, use_default_offset=True - ) - # override command generator body - # end-effector is along z-direction - self.commands.ee_pose.body_name = "panda_hand" - self.commands.ee_pose.ranges.pitch = (math.pi, math.pi) - - -@configclass -class FrankaReachEnvCfg_PLAY(FrankaReachEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # disable randomization for play - self.observations.policy.enable_corruption = False diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/joint_pos_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/joint_pos_env_cfg.py deleted file mode 100644 index 42104e23b1b2..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/joint_pos_env_cfg.py +++ /dev/null @@ -1,74 +0,0 @@ -# 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 - -import math - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.manipulation.reach.mdp as mdp -from isaaclab_tasks_experimental.manager_based.manipulation.reach.reach_env_cfg import ReachEnvCfg - -## -# Pre-defined configs -## -from isaaclab_assets import UR10_CFG # isort: skip - - -## -# Environment configuration -## - - -@configclass -class UR10ReachEnvCfg(ReachEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=50, - nconmax=20, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - # post init of parent - super().__post_init__() - - # switch robot to ur10 - self.scene.robot = UR10_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - # override events - self.events.reset_robot_joints.params["position_range"] = (0.75, 1.25) - # override rewards - self.rewards.end_effector_position_tracking.params["asset_cfg"].body_names = ["ee_link"] - self.rewards.end_effector_position_tracking_fine_grained.params["asset_cfg"].body_names = ["ee_link"] - self.rewards.end_effector_orientation_tracking.params["asset_cfg"].body_names = ["ee_link"] - # override actions - self.actions.arm_action = mdp.JointPositionActionCfg( - asset_name="robot", joint_names=[".*"], scale=0.5, use_default_offset=True - ) - # override command generator body - # end-effector is along x-direction - self.commands.ee_pose.body_name = "ee_link" - self.commands.ee_pose.ranges.pitch = (math.pi / 2, math.pi / 2) - - -@configclass -class UR10ReachEnvCfg_PLAY(UR10ReachEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # disable randomization for play - self.observations.policy.enable_corruption = False diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py deleted file mode 100644 index e83a27e44db5..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py +++ /dev/null @@ -1,206 +0,0 @@ -# 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 - -from dataclasses import MISSING - -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg, AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import ActionTermCfg as ActionTerm -from isaaclab.managers import CurriculumTermCfg as CurrTerm -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.utils.configclass import configclass -from isaaclab.utils.noise import UniformNoiseCfg as Unoise - -import isaaclab_tasks_experimental.manager_based.manipulation.reach.mdp as mdp - -## -# Scene definition -## - - -@configclass -class ReachSceneCfg(InteractiveSceneCfg): - """Configuration for the scene with a robotic arm.""" - - # world - ground = AssetBaseCfg( - prim_path="/World/ground", - spawn=sim_utils.GroundPlaneCfg(), - init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, -1.05)), - ) - - # robots - robot: ArticulationCfg = MISSING - - # lights - light = AssetBaseCfg( - prim_path="/World/light", - spawn=sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=2500.0), - ) - - -## -# MDP settings -## - - -@configclass -class CommandsCfg: - """Command terms for the MDP.""" - - ee_pose = mdp.UniformPoseCommandCfg( - asset_name="robot", - body_name=MISSING, - resampling_time_range=(4.0, 4.0), - debug_vis=True, - ranges=mdp.UniformPoseCommandCfg.Ranges( - pos_x=(0.35, 0.65), - pos_y=(-0.2, 0.2), - pos_z=(0.15, 0.5), - roll=(0.0, 0.0), - pitch=MISSING, # depends on end-effector axis - yaw=(-3.14, 3.14), - ), - ) - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - arm_action: ActionTerm = MISSING - gripper_action: ActionTerm | None = None - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - @configclass - class PolicyCfg(ObsGroup): - """Observations for policy group.""" - - # observation terms (order preserved) - joint_pos = ObsTerm(func=mdp.joint_pos_rel, noise=Unoise(n_min=-0.01, n_max=0.01)) - joint_vel = ObsTerm(func=mdp.joint_vel_rel, noise=Unoise(n_min=-0.01, n_max=0.01)) - pose_command = ObsTerm(func=mdp.generated_commands, params={"command_name": "ee_pose"}) - actions = ObsTerm(func=mdp.last_action) - - def __post_init__(self): - self.enable_corruption = True - self.concatenate_terms = True - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - reset_base = EventTerm( - func=mdp.reset_root_state_uniform, - mode="reset", - params={"pose_range": {}, "velocity_range": {}}, - ) - - reset_robot_joints = EventTerm( - func=mdp.reset_joints_by_scale, - mode="reset", - params={ - "position_range": (0.5, 1.5), - "velocity_range": (0.0, 0.0), - }, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # task terms - end_effector_position_tracking = RewTerm( - func=mdp.position_command_error, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", body_names=MISSING), "command_name": "ee_pose"}, - ) - end_effector_position_tracking_fine_grained = RewTerm( - func=mdp.position_command_error_tanh, - weight=0.1, - params={"asset_cfg": SceneEntityCfg("robot", body_names=MISSING), "std": 0.1, "command_name": "ee_pose"}, - ) - end_effector_orientation_tracking = RewTerm( - func=mdp.orientation_command_error, - weight=-0.1, - params={"asset_cfg": SceneEntityCfg("robot", body_names=MISSING), "command_name": "ee_pose"}, - ) - - # action penalty - action_rate = RewTerm(func=mdp.action_rate_l2, weight=-0.0001) - joint_vel = RewTerm( - func=mdp.joint_vel_l2, - weight=-0.0001, - params={"asset_cfg": SceneEntityCfg("robot")}, - ) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - time_out = DoneTerm(func=mdp.time_out, time_out=True) - - -@configclass -class CurriculumCfg: - """Curriculum terms for the MDP.""" - - action_rate = CurrTerm( - func=mdp.modify_reward_weight, params={"term_name": "action_rate", "weight": -0.005, "num_steps": 4500} - ) - - joint_vel = CurrTerm( - func=mdp.modify_reward_weight, params={"term_name": "joint_vel", "weight": -0.001, "num_steps": 4500} - ) - - -## -# Environment configuration -## - - -@configclass -class ReachEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the reach end-effector pose tracking environment.""" - - # Scene settings - scene: ReachSceneCfg = ReachSceneCfg(num_envs=4096, env_spacing=2.5) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - commands: CommandsCfg = CommandsCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() - curriculum: CurriculumCfg = CurriculumCfg() - - def __post_init__(self): - """Post initialization.""" - # general settings - self.decimation = 2 - self.sim.render_interval = self.decimation - self.episode_length_s = 12.0 - self.viewer.eye = (3.5, 3.5, 3.5) - # simulation settings - self.sim.dt = 1.0 / 60.0 From 817758733266126023d72147cbfbbdd1853309e3 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Fri, 5 Jun 2026 00:11:10 +0000 Subject: [PATCH 17/58] Limit warp cfg adaptation to warp-managed term groups + coverage test Only observation/reward/termination/action terms are swapped to warp twins and have their SceneEntityCfg promoted; event, curriculum, recorder and command terms run on the stable (torch) managers, so their stable funcs are left in place. This lets the warp env reuse the full stable cfg (domain randomization events, reward curricula, etc.) instead of a trimmed copy. Add a parametrized coverage test that, for every manager-based *-Warp-v0 task, loads its stable cfg and asserts adapt_cfg_for_warp succeeds (all warp-managed twins resolve). Three tasks are xfail(strict) pending warp twins: cartpole survival_success_rate, ant/humanoid body_incoming_wrench. Also add the missing core/__init__.py so the registration walk descends into the migrated package. --- .../isaaclab_experimental/envs/frontend.py | 12 +++ .../test/envs/test_frontend.py | 6 +- .../test/envs/test_frontend_cfg_conversion.py | 92 +++++++++++++++++++ .../core/reach/config/ur_10/__init__.py | 2 +- 4 files changed, 110 insertions(+), 2 deletions(-) create mode 100644 source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index 2b85e474da8f..7a3eebd95f28 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -76,6 +76,14 @@ class FrontendIncompatibleError(RuntimeError): # the warp packages. Used by the direct-workflow guard. _WARP_ROOT_PREFIXES: tuple[str, ...] = ("isaaclab_experimental", "isaaclab_tasks_experimental") +# Top-level cfg groups whose managers run warp-first. Only terms under these are +# adapted (SceneEntityCfg promotion + MDP twin swap); the event, curriculum, +# recorder and command managers run on the stable (torch) implementation, so +# their stable funcs/SceneEntityCfgs are left untouched. A stable term left on a +# warp manager would break, so a missing twin in these groups is a hard error; +# a stable term on a stable manager is correct, so those groups are skipped. +_WARP_MANAGED_GROUPS: frozenset[str] = frozenset({"observations", "rewards", "terminations", "actions"}) + def build( frontend: Frontend | str, @@ -193,6 +201,8 @@ def _promote_scene_entity_cfgs(cfg: Any) -> None: promoted: list[str] = [] for path, term in _walk_terms(cfg): + if not path or path[0] not in _WARP_MANAGED_GROUPS: + continue params = getattr(term, "params", None) if not isinstance(params, dict): continue @@ -237,6 +247,8 @@ def _swap_mdp(cfg: Any, label: str) -> None: swapped = 0 missing: list[tuple[str, str, str]] = [] # (location, attr, symbol) for path, term in _walk_terms(cfg): + if not path or path[0] not in _WARP_MANAGED_GROUPS: + continue location = ".".join(path) for attr in ("func", "class_type"): stable = getattr(term, attr, None) diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py index 28666ed78e34..bd43f3c42611 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend.py +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -313,12 +313,16 @@ def test_walks_all_term_groups(self): ), ) _promote_scene_entity_cfgs(cfg) + # Warp-managed groups (rewards, observations incl. sub-groups) are promoted. self.assertIsInstance(cfg.rewards.r1.params["asset_cfg"], WarpSceneEntityCfg) - self.assertIsInstance(cfg.events.e1.params["asset_cfg"], WarpSceneEntityCfg) self.assertIsInstance(cfg.observations.policy.o1.params["asset_cfg"], WarpSceneEntityCfg) # The perception sub-group is reached even though its attribute name is # not hardcoded in the framework. self.assertIsInstance(cfg.observations.perception.o3.params["asset_cfg"], WarpSceneEntityCfg) + # Events run on the stable (torch) manager, so their SceneEntityCfg is + # left untouched — promoting it would hand a warp variant to a torch manager. + self.assertIsInstance(cfg.events.e1.params["asset_cfg"], StableSceneEntityCfg) + self.assertNotIsInstance(cfg.events.e1.params["asset_cfg"], WarpSceneEntityCfg) # ====================================================================== diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py new file mode 100644 index 000000000000..14af767d8d83 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -0,0 +1,92 @@ +# 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 + +"""Coverage test: every manager-based ``*-Warp-v0`` task adapts cleanly from its stable cfg. + +For each registered manager-based warp task we load the *stable* env cfg it points +at, force Newton physics, and run :func:`adapt_cfg_for_warp` — the same adaptation +the warp env runs in its ``__init__``. A pass proves the SceneEntityCfg promotion +and the MDP twin swap succeed for every warp-managed term (observations / rewards / +terminations / actions). This is the guard against drift: if a stable cfg later +gains a warp-managed term with no warp twin, the matching case fails here instead +of at training time. +""" + +from __future__ import annotations + +import importlib + +import gymnasium as gym +import isaaclab_tasks_experimental # noqa: F401 +import pytest +from isaaclab_experimental.envs.frontend import adapt_cfg_for_warp +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg + +# Registering the task packages is the whole point — import for side effects. +import isaaclab_tasks # noqa: F401 + +# Manager-based warp tasks are exactly those whose entry point is the shared warp +# env class; direct ``*-Direct-Warp-v0`` tasks register their own env class and +# are not cfg-adapted, so they are excluded here. +_MANAGER_WARP_ENTRY_POINT = "isaaclab_experimental.envs:ManagerBasedRLEnvWarp" + + +def _manager_warp_tasks() -> list[tuple[str, str]]: + """Return ``(task_id, env_cfg_entry_point)`` for every manager-based warp task.""" + tasks: list[tuple[str, str]] = [] + for task_id, spec in gym.registry.items(): + if spec.entry_point != _MANAGER_WARP_ENTRY_POINT: + continue + cfg_entry = (spec.kwargs or {}).get("env_cfg_entry_point") + if isinstance(cfg_entry, str): + tasks.append((task_id, cfg_entry)) + return sorted(tasks) + + +def _instantiate_cfg(cfg_entry_point: str): + """Import ``module:Class`` and instantiate the stable env cfg.""" + module_path, class_name = cfg_entry_point.split(":") + return getattr(importlib.import_module(module_path), class_name)() + + +_MANAGER_WARP_TASKS = _manager_warp_tasks() + +# Tasks whose stable cfg still includes a warp-managed term with no warp twin. +# These are tracked for implementation; delete the entry once the twin lands +# (xfail is strict, so an unexpected pass fails the suite and flags the cleanup). +_PENDING_WARP_TWINS = { + "Isaac-Cartpole-Warp-v0": "reward 'survival_success_rate'", + "Isaac-Ant-Warp-v0": "observation 'body_incoming_wrench'", + "Isaac-Humanoid-Warp-v0": "observation 'body_incoming_wrench'", +} + + +def _params(): + cases = [] + for task_id, cfg_entry_point in _MANAGER_WARP_TASKS: + marks = () + if task_id in _PENDING_WARP_TWINS: + marks = pytest.mark.xfail( + strict=True, + reason=f"pending warp twin for {_PENDING_WARP_TWINS[task_id]}", + ) + cases.append(pytest.param(task_id, cfg_entry_point, marks=marks, id=task_id)) + return cases + + +def test_manager_warp_tasks_are_registered(): + """Sanity: the package actually registered manager-based warp tasks.""" + assert _MANAGER_WARP_TASKS, "no manager-based '*-Warp-v0' tasks registered" + + +@pytest.mark.parametrize("task_id, cfg_entry_point", _params()) +def test_stable_cfg_adapts_to_warp(task_id: str, cfg_entry_point: str): + """The stable cfg behind each warp task adapts without a missing twin.""" + cfg = _instantiate_cfg(cfg_entry_point) + # The warp env requires Newton physics (normally via ``presets=newton_mjwarp``); + # set it directly so the test does not depend on Hydra preset resolution. + cfg.sim.physics = NewtonCfg(solver_cfg=MJWarpSolverCfg(), num_substeps=1) + # Raises FrontendIncompatibleError if any warp-managed term lacks a warp twin. + adapt_cfg_for_warp(cfg) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/ur_10/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/ur_10/__init__.py index 03843620488d..afeaffba8433 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/ur_10/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/ur_10/__init__.py @@ -13,7 +13,7 @@ # Warp tasks reuse the stable env cfgs; only physics (presets=newton_mjwarp) # and MDP twins are swapped at construction (see adapt_cfg_for_warp). -_stable_pkg = agents.__name__.rsplit(".", 1)[0] +# _stable_pkg = agents.__name__.rsplit(".", 1)[0] # gym.register( # id="Isaac-Reach-UR10-Warp-v0", From 28033c89b992cbfb6654eef1545154ee6bd6cb56 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Fri, 5 Jun 2026 00:16:58 +0000 Subject: [PATCH 18/58] Match consolidated cartpole IDs: drop -v0 from warp cartpole tasks Stable cartpole was consolidated to unversioned ids (Isaac-Cartpole, Isaac-Cartpole-Direct); mirror that on the warp side so each warp task pairs 1:1 with its stable counterpart. Other tasks keep -Warp-v0 because their stable ids still carry -v0 on develop. --- .../isaaclab_experimental/envs/frontend.py | 2 +- .../test/envs/test_frontend_cfg_conversion.py | 2 +- .../isaaclab_tasks_experimental/core/cartpole/__init__.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index 7a3eebd95f28..dc79bca5a5d8 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -115,7 +115,7 @@ def build( workflow = _detect_workflow(env_cfg) if workflow is Workflow.DIRECT: # Direct workflows aren't adapted — they must already be registered - # under the warp packages (e.g. ``Isaac-Cartpole-Direct-Warp-v0``). + # under the warp packages (e.g. ``Isaac-Cartpole-Direct-Warp``). _assert_direct_warp_registration(task_id) return gym.make(task_id, cfg=env_cfg, **construct_kwargs) diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index 14af767d8d83..88fc1a424906 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -57,7 +57,7 @@ def _instantiate_cfg(cfg_entry_point: str): # These are tracked for implementation; delete the entry once the twin lands # (xfail is strict, so an unexpected pass fails the suite and flags the cleanup). _PENDING_WARP_TWINS = { - "Isaac-Cartpole-Warp-v0": "reward 'survival_success_rate'", + "Isaac-Cartpole-Warp": "reward 'survival_success_rate'", "Isaac-Ant-Warp-v0": "observation 'body_incoming_wrench'", "Isaac-Humanoid-Warp-v0": "observation 'body_incoming_wrench'", } diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py index d5d93b3dba48..fd5463ad023f 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py @@ -26,7 +26,7 @@ # Manager-based: stable cfg adapted to warp at construction. gym.register( - id="Isaac-Cartpole-Warp-v0", + id="Isaac-Cartpole-Warp", entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ @@ -40,7 +40,7 @@ # Direct: dedicated warp env class + local cfg (not adaptable from a cfg). gym.register( - id="Isaac-Cartpole-Direct-Warp-v0", + id="Isaac-Cartpole-Direct-Warp", entry_point=f"{__name__}.cartpole_warp_env:CartpoleWarpEnv", disable_env_checker=True, kwargs={ From 87a03b7d01f7c4428d80cfcb56fed73b88151540 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Fri, 5 Jun 2026 00:38:59 +0000 Subject: [PATCH 19/58] Add warp twins: cartpole survival_success_rate, body_incoming_wrench MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements the two missing warp MDP twins so the full stable cfg adapts for every manager-based *-Warp task (clears the xfails in the conversion test): - cartpole `survival_success_rate`: ManagerTermBase reward mirroring stable — zero reward contribution, logs Metrics/success_rate (time-out rate) on reset. - generic `body_incoming_wrench`: observation reading the Newton joint-wrench sensor and gathering per-body force+torque, following the existing sensor-reading twins (undesired_contacts/illegal_contact). Lives in the shared fallback mdp so it is task-agnostic. Extends the observation manager's `body:N` out_dim resolution to also read `sensor_cfg` (not just `asset_cfg`), since wrench obs select bodies via a sensor entity. --- .../envs/mdp/observations.py | 46 +++++++++++++++++++ .../managers/observation_manager.py | 9 ++-- .../test/envs/test_frontend_cfg_conversion.py | 6 +-- .../core/cartpole/mdp/__init__.pyi | 3 +- .../core/cartpole/mdp/rewards.py | 30 ++++++++++++ 5 files changed, 84 insertions(+), 10 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py index 808231c8faf6..e36ddd189f27 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py @@ -367,3 +367,49 @@ def generated_commands(env: ManagerBasedEnv, out, command_name: str) -> None: fn._cmd_name = command_name fn._is_warmed_up = True wp.copy(out, fn._cmd_wp) + + +""" +Sensors. +""" + + +@wp.kernel +def _body_incoming_wrench_kernel( + force: wp.array(dtype=wp.vec3f, ndim=2), + torque: wp.array(dtype=wp.vec3f, ndim=2), + body_ids: wp.array(dtype=wp.int32), + out: wp.array(dtype=wp.float32, ndim=2), +): + env_id = wp.tid() + for k in range(body_ids.shape[0]): + b = body_ids[k] + f = force[env_id, b] + t = torque[env_id, b] + out[env_id, k * 6 + 0] = f[0] + out[env_id, k * 6 + 1] = f[1] + out[env_id, k * 6 + 2] = f[2] + out[env_id, k * 6 + 3] = t[0] + out[env_id, k * 6 + 4] = t[1] + out[env_id, k * 6 + 5] = t[2] + + +@generic_io_descriptor_warp(out_dim="body:6", observation_type="BodyState", on_inspect=[record_shape, record_dtype]) +def body_incoming_wrench(env: ManagerBasedEnv, out, sensor_cfg: SceneEntityCfg) -> None: + """Incoming spatial wrench [N, N·m] on selected bodies, laid out per body as ``[fx, fy, fz, tx, ty, tz]``. + + Warp-first override of :func:`isaaclab.envs.mdp.observations.body_incoming_wrench`. Reads the + Newton joint-wrench sensor and gathers the force/torque of the bodies selected by ``sensor_cfg`` + into ``out`` (shape ``(num_envs, 6 * num_selected_bodies)``). + """ + sensor = env.scene.sensors[sensor_cfg.name] + force = sensor.data.force + torque = sensor.data.torque + if force is None or torque is None: + raise RuntimeError("Joint wrench sensor data is not initialized. Call sim.reset() before reading observations.") + wp.launch( + kernel=_body_incoming_wrench_kernel, + dim=env.num_envs, + inputs=[force.warp, torque.warp, sensor_cfg.body_ids_wp, out], + device=env.device, + ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index 059ae2afdce1..46c93f891fa5 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -922,11 +922,12 @@ def _resolve_out_dim(self, out_dim: int | str, term_cfg: ObservationTermCfg) -> if isinstance(out_dim, str) and out_dim.startswith("body:"): per_body = int(out_dim.split(":")[1]) - asset_cfg = term_cfg.params.get("asset_cfg") - body_ids = getattr(asset_cfg, "body_ids", None) + # Body selection may live on an asset (articulation) or a sensor entity. + entity_cfg = term_cfg.params.get("asset_cfg") or term_cfg.params.get("sensor_cfg") + body_ids = getattr(entity_cfg, "body_ids", None) if body_ids is None or body_ids == slice(None): - asset = self._env.scene[asset_cfg.name] - return per_body * len(asset.body_names) + entity = self._env.scene[entity_cfg.name] + return per_body * len(entity.body_names) return per_body * len(body_ids) if out_dim == "command": diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index 88fc1a424906..ea3afda6d300 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -56,11 +56,7 @@ def _instantiate_cfg(cfg_entry_point: str): # Tasks whose stable cfg still includes a warp-managed term with no warp twin. # These are tracked for implementation; delete the entry once the twin lands # (xfail is strict, so an unexpected pass fails the suite and flags the cleanup). -_PENDING_WARP_TWINS = { - "Isaac-Cartpole-Warp": "reward 'survival_success_rate'", - "Isaac-Ant-Warp-v0": "observation 'body_incoming_wrench'", - "Isaac-Humanoid-Warp-v0": "observation 'body_incoming_wrench'", -} +_PENDING_WARP_TWINS: dict[str, str] = {} def _params(): diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.pyi index 377068405c79..657186c2539b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.pyi +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.pyi @@ -5,10 +5,11 @@ __all__ = [ "joint_pos_target_l2", + "survival_success_rate", ] # Forward stable MDP terms and experimental Warp-first overrides lazily, then # override with cartpole-specific terms below. from isaaclab_experimental.envs.mdp import * # noqa: F401, F403 -from .rewards import joint_pos_target_l2 +from .rewards import joint_pos_target_l2, survival_success_rate diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py index 2678ef03de65..ced0dadaad91 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py @@ -9,9 +9,12 @@ import warp as wp from isaaclab_experimental.managers import SceneEntityCfg +from isaaclab_experimental.managers.manager_base import ManagerTermBase from isaaclab_experimental.utils.warp.utils import wrap_to_pi if TYPE_CHECKING: + from isaaclab_experimental.managers.manager_term_cfg import RewardTermCfg + from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedRLEnv @@ -43,3 +46,30 @@ def joint_pos_target_l2(env: ManagerBasedRLEnv, out, target: float, asset_cfg: S inputs=[asset.data.joint_pos.warp, asset_cfg.joint_mask, out, target], device=env.device, ) + + +class survival_success_rate(ManagerTermBase): + """Logs the mean time-out (survival) rate of the resetting envs; contributes zero reward. + + Mirrors the stable term: the reward value is always zero (so it is registered with + ``weight=0.0``) and the only effect is logging ``Metrics/success_rate`` on reset, where + success is defined as the episode ending by time-out rather than an early termination. + """ + + def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + + def reset(self, env_mask: wp.array | None = None) -> None: + # ``time_outs`` is exposed as a torch tensor by the warp termination manager. + time_outs = self._env.termination_manager.time_outs + if env_mask is None: + survived = time_outs.float().mean() + else: + selected = time_outs[wp.to_torch(env_mask).bool()] + survived = selected.float().mean() if selected.numel() > 0 else time_outs.new_zeros(()) + self._env.extras.setdefault("log", {})["Metrics/success_rate"] = float(survived.item()) + + def __call__(self, env: ManagerBasedRLEnv, out) -> None: + # Pure logging term: the reward contribution is zero. The reward manager pre-zeroes + # ``out`` each step, but zero it explicitly so the term is correct independent of that. + out.zero_() From 3c8a8ae63a980cdec0fc2834c5633735f97ceca0 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Fri, 5 Jun 2026 00:48:06 +0000 Subject: [PATCH 20/58] Swap action-term class_type in warp adaptation; fix cartpole agents The cfg walker only matched ManagerTermBaseCfg, but ActionTermCfg is a separate base (not a subclass) carrying a swappable class_type. As a result action terms kept their stable class and the warp ActionManager rejected them ("not of type ActionType") when a *-Warp task reused a stable cfg. Match ActionTermCfg in _walk_terms so action class_type is swapped too, and guard it in the conversion test. Also fix the cartpole warp registration agent entry points to the current stable modules (rsl_rl_ppo_cfg:Cartpole[Direct]PPORunnerCfg, sb3_ppo_cfg) after the cartpole consolidation renamed them. --- .../isaaclab_experimental/envs/frontend.py | 23 +++++++++++-------- .../test/envs/test_frontend_cfg_conversion.py | 16 +++++++++++++ .../core/cartpole/__init__.py | 8 +++---- 3 files changed, 34 insertions(+), 13 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index dc79bca5a5d8..f2d379ece46a 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -384,26 +384,31 @@ def _is_swap_candidate(value: Any) -> bool: def _walk_terms(node: Any, path: tuple[str, ...] = ()) -> Iterator[tuple[tuple[str, ...], Any]]: - """Yield ``(path, term)`` for every :class:`ManagerTermBaseCfg` in the cfg tree. + """Yield ``(path, term)`` for every MDP term cfg in the cfg tree. + + A "term" is a :class:`ManagerTermBaseCfg` (observation/reward/termination/ + event/curriculum) *or* an :class:`ActionTermCfg` — the latter is a separate + base that is **not** a ``ManagerTermBaseCfg`` subclass, yet carries a + swappable ``class_type``, so it must be matched explicitly. Behavior at each node: - * Match (:class:`ManagerTermBaseCfg` instance): yield ``(path, node)`` and stop — - do not descend into ``term.params`` / ``term.func``. + * Match (a term cfg instance): yield ``(path, node)`` and stop — do not + descend into ``term.params`` / ``term.func`` / ``term.class_type``. * Configclass: don't yield; recurse into every non-underscore attribute, - extending the path. ``observations``, ``rewards``, ``events``, sub-groups - like ``observations.policy`` / ``observations.perception``, and anything - nested deeper are reached transparently. + extending the path. ``observations``, ``rewards``, ``events``, ``actions``, + sub-groups like ``observations.policy`` / ``observations.perception``, and + anything nested deeper are reached transparently. * Anything else (plain Python data, callables, non-configclass objects): stop. No yield, no recursion. Driven entirely by type — no attribute names are hardcoded — so future cfg layouts (extra observation groups, new nesting, etc.) are picked up - automatically as long as their terms subclass :class:`ManagerTermBaseCfg`. + automatically as long as their terms subclass one of the term base cfgs. """ - from isaaclab.managers.manager_term_cfg import ManagerTermBaseCfg + from isaaclab.managers.manager_term_cfg import ActionTermCfg, ManagerTermBaseCfg - if isinstance(node, ManagerTermBaseCfg): + if isinstance(node, (ManagerTermBaseCfg, ActionTermCfg)): yield path, node return if not hasattr(node, "__dataclass_fields__"): diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index ea3afda6d300..934542d5d525 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -77,6 +77,9 @@ def test_manager_warp_tasks_are_registered(): assert _MANAGER_WARP_TASKS, "no manager-based '*-Warp-v0' tasks registered" +_WARP_ROOTS = ("isaaclab_experimental", "isaaclab_tasks_experimental") + + @pytest.mark.parametrize("task_id, cfg_entry_point", _params()) def test_stable_cfg_adapts_to_warp(task_id: str, cfg_entry_point: str): """The stable cfg behind each warp task adapts without a missing twin.""" @@ -86,3 +89,16 @@ def test_stable_cfg_adapts_to_warp(task_id: str, cfg_entry_point: str): cfg.sim.physics = NewtonCfg(solver_cfg=MJWarpSolverCfg(), num_substeps=1) # Raises FrontendIncompatibleError if any warp-managed term lacks a warp twin. adapt_cfg_for_warp(cfg) + + # Action terms carry a ``class_type`` (not a ``func``) and live on a base that + # is not a ManagerTermBaseCfg; guard that the adapter still swaps them to the + # warp ActionTerm, otherwise the warp ActionManager rejects them at runtime. + actions = getattr(cfg, "actions", None) + if actions is not None: + for name, term in vars(actions).items(): + class_type = getattr(term, "class_type", None) + if class_type is not None: + assert class_type.__module__.startswith(_WARP_ROOTS), ( + f"{task_id}: action term '{name}' class_type was not swapped to a warp twin" + f" (got {class_type.__module__}.{class_type.__name__})" + ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py index fd5463ad023f..1ebde1b42bbd 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py @@ -32,9 +32,9 @@ kwargs={ "env_cfg_entry_point": f"{_stable_pkg}.cartpole_manager_env_cfg:CartpoleEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_manager_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_manager_ppo_cfg:CartpolePPORunnerCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_manager_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{agents.__name__}:sb3_manager_ppo_cfg.yaml", + "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml", }, ) @@ -46,8 +46,8 @@ kwargs={ "env_cfg_entry_point": f"{__name__}.cartpole_warp_env_cfg:CartpoleWarpEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_direct_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_direct_ppo_cfg:CartpolePPORunnerCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpoleDirectPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_direct_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{agents.__name__}:sb3_direct_ppo_cfg.yaml", + "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml", }, ) From d587a49a1bbf00ef86e8334c925a08bd6025d0e1 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Fri, 5 Jun 2026 00:57:22 +0000 Subject: [PATCH 21/58] Treat events as warp-managed; make survival_success_rate graph-safe The warp EventManager invokes term funcs with a Warp env-mask, so a stable event func breaks at runtime (torch index by wp.array). Add events to the warp-managed groups so their funcs are swapped to warp twins; only curriculum, recorder and command managers fall back to stable. Cartpole (reset_joints_by_offset) and reach now adapt and cartpole trains end-to-end on warp. Rewrite cartpole survival_success_rate as a plain zero-reward term func: its stable form logs a metric on reset via a host readback, which is incompatible with the reward manager's CUDA-graph-captured reset; the warp twin keeps the zero reward and omits the diagnostic metric. Velocity tasks remain xfail in the conversion test pending warp twins for randomize_rigid_body_mass / randomize_rigid_body_material (domain randomization). --- .../isaaclab_experimental/envs/frontend.py | 14 ++++---- .../test/envs/test_frontend.py | 19 +++++++--- .../test/envs/test_frontend_cfg_conversion.py | 24 +++++++------ .../core/cartpole/mdp/rewards.py | 35 +++++-------------- 4 files changed, 46 insertions(+), 46 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index f2d379ece46a..a3882949d826 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -77,12 +77,14 @@ class FrontendIncompatibleError(RuntimeError): _WARP_ROOT_PREFIXES: tuple[str, ...] = ("isaaclab_experimental", "isaaclab_tasks_experimental") # Top-level cfg groups whose managers run warp-first. Only terms under these are -# adapted (SceneEntityCfg promotion + MDP twin swap); the event, curriculum, -# recorder and command managers run on the stable (torch) implementation, so -# their stable funcs/SceneEntityCfgs are left untouched. A stable term left on a -# warp manager would break, so a missing twin in these groups is a hard error; -# a stable term on a stable manager is correct, so those groups are skipped. -_WARP_MANAGED_GROUPS: frozenset[str] = frozenset({"observations", "rewards", "terminations", "actions"}) +# adapted (SceneEntityCfg promotion + MDP twin swap). The event manager is +# warp-first too — it invokes term funcs with a Warp env-mask, so a stable event +# func (which expects torch ``env_ids``) breaks at runtime; its funcs must be +# swapped to warp twins. The curriculum, recorder and command managers run on the +# stable (torch) implementation, so their terms are left untouched. A stable term +# left on a warp manager would break, so a missing twin in these groups is a hard +# error; a stable term on a stable manager is correct, so those groups are skipped. +_WARP_MANAGED_GROUPS: frozenset[str] = frozenset({"observations", "rewards", "terminations", "actions", "events"}) def build( diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py index bd43f3c42611..5ba58490c385 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend.py +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -206,11 +206,17 @@ class _EventsCfg: e1: EventTermCfg | None = None +@configclass +class _CurriculumCfg: + c1: EventTermCfg | None = None + + @configclass class _CfgFixture: observations: _ObservationsCfg | None = None rewards: _RewardsCfg | None = None events: _EventsCfg | None = None + curriculum: _CurriculumCfg | None = None def _term(func=None, params: dict | None = None) -> RewardTermCfg: @@ -311,18 +317,23 @@ def test_walks_all_term_groups(self): o3=ObservationTermCfg(func=_stable_func, params={"asset_cfg": StableSceneEntityCfg(name="o-per")}) ), ), + curriculum=_CurriculumCfg( + c1=EventTermCfg(func=_stable_func, mode="reset", params={"asset_cfg": StableSceneEntityCfg(name="c")}) + ), ) _promote_scene_entity_cfgs(cfg) - # Warp-managed groups (rewards, observations incl. sub-groups) are promoted. + # Warp-managed groups (rewards, observations incl. sub-groups, events) are promoted. self.assertIsInstance(cfg.rewards.r1.params["asset_cfg"], WarpSceneEntityCfg) self.assertIsInstance(cfg.observations.policy.o1.params["asset_cfg"], WarpSceneEntityCfg) # The perception sub-group is reached even though its attribute name is # not hardcoded in the framework. self.assertIsInstance(cfg.observations.perception.o3.params["asset_cfg"], WarpSceneEntityCfg) - # Events run on the stable (torch) manager, so their SceneEntityCfg is + # The event manager is warp-first, so its terms are promoted too. + self.assertIsInstance(cfg.events.e1.params["asset_cfg"], WarpSceneEntityCfg) + # Curriculum runs on the stable (torch) manager, so its SceneEntityCfg is # left untouched — promoting it would hand a warp variant to a torch manager. - self.assertIsInstance(cfg.events.e1.params["asset_cfg"], StableSceneEntityCfg) - self.assertNotIsInstance(cfg.events.e1.params["asset_cfg"], WarpSceneEntityCfg) + self.assertIsInstance(cfg.curriculum.c1.params["asset_cfg"], StableSceneEntityCfg) + self.assertNotIsInstance(cfg.curriculum.c1.params["asset_cfg"], WarpSceneEntityCfg) # ====================================================================== diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index 934542d5d525..6d0bdba4166e 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -53,21 +53,25 @@ def _instantiate_cfg(cfg_entry_point: str): _MANAGER_WARP_TASKS = _manager_warp_tasks() -# Tasks whose stable cfg still includes a warp-managed term with no warp twin. -# These are tracked for implementation; delete the entry once the twin lands -# (xfail is strict, so an unexpected pass fails the suite and flags the cleanup). -_PENDING_WARP_TWINS: dict[str, str] = {} + +def _pending_reason(task_id: str) -> str | None: + """Reason a task's cfg cannot yet fully adapt, or ``None`` if it should pass. + + Tracked for implementation; remove an entry once its warp twin lands (xfail is + strict, so an unexpected pass fails the suite and flags the cleanup). + """ + # Velocity locomotion tasks randomize body mass + material in their event terms, + # which have no warp twins yet (would need Newton articulation mass/material setters). + if task_id.startswith("Isaac-Velocity-"): + return "pending warp event twins: randomize_rigid_body_mass / randomize_rigid_body_material" + return None def _params(): cases = [] for task_id, cfg_entry_point in _MANAGER_WARP_TASKS: - marks = () - if task_id in _PENDING_WARP_TWINS: - marks = pytest.mark.xfail( - strict=True, - reason=f"pending warp twin for {_PENDING_WARP_TWINS[task_id]}", - ) + reason = _pending_reason(task_id) + marks = pytest.mark.xfail(strict=True, reason=reason) if reason else () cases.append(pytest.param(task_id, cfg_entry_point, marks=marks, id=task_id)) return cases diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py index ced0dadaad91..7830edabbae1 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py @@ -9,12 +9,9 @@ import warp as wp from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_experimental.managers.manager_base import ManagerTermBase from isaaclab_experimental.utils.warp.utils import wrap_to_pi if TYPE_CHECKING: - from isaaclab_experimental.managers.manager_term_cfg import RewardTermCfg - from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedRLEnv @@ -48,28 +45,14 @@ def joint_pos_target_l2(env: ManagerBasedRLEnv, out, target: float, asset_cfg: S ) -class survival_success_rate(ManagerTermBase): - """Logs the mean time-out (survival) rate of the resetting envs; contributes zero reward. +def survival_success_rate(env: ManagerBasedRLEnv, out) -> None: + """Zero-reward placeholder mirroring the stable ``survival_success_rate`` term. - Mirrors the stable term: the reward value is always zero (so it is registered with - ``weight=0.0``) and the only effect is logging ``Metrics/success_rate`` on reset, where - success is defined as the episode ending by time-out rather than an early termination. + The stable term contributes zero reward (registered with ``weight=0.0``) and its only + effect is logging ``Metrics/success_rate`` on reset. That logging needs a host readback + of the time-out mask, which is incompatible with the warp reward manager's + CUDA-graph-captured reset, so the warp twin keeps the zero reward and omits the diagnostic + metric. Implemented as a plain term func (not a :class:`ManagerTermBase`) since there is no + graph-safe per-reset state to maintain. """ - - def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): - super().__init__(cfg, env) - - def reset(self, env_mask: wp.array | None = None) -> None: - # ``time_outs`` is exposed as a torch tensor by the warp termination manager. - time_outs = self._env.termination_manager.time_outs - if env_mask is None: - survived = time_outs.float().mean() - else: - selected = time_outs[wp.to_torch(env_mask).bool()] - survived = selected.float().mean() if selected.numel() > 0 else time_outs.new_zeros(()) - self._env.extras.setdefault("log", {})["Metrics/success_rate"] = float(survived.item()) - - def __call__(self, env: ManagerBasedRLEnv, out) -> None: - # Pure logging term: the reward contribution is zero. The reward manager pre-zeroes - # ``out`` each step, but zero it explicitly so the term is correct independent of that. - out.zero_() + out.zero_() From 855588fda941b857759c445bb00428c906a7ea61 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 9 Jun 2026 06:17:11 +0000 Subject: [PATCH 22/58] Align warp env ids with develop torch names after cleanup renames Develop's task-cleanup PRs dropped the -v0 suffix from the core Ant, Humanoid, and Reach-Franka torch ids. The warp variants follow a fixed pattern (torch id with -Warp inserted before any -Play/-v0 suffix), so sync the six affected registrations and their doc references: Isaac-Ant-Warp-v0 -> Isaac-Ant-Warp Isaac-Ant-Direct-Warp-v0 -> Isaac-Ant-Direct-Warp Isaac-Humanoid-Warp-v0 -> Isaac-Humanoid-Warp Isaac-Humanoid-Direct-Warp-v0 -> Isaac-Humanoid-Direct-Warp Isaac-Reach-Franka-Warp-v0 -> Isaac-Reach-Franka-Warp Isaac-Reach-Franka-Warp-Play-v0 -> Isaac-Reach-Franka-Warp-Play Cartpole already conformed; contrib and core-velocity tasks kept -v0 upstream, so their warp ids are unchanged. --- .../physical-backends/newton/warp-environments.rst | 10 +++++----- docs/source/overview/environments.rst | 6 +++--- .../isaaclab_tasks_experimental/core/__init__.py | 12 ++++++++++++ .../core/direct_ant/__init__.py | 2 +- .../core/direct_humanoid/__init__.py | 2 +- .../core/manager_ant/__init__.py | 2 +- .../core/manager_humanoid/__init__.py | 2 +- .../core/reach/config/franka/__init__.py | 4 ++-- 8 files changed, 26 insertions(+), 14 deletions(-) create mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/__init__.py diff --git a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst index 94a2dad20c87..86068a232bcb 100644 --- a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst +++ b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst @@ -32,8 +32,8 @@ Direct Warp Environments ^^^^^^^^^^^^^^^^^^^^^^^^ - ``Isaac-Cartpole-Direct-Warp-v0`` — Cartpole balance -- ``Isaac-Ant-Direct-Warp-v0`` — Ant locomotion -- ``Isaac-Humanoid-Direct-Warp-v0`` — Humanoid locomotion +- ``Isaac-Ant-Direct-Warp`` — Ant locomotion +- ``Isaac-Humanoid-Direct-Warp`` — Humanoid locomotion - ``Isaac-Repose-Cube-Allegro-Direct-Warp-v0`` — Allegro hand cube repose @@ -43,8 +43,8 @@ Manager-Based Warp Environments **Classic** - ``Isaac-Cartpole-Warp-v0`` -- ``Isaac-Ant-Warp-v0`` -- ``Isaac-Humanoid-Warp-v0`` +- ``Isaac-Ant-Warp`` +- ``Isaac-Humanoid-Warp`` **Locomotion (Flat)** @@ -61,7 +61,7 @@ Manager-Based Warp Environments **Manipulation** -- ``Isaac-Reach-Franka-Warp-v0`` +- ``Isaac-Reach-Franka-Warp`` - ``Isaac-Reach-UR10-Warp-v0`` diff --git a/docs/source/overview/environments.rst b/docs/source/overview/environments.rst index 08ca2372f504..0b4eac1144b0 100644 --- a/docs/source/overview/environments.rst +++ b/docs/source/overview/environments.rst @@ -857,7 +857,7 @@ inferencing, including reading from an already trained checkpoint and disabling - **Workflow** - **RL Library** - **Presets** - * - Isaac-Ant-Direct-Warp-v0 + * - Isaac-Ant-Direct-Warp - - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) @@ -867,7 +867,7 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``ovphysx``, ``physx`` - * - Isaac-Ant-Warp-v0 + * - Isaac-Ant-Warp - - Manager Based - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) @@ -1055,7 +1055,7 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **skrl** (AMP) - - * - Isaac-Humanoid-Direct-Warp-v0 + * - Isaac-Humanoid-Direct-Warp - - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/__init__.py new file mode 100644 index 000000000000..5491e12e3855 --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/__init__.py @@ -0,0 +1,12 @@ +# 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 + +"""Warp entry points for the core task families. + +Mirrors the ``isaaclab_tasks.core`` layout. Manager-based tasks register +``*-Warp-v0`` ids that reuse the stable env cfgs (adapted to the warp runtime at +construction); direct tasks register their own warp env classes here. Task +definitions (cfgs) live in ``isaaclab_tasks`` to avoid duplication. +""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_ant/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_ant/__init__.py index 33d5ac926387..79106ae272bd 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_ant/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_ant/__init__.py @@ -16,7 +16,7 @@ stable_agents = "isaaclab_tasks.core.locomotion.ant.agents" gym.register( - id="Isaac-Ant-Direct-Warp-v0", + id="Isaac-Ant-Direct-Warp", entry_point=f"{__name__}.ant_env_warp:AntWarpEnv", disable_env_checker=True, kwargs={ diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_humanoid/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_humanoid/__init__.py index c38d7dd955e5..a9ab94d9f4ac 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_humanoid/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/direct_humanoid/__init__.py @@ -16,7 +16,7 @@ stable_agents = "isaaclab_tasks.core.locomotion.humanoid.agents" gym.register( - id="Isaac-Humanoid-Direct-Warp-v0", + id="Isaac-Humanoid-Direct-Warp", entry_point=f"{__name__}.humanoid_warp_env:HumanoidWarpEnv", disable_env_checker=True, kwargs={ diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_ant/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_ant/__init__.py index 76a22111d486..a55544cb793d 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_ant/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_ant/__init__.py @@ -21,7 +21,7 @@ ## gym.register( - id="Isaac-Ant-Warp-v0", + id="Isaac-Ant-Warp", entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/__init__.py index 537d2a9aceaf..ebf7f2909a9e 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/manager_humanoid/__init__.py @@ -21,7 +21,7 @@ ## gym.register( - id="Isaac-Humanoid-Warp-v0", + id="Isaac-Humanoid-Warp", entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/franka/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/franka/__init__.py index a82f977588fd..b94ba1315fdc 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/franka/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/config/franka/__init__.py @@ -21,7 +21,7 @@ ## gym.register( - id="Isaac-Reach-Franka-Warp-v0", + id="Isaac-Reach-Franka-Warp", entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ @@ -33,7 +33,7 @@ ) gym.register( - id="Isaac-Reach-Franka-Warp-Play-v0", + id="Isaac-Reach-Franka-Warp-Play", entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", disable_env_checker=True, kwargs={ From b6341271983ae633fbc663e79da3321d69d52eae Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 15 Jul 2026 17:28:20 -0700 Subject: [PATCH 23/58] Register --frontend in the shared RL training CLI Define the flag once in add_common_train_args so every RL library's train and benchmark entrypoint exposes it, instead of registering it per-script for RSL-RL only. create_isaaclab_env now reads the argument directly, so a caller that misses the shared registration fails loudly rather than silently defaulting to the torch runtime. Also drop the frontend dispatch from the deprecated rsl_rl/train.py wrapper; deprecated entrypoints should not grow new features. --- scripts/reinforcement_learning/common.py | 14 ++++++- .../reinforcement_learning/rsl_rl/train.py | 38 +------------------ .../rsl_rl/train_rsl_rl.py | 10 ----- .../test_reinforcement_learning_common.py | 13 ++++++- .../changelog.d/warp-manager-bridge.rst | 5 ++- 5 files changed, 30 insertions(+), 50 deletions(-) diff --git a/scripts/reinforcement_learning/common.py b/scripts/reinforcement_learning/common.py index 8edc124b469b..a36b5ac9266b 100644 --- a/scripts/reinforcement_learning/common.py +++ b/scripts/reinforcement_learning/common.py @@ -257,6 +257,18 @@ def add_common_train_args( ) 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( + "--frontend", + type=str, + choices=["torch", "warp"], + default="torch", + help=( + "Runtime that constructs the environment. 'torch' uses the registered stable environment via" + " gym.make. 'warp' (experimental) adapts a manager-based task config onto the Warp runtime, or" + " dispatches a direct task to its registered Warp environment; requires isaaclab_experimental" + " and `presets=newton_mjwarp`." + ), + ) if include_agent: parser.add_argument("--agent", type=str, default=agent_default, help=agent_help) parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment") @@ -415,7 +427,7 @@ def create_isaaclab_env( The created Gymnasium environment. """ render_mode = "rgb_array" if args_cli.video else None - frontend = getattr(args_cli, "frontend", "torch") + frontend = args_cli.frontend if frontend == "torch": env = gym.make(task, cfg=env_cfg, render_mode=render_mode) else: diff --git a/scripts/reinforcement_learning/rsl_rl/train.py b/scripts/reinforcement_learning/rsl_rl/train.py index a494e375f470..c9be857d41f5 100644 --- a/scripts/reinforcement_learning/rsl_rl/train.py +++ b/scripts/reinforcement_learning/rsl_rl/train.py @@ -79,19 +79,6 @@ "--ray-proc-id", "-rid", type=int, default=None, help="Automatically configured by Ray integration, otherwise None." ) parser.add_argument("--external_callback", default=None, help="Fully qualified path to an externally defined callback.") -parser.add_argument( - "--frontend", - type=str, - default="torch", - choices=["torch", "warp"], - help=( - "Runtime backend for the env. 'torch' uses isaaclab.envs.* (PhysX or Newton via" - " factory dispatch); 'warp' routes through the WarpFrontend, which adapts a" - " manager-based cfg onto ManagerBasedRLEnvWarp or dispatches a direct task to its" - " registered warp env class. See isaaclab_experimental.envs.frontend for the" - " pluggable rule pipeline." - ), -) cli_args.add_rsl_rl_args(parser) add_launcher_args(parser) args_cli, remaining_args = setup_preset_cli(parser) @@ -113,22 +100,6 @@ # intersection share the same token vocabulary (the callback reads the user's # original sys.argv), so preset tokens like ``physics=NAME`` compare correctly. remaining_args = list_intersection(remaining_args, remaining_args_env_registration) - -# Build path for ``--frontend=warp`` lives in ``isaaclab_experimental``, which is an -# optional package — torch (the default) must still work without it, so import lazily. -# Warp callers are required to pass ``presets=newton_mjwarp`` themselves; the frontend hard-checks -# it at build time rather than mutating Hydra args here. -_frontend_build = None -try: - from isaaclab_experimental.envs.frontend import build as _frontend_build # noqa: E402 -except ImportError as _frontend_import_err: - if args_cli.frontend != "torch": - raise SystemExit( - f"--frontend={args_cli.frontend} requires isaaclab_experimental to be installed," - f" but the import failed: {_frontend_import_err}" - ) from _frontend_import_err - logger.info("isaaclab_experimental not available; --frontend=torch will dispatch via plain gym.make.") - sys.argv = [sys.argv[0]] + remaining_args # -- check RSL-RL version ---------------------------------------------------- @@ -213,13 +184,8 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen # set the log directory for the environment (works for all environment types) env_cfg.log_dir = log_dir - # Build the env via the selected frontend. Warp adapts env_cfg in place; - # missing warp twins are a hard failure. - render_mode = "rgb_array" if args_cli.video else None - if _frontend_build is not None: - env = _frontend_build(args_cli.frontend, env_cfg, args_cli.task, render_mode=render_mode) - else: - env = gym.make(args_cli.task, cfg=env_cfg, render_mode=render_mode) + # create isaac environment + env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) # convert to single-agent instance if required by the RL algorithm if isinstance(env.unwrapped.cfg, DirectMARLEnvCfg): diff --git a/scripts/reinforcement_learning/rsl_rl/train_rsl_rl.py b/scripts/reinforcement_learning/rsl_rl/train_rsl_rl.py index 670da05b7f97..8f5c797cad40 100644 --- a/scripts/reinforcement_learning/rsl_rl/train_rsl_rl.py +++ b/scripts/reinforcement_learning/rsl_rl/train_rsl_rl.py @@ -82,16 +82,6 @@ def _parse_args(argv: list[str]) -> argparse.Namespace: default=None, help="Fully qualified path to an externally defined callback.", ) - parser.add_argument( - "--frontend", - choices=["torch", "warp"], - default="torch", - help=( - "Environment runtime. 'torch' uses the registered stable environment; " - "'warp' adapts a manager-based config to ManagerBasedRLEnvWarp or requires " - "a registered Warp environment for direct tasks." - ), - ) CLI_ARGS.add_rsl_rl_args(parser) add_isaaclab_launcher_args(parser) # setup_preset_cli registers preset-selection help text + runs parse_known_args diff --git a/source/isaaclab/test/test_reinforcement_learning_common.py b/source/isaaclab/test/test_reinforcement_learning_common.py index 9b3a5f3836ac..9de4173bd8f7 100644 --- a/source/isaaclab/test/test_reinforcement_learning_common.py +++ b/source/isaaclab/test/test_reinforcement_learning_common.py @@ -221,6 +221,17 @@ def test_enable_cameras_for_video_enables_cameras_for_sensor_capture() -> None: assert args_cli.enable_cameras +def test_common_train_args_register_frontend_with_torch_default() -> None: + """Every RL library CLI exposes ``--frontend`` and defaults to the torch runtime.""" + parser = argparse.ArgumentParser() + add_common_train_args(parser, agent_default=None, agent_help="", include_agent=False) + + assert parser.parse_args([]).frontend == "torch" + assert parser.parse_args(["--frontend", "warp"]).frontend == "warp" + with pytest.raises(SystemExit): + parser.parse_args(["--frontend", "tensorflow"]) + + def test_create_isaaclab_env_uses_registered_torch_env_by_default(monkeypatch: pytest.MonkeyPatch) -> None: """The shared factory preserves the existing Gym path when no frontend is selected.""" expected_env = object() @@ -232,7 +243,7 @@ def fake_make(task: str, **kwargs: Any) -> Any: return expected_env monkeypatch.setattr(_rl_common.gym, "make", fake_make) - args_cli = argparse.Namespace(video=False) + args_cli = argparse.Namespace(video=False, frontend="torch") env = create_isaaclab_env("Isaac-Test", env_cfg, args_cli, convert_marl_to_single_agent=False) diff --git a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst index a645acf116ca..4b3996529ba8 100644 --- a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst +++ b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst @@ -1,8 +1,9 @@ Added ^^^^^ -* Added ``--frontend {torch,warp}`` to the RSL-RL training entrypoint for - selecting the environment runtime; default ``torch`` is unchanged. +* Added ``--frontend {torch,warp}`` to the shared reinforcement-learning + training CLI (all supported RL libraries) for selecting the environment + runtime; default ``torch`` is unchanged. * Added :mod:`isaaclab_experimental.envs.frontend` runtime selector and :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable` used by ``--frontend=warp`` to adapt stable cfgs onto the warp runtime. From 826779f328a30c8932c0eacb83a0386a59633149 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 15 Jul 2026 17:34:39 -0700 Subject: [PATCH 24/58] Report success-rate metric from the Warp Cartpole reward twin The Warp survival_success_rate twin wrote zeros and dropped the Metrics/success_rate value that the stable class-based term flushes into extras on reset, so torch and warp runs of the same task were not comparable on the success metric. Rebuild the twin as a Warp ManagerTermBase class that accumulates the timed-out fraction of just-reset envs in device buffers and exposes it as a persistent tensor view through the reward manager's reset extras. Class reward terms may now return such views from reset(); the reward manager merges them, which keeps the whole reset stage CUDA-graph capturable (no host readback). --- .../managers/manager_base.py | 12 ++- .../managers/reward_manager.py | 9 ++- .../changelog.d/warp-manager-bridge.rst | 3 + .../classic/cartpole/mdp/rewards.py | 73 +++++++++++++++++-- 4 files changed, 85 insertions(+), 12 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py index 699e27170e02..7954800d87c9 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py @@ -37,6 +37,8 @@ if TYPE_CHECKING: + import torch + from isaaclab.envs import ManagerBasedEnv # import logger @@ -107,14 +109,20 @@ def __name__(self) -> str: Operations. """ - def reset(self, env_mask: wp.array | None = None) -> None: + def reset(self, env_mask: wp.array | None = None) -> dict[str, torch.Tensor] | None: """Resets the manager term (mask-based). Args: env_mask: Boolean mask of shape (num_envs,) indicating which envs to reset. If None, all envs are considered. + + Returns: + An optional dictionary of logging values to merge into the owning manager's + reset extras. To stay CUDA-graph capturable, values must be persistent + tensor views over device buffers written by kernels (no host readback); + the same objects must be returned on every call. """ - pass + return None def serialize(self) -> dict: """General serialization call. Includes the configuration dict.""" diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py index 99e3d9af979c..7dd6b85bd6ed 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py @@ -272,9 +272,14 @@ def reset( device=self.device, ) - # reset all the reward terms + # reset all the reward terms; class terms may expose logging values as + # persistent tensor views (see :meth:`ManagerTermBase.reset`), which are + # merged into the manager's reset extras. Merging the same views again on + # subsequent resets is a no-op, so this stays capture-safe. for term_cfg in self._class_term_cfgs: - term_cfg.func.reset(env_mask=env_mask) + term_extras = term_cfg.func.reset(env_mask=env_mask) + if term_extras: + self._reset_extras.update(term_extras) return self._reset_extras diff --git a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.rst index 86ae1af09452..4ac20a99187a 100644 --- a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.rst +++ b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.rst @@ -3,3 +3,6 @@ Fixed * Fixed the direct Warp Cartpole task to match the stable task's observations, reset ranges, termination condition, reward scaling, and scene configuration. +* Fixed the Warp Cartpole ``survival_success_rate`` twin to report the + ``Metrics/success_rate`` value on-device instead of silently dropping the + metric. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py index b54267ac5ea8..6049af38ad78 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py @@ -8,12 +8,15 @@ from typing import TYPE_CHECKING import warp as wp -from isaaclab_experimental.managers import SceneEntityCfg +from isaaclab_experimental.managers import ManagerTermBase, SceneEntityCfg from isaaclab_experimental.utils.warp.utils import wrap_to_pi if TYPE_CHECKING: + import torch + from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.managers import RewardTermCfg @wp.kernel @@ -45,12 +48,66 @@ def joint_pos_target_l2(env: ManagerBasedRLEnv, out, target: float, asset_cfg: S ) -def survival_success_rate(env: ManagerBasedRLEnv, out) -> None: - """Write a zero reward for the stable Cartpole success-rate metric term. +@wp.kernel +def _survival_counts_kernel( + env_mask: wp.array(dtype=wp.bool), + time_outs: wp.array(dtype=wp.bool), + counts: wp.array(dtype=wp.int32), +): + env_index = wp.tid() + if env_mask[env_index]: + wp.atomic_add(counts, 0, 1) + if time_outs[env_index]: + wp.atomic_add(counts, 1, 1) + + +@wp.kernel +def _survival_rate_kernel( + counts: wp.array(dtype=wp.int32), + success_rate: wp.array(dtype=wp.float32), +): + if counts[0] > 0: + success_rate[0] = wp.float32(counts[1]) / wp.float32(counts[0]) + else: + success_rate[0] = 0.0 + + +class survival_success_rate(ManagerTermBase): + """Tracks episode survival as the success metric (Warp-first). - The stable term logs the episode survival rate during reset. That logging - requires a host readback of the time-out mask and is not compatible with - the Warp reward manager's CUDA-graph-captured reset. This twin preserves - the term's zero reward without performing the diagnostic readback. + Twin of :class:`isaaclab_tasks.core.cartpole.mdp.rewards.survival_success_rate`. + Returns zero reward (pure metric tracking). On reset, computes the fraction of + just-reset environments that timed out (survived the full episode) entirely + on-device and exposes it as ``Metrics/success_rate`` through the reward + manager's reset extras. Unlike the stable term, there is no host readback, so + the computation stays CUDA-graph capturable. """ - out.zero_() + + def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + # [0] number of just-reset envs, [1] number of those that timed out + self._counts_wp = wp.zeros((2,), dtype=wp.int32, device=env.device) + self._success_rate_wp = wp.zeros((1,), dtype=wp.float32, device=env.device) + # persistent 0-dim tensor view: kernels refresh the value on every reset/replay + self._reset_extras = {"Metrics/success_rate": wp.to_torch(self._success_rate_wp)[0]} + + def reset(self, env_mask: wp.array | None = None) -> dict[str, torch.Tensor]: + if env_mask is None: + env_mask = self._env.resolve_env_mask() + self._counts_wp.zero_() + wp.launch( + kernel=_survival_counts_kernel, + dim=self.num_envs, + inputs=[env_mask, self._env.termination_manager.time_outs_wp, self._counts_wp], + device=self.device, + ) + wp.launch( + kernel=_survival_rate_kernel, + dim=1, + inputs=[self._counts_wp, self._success_rate_wp], + device=self.device, + ) + return self._reset_extras + + def __call__(self, env: ManagerBasedRLEnv, out: wp.array) -> None: + out.zero_() From ac01de7e48bf54fdfda75a938cc74b98991d6d3b Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 15 Jul 2026 17:52:08 -0700 Subject: [PATCH 25/58] Resolve warp MDP twins through registered task routes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the frontend's hardcoded stable-to-experimental module table with a registry that experimental task packages populate at import time via register_mdp_route(), so adding a task family no longer requires editing the frontend module. Twin lookup previously keyed only off the symbol's defining module, which missed every twin that overrides a symbol defined in a core or shared package: the stable reach rewards live in isaaclab.envs.mdp and the humanoid observations in the shared locomotion package, so their task-specific warp twins were never consulted and the stable Humanoid, Ant, and Reach tasks could not adapt at all. Resolution now consults the warp mirror of the cfg's own task MDP namespace first (routed from the cfg class hierarchy), then the mirror of the symbol's package, then the shared fallback — mirroring how a stable cfg consumes its mdp namespace. This makes the stable Ant, Humanoid, and Reach joint-pos tasks adapt cleanly under --frontend=warp. Also convert test_frontend.py to pytest style and cover the new registry (longest-prefix match, conflict rejection, broken-target error, cfg-hierarchy resolution). --- .../isaaclab_experimental/envs/frontend.py | 244 +++--- .../test/envs/test_frontend.py | 727 +++++++++--------- .../manager_based/classic/ant/__init__.py | 9 + .../classic/cartpole/__init__.py | 5 + .../classic/humanoid/__init__.py | 6 + .../locomotion/velocity/__init__.py | 5 + .../manipulation/reach/__init__.py | 5 + 7 files changed, 564 insertions(+), 437 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index 0af62048df02..780685198e7c 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -41,6 +41,7 @@ "Workflow", "adapt_cfg_for_warp", "build", + "register_mdp_route", ] @@ -76,29 +77,54 @@ class FrontendIncompatibleError(RuntimeError): # the warp packages. Used by the direct-workflow guard. _WARP_ROOT_PREFIXES: tuple[str, ...] = ("isaaclab_experimental", "isaaclab_tasks_experimental") -# Stable task MDP packages and their Warp-first counterparts. The stable task -# package is organized by public task domain while the experimental package -# separates manager-based workflows by family, so the module trees are not -# one-to-one. Keep the routing at the package boundary and preserve any -# submodule suffix (for example, ``.rewards``). -_WARP_MDP_MODULE_ROUTES: tuple[tuple[str, str], ...] = ( - ( - "isaaclab_tasks.core.cartpole.mdp", - "isaaclab_tasks_experimental.manager_based.classic.cartpole.mdp", - ), - ( - "isaaclab_tasks.core.locomotion.humanoid.mdp", - "isaaclab_tasks_experimental.manager_based.classic.humanoid.mdp", - ), - ( - "isaaclab_tasks.core.velocity.mdp", - "isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp", - ), - ( - "isaaclab_tasks.core.reach.mdp", - "isaaclab_tasks_experimental.manager_based.manipulation.reach.mdp", - ), -) +# Stable package prefixes mapped to the warp MDP module that twins their terms. +# The stable task package is organized by public task domain while the +# experimental package separates manager-based workflows by family, so the +# module trees are not one-to-one. The registry is populated by the experimental +# task packages themselves via :func:`register_mdp_route` at import time, so +# adding a new task family never requires editing this module. +_WARP_MDP_MODULE_ROUTES: dict[str, str] = {} + +# Root package that provides task-specific warp MDP twins. Imported (lazily, at +# swap time) so its task packages run their register_mdp_route() calls even when +# the caller only imported the stable task package. +_TWIN_PROVIDER_PACKAGE = "isaaclab_tasks_experimental" + + +def register_mdp_route(stable_package: str, warp_mdp_module: str) -> None: + """Register the warp MDP module that twins a stable package's terms. + + Experimental task packages call this at import time to declare where the + warp twins of a stable task's MDP terms live. :func:`_swap_mdp` consults + the registry to resolve task-specific twins for ``--frontend=warp``. Twin + lookup is by symbol name on ``warp_mdp_module`` itself, mirroring how a + stable cfg consumes its ``mdp`` namespace — so the module should re-export + everything the task needs (task-specific twins plus the generic forwards). + + Two kinds of keys are useful: + + * A *task package* (e.g. ``"isaaclab_tasks.core.reach"``) routes every cfg + class defined under it. This also covers terms whose stable functions are + defined in core/shared packages but overridden by task-specific twins. + * An *MDP package* (e.g. ``"isaaclab_tasks.core.locomotion.humanoid.mdp"``) + routes symbols *defined* under it, which serves other tasks that borrow + those terms. + + Args: + stable_package: Stable package prefix to route (longest match wins). + warp_mdp_module: Warp MDP module providing the twins, e.g. + ``"isaaclab_tasks_experimental.manager_based.classic.cartpole.mdp"``. + + Raises: + ValueError: If ``stable_package`` is already routed to a different module. + """ + existing = _WARP_MDP_MODULE_ROUTES.get(stable_package) + if existing is not None and existing != warp_mdp_module: + raise ValueError( + f"MDP route conflict for {stable_package!r}: already routed to {existing!r}, got {warp_mdp_module!r}." + ) + _WARP_MDP_MODULE_ROUTES[stable_package] = warp_mdp_module + # Top-level cfg groups whose managers run warp-first. Only terms under these are # adapted (SceneEntityCfg promotion + MDP twin swap). The event manager is @@ -248,23 +274,26 @@ def _swap_mdp(cfg: Any, label: str) -> None: Iterates every term cfg in the tree (via :func:`_walk_terms`) and on each term swaps whichever of ``func`` / ``class_type`` is a stable-origin - callable. Twin lookup is name-based against the warp mirror of the *stable - symbol's own module* (:func:`_warp_mdp_modules`) plus the - :mod:`isaaclab_experimental.envs.mdp` fallback. Keying off the symbol's - module — not the cfg's — means a task that borrows another task's MDP (e.g. - ``manager_ant`` reuses ``manager_humanoid.mdp``) resolves to the right warp - twins without a per-task shim, mirroring the stable ``core/`` layout. Any - missing twin raises :class:`FrontendIncompatibleError` listing every - affected term — partial swaps would leave torch-style callables in the cfg - and the warp managers would call them with the wrong signature. - - The warp-side side declarations (``out_dim``, ``axes``, ``observation_type``) + callable. Twin lookup is name-based, in the order given by + :func:`_twin_modules`: the warp mirror of the cfg's own task MDP namespace + first (task twins win, even for symbols defined in core packages), then the + mirror of the symbol's defining package (covers terms borrowed from another + task family, e.g. the Ant task reusing Humanoid MDP terms), then the shared + :mod:`isaaclab_experimental.envs.mdp` fallback. Any missing twin raises + :class:`FrontendIncompatibleError` listing every affected term — partial + swaps would leave torch-style callables in the cfg and the warp managers + would call them with the wrong signature. + + The warp-side declarations (``out_dim``, ``axes``, ``observation_type``) that the warp managers need at init are *not* supplied by this swap; they travel with the warp twin function itself via its own ``@generic_io_descriptor_warp(out_dim=…)`` decorator. This function only substitutes the callable; the manager reads the new func's annotations when it parses the term cfg. """ + _ensure_twin_providers_imported() + + cfg_route_modules = _cfg_route_modules(cfg) module_cache: dict[str, list[ModuleType]] = {} searched: set[str] = set() @@ -280,7 +309,7 @@ def _swap_mdp(cfg: Any, label: str) -> None: continue origin = getattr(stable, "__module__", "") or "" if origin not in module_cache: - module_cache[origin] = _warp_mdp_modules(origin) + module_cache[origin] = _twin_modules(origin, cfg_route_modules) searched.update(m.__name__ for m in module_cache[origin]) twin = _resolve_warp_twin(stable.__name__, module_cache[origin]) if twin is None: @@ -337,62 +366,94 @@ def _assert_direct_warp_registration(task_id: str) -> None: ) -def _warp_mdp_modules(symbol_module: str) -> list[ModuleType]: - """Locate warp MDP modules to consult for a stable symbol's twin. +def _match_route(module: str) -> str | None: + """Return the routed warp MDP module for ``module``, longest prefix wins.""" + if not isinstance(module, str) or not module: + return None + best_key: str | None = None + for stable_prefix in _WARP_MDP_MODULE_ROUTES: + if module == stable_prefix or module.startswith(f"{stable_prefix}."): + if best_key is None or len(stable_prefix) > len(best_key): + best_key = stable_prefix + return _WARP_MDP_MODULE_ROUTES[best_key] if best_key is not None else None + + +def _import_routed_module(target: str) -> ModuleType: + """Import a registered route target; a broken registration is a hard error.""" + try: + return importlib.import_module(target) + except ImportError as exc: + raise FrontendIncompatibleError(f"registered warp MDP route target {target!r} failed to import: {exc}") from exc - The lookup is keyed off the *stable symbol's own module* - (``func.__module__`` / ``class_type.__module__``), e.g. - ``isaaclab_tasks.core.locomotion.humanoid.mdp.rewards``. The stable module - is routed to its manager-based experimental counterpart, then walked up to - the nearest importable module. Keying off the symbol's module — not the - cfg's — resolves the right twins even when a task borrows another task's - MDP (the Ant task reuses Humanoid MDP terms). - Order of preference: +def _cfg_route_modules(cfg: Any) -> list[ModuleType]: + """Warp MDP modules routed from the cfg's class hierarchy, subclass first. - 1. The routed Warp counterpart of ``symbol_module`` (or its nearest importable ancestor, - e.g. the ``...mdp`` package when the exact submodule isn't mirrored). - 2. The shared :mod:`isaaclab_experimental.envs.mdp` fallback (where - generic warp twins live). + A stable cfg consumes terms through its task's ``mdp`` namespace, which + re-exports symbols that may be *defined* in core or shared packages. The + warp mirror of that namespace is therefore the primary place to resolve + twins, and it is found by routing the modules of ``type(cfg).__mro__`` — + the subclass module first (robot-specific cfgs), then base-cfg modules + (task-family bases). """ modules: list[ModuleType] = [] - warp_mod = _warp_mdp_module_name(symbol_module) - if warp_mod is not None: - parts = warp_mod.split(".") - for depth in range(len(parts), 0, -1): - target = ".".join(parts[:depth]) - try: - modules.append(importlib.import_module(target)) - break - except ModuleNotFoundError as exc: - # Keep walking up while the missing module is part of the path - # we're probing; a genuine import error inside an existing - # module (missing third-party dep, etc.) must surface. - if exc.name and warp_mod.startswith(exc.name): - continue - raise - # Generic fallback. + for klass in type(cfg).__mro__: + target = _match_route(getattr(klass, "__module__", "") or "") + if target is None: + continue + module = _import_routed_module(target) + if module not in modules: + modules.append(module) + return modules + + +def _twin_modules(symbol_module: str, cfg_route_modules: list[ModuleType]) -> list[ModuleType]: + """Warp modules to consult for a stable symbol's twin, in preference order. + + 1. The warp mirrors of the cfg's own task MDP namespace + (:func:`_cfg_route_modules`) — task-specific twins win, including twins + for symbols defined in core/shared packages. + 2. The warp mirror routed from the symbol's defining package — covers terms + borrowed from another task family's MDP package. + 3. The shared :mod:`isaaclab_experimental.envs.mdp` fallback (where generic + warp twins live). + """ + modules = list(cfg_route_modules) + target = _match_route(symbol_module) + if target is not None: + module = _import_routed_module(target) + if module not in modules: + modules.append(module) fallback = "isaaclab_experimental.envs.mdp" try: - modules.append(importlib.import_module(fallback)) + fallback_module = importlib.import_module(fallback) except ModuleNotFoundError as exc: - if exc.name == fallback: - logger.warning("frontend.warp: fallback mdp module %r not importable", fallback) - else: + if exc.name != fallback: raise + logger.warning("frontend.warp: fallback mdp module %r not importable", fallback) + else: + if fallback_module not in modules: + modules.append(fallback_module) return modules -def _warp_mdp_module_name(symbol_module: str) -> str | None: - """Return the routed experimental MDP module for a stable task symbol.""" - if not isinstance(symbol_module, str): - return None - for stable_prefix, warp_prefix in _WARP_MDP_MODULE_ROUTES: - if symbol_module == stable_prefix: - return warp_prefix - if symbol_module.startswith(f"{stable_prefix}."): - return f"{warp_prefix}{symbol_module[len(stable_prefix) :]}" - return None +def _ensure_twin_providers_imported() -> None: + """Import the twin-provider package so its MDP routes are registered. + + Route registration happens as an import side effect of the experimental + task packages. A caller running ``--frontend=warp`` on a stable task id + has no other reason to import :mod:`isaaclab_tasks_experimental`, so the + swap triggers it here. A missing provider package is not an error by + itself — the swap then fails with the explicit missing-twin report. + """ + try: + importlib.import_module(_TWIN_PROVIDER_PACKAGE) + except ModuleNotFoundError as exc: + # Only swallow "the provider package is not installed"; a genuine import + # error inside an existing package must surface. + if exc.name != _TWIN_PROVIDER_PACKAGE: + raise + logger.warning("frontend.warp: twin provider package %r is not installed", _TWIN_PROVIDER_PACKAGE) def _resolve_warp_twin(name: str, modules: Iterable[ModuleType]) -> Any | None: @@ -429,13 +490,20 @@ def _walk_terms(node: Any, path: tuple[str, ...] = ()) -> Iterator[tuple[tuple[s * Match (a term cfg instance): yield ``(path, node)`` and stop — do not descend into ``term.params`` / ``term.func`` / ``term.class_type``. - * Configclass: don't yield; recurse into every non-underscore attribute, - extending the path. ``observations``, ``rewards``, ``events``, ``actions``, - sub-groups like ``observations.policy`` / ``observations.perception``, and - anything nested deeper are reached transparently. + * Configclass: don't yield; recurse into every non-underscore *instance* + attribute (``vars(node)``), extending the path. ``observations``, + ``rewards``, ``events``, ``actions``, sub-groups like + ``observations.policy`` / ``observations.perception``, and anything + nested deeper are reached transparently. * Anything else (plain Python data, callables, non-configclass objects): stop. No yield, no recursion. + Iterating the instance ``__dict__`` mirrors how the warp managers consume + group cfgs (they iterate ``cfg.__dict__.items()``), so the walker sees + exactly the terms the managers will see — including terms assigned in + ``__post_init__`` — while never descending into methods or nested class + objects, which live on the class rather than the instance. + Driven entirely by type — no attribute names are hardcoded — so future cfg layouts (extra observation groups, new nesting, etc.) are picked up automatically as long as their terms subclass one of the term base cfgs. @@ -447,13 +515,7 @@ def _walk_terms(node: Any, path: tuple[str, ...] = ()) -> Iterator[tuple[tuple[s return if not hasattr(node, "__dataclass_fields__"): return - for name in dir(node): - if name.startswith("_"): - continue - try: - value = getattr(node, name, None) - except Exception: # noqa: BLE001 — defensive; some descriptors can raise on attribute access - continue - if value is None: + for name, value in vars(node).items(): + if name.startswith("_") or value is None: continue yield from _walk_terms(value, path + (name,)) diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py index a7af62ee146e..03e39287e679 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend.py +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -10,12 +10,13 @@ * :class:`Frontend` / :class:`Workflow` enum surface. * :func:`SceneEntityCfg.from_stable` field copy. * :func:`_require_newton_physics` hard-check. -* :func:`_walk_terms` recursive ManagerTermBaseCfg discovery. +* :func:`_walk_terms` recursive ManagerTermBaseCfg discovery over instance + attributes (nested class objects are not descended into). +* :func:`register_mdp_route` / :func:`_match_route` route registry. +* :func:`_cfg_route_modules` cfg-hierarchy route resolution. * :func:`_promote_scene_entity_cfgs` walks ``term.params`` dicts. * :func:`_swap_mdp` swaps ``func`` *and* ``class_type``; raises with a path list when twins are missing. -* :func:`_warp_mdp_module_name` routes stable task modules to the current - experimental manager-based package layout. * :func:`_resolve_warp_twin` rejects stable-origin re-exports. * :func:`_assert_direct_warp_registration` accepts warp-rooted entry points and rejects stable ones. @@ -23,25 +24,21 @@ from __future__ import annotations +import contextlib +import importlib import types -import unittest from typing import Any from unittest.mock import patch import gymnasium as gym +import isaaclab_experimental.envs.frontend as fe +import pytest from isaaclab_experimental.envs.frontend import ( Frontend, FrontendIncompatibleError, Workflow, - _assert_direct_warp_registration, - _is_swap_candidate, - _promote_scene_entity_cfgs, - _require_newton_physics, - _resolve_warp_twin, - _swap_mdp, - _walk_terms, - _warp_mdp_module_name, build, + register_mdp_route, ) from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as WarpSceneEntityCfg from isaaclab_newton.physics import NewtonCfg @@ -53,133 +50,13 @@ from isaaclab.utils.configclass import configclass # ====================================================================== -# Enums -# ====================================================================== - - -class TestEnums(unittest.TestCase): - def test_frontend_values(self): - self.assertEqual(Frontend.TORCH, "torch") - self.assertEqual(Frontend.WARP, "warp") - - def test_frontend_coercion(self): - self.assertIs(Frontend("torch"), Frontend.TORCH) - self.assertIs(Frontend("warp"), Frontend.WARP) - with self.assertRaises(ValueError): - Frontend("kit") - - def test_workflow_values(self): - self.assertEqual(Workflow.MANAGER_BASED, "manager_based") - self.assertEqual(Workflow.DIRECT, "direct") - - -class TestBuild(unittest.TestCase): - def test_warp_manager_build_adapts_cfg_before_construction(self): - import isaaclab_experimental.envs as envs - import isaaclab_experimental.envs.frontend as frontend_module - - cfg = object.__new__(ManagerBasedRLEnvCfg) - expected_env = object() - calls: list[tuple[str, Any]] = [] - - def fake_adapt(received_cfg: Any) -> None: - calls.append(("adapt", received_cfg)) - - def fake_env(*, cfg: Any, **kwargs: Any) -> Any: - calls.append(("construct", cfg)) - return expected_env - - with ( - patch.object(frontend_module, "adapt_cfg_for_warp", fake_adapt), - patch.object(envs, "ManagerBasedRLEnvWarp", fake_env), - ): - env = build("warp", cfg, "Isaac-Test") - - self.assertIs(env, expected_env) - self.assertEqual(calls, [("adapt", cfg), ("construct", cfg)]) - - -# ====================================================================== -# SceneEntityCfg.from_stable -# ====================================================================== - - -class TestFromStable(unittest.TestCase): - def test_copies_minimum_fields(self): - stable = StableSceneEntityCfg(name="robot") - warp = WarpSceneEntityCfg.from_stable(stable) - self.assertIsInstance(warp, WarpSceneEntityCfg) - self.assertEqual(warp.name, "robot") - # Warp-only fields stay None until :meth:`resolve` runs. - self.assertIsNone(warp.joint_mask) - self.assertIsNone(warp.joint_ids_wp) - self.assertIsNone(warp.body_ids_wp) - - def test_copies_all_selection_fields(self): - stable = StableSceneEntityCfg( - name="robot", - joint_names=["lf_hip"], - joint_ids=[0, 1, 2], - fixed_tendon_names=["tendon_a"], - fixed_tendon_ids=[5], - body_names=["base", "lf_foot"], - body_ids=[0, 4], - object_collection_names=["objs"], - object_collection_ids=[7], - preserve_order=True, - ) - warp = WarpSceneEntityCfg.from_stable(stable) - for field in ( - "name", - "joint_names", - "joint_ids", - "fixed_tendon_names", - "fixed_tendon_ids", - "body_names", - "body_ids", - "object_collection_names", - "object_collection_ids", - "preserve_order", - ): - self.assertEqual(getattr(warp, field), getattr(stable, field), msg=f"field {field!r} mismatch") - - -# ====================================================================== -# _require_newton_physics -# ====================================================================== - - -class TestRequireNewtonPhysics(unittest.TestCase): - def _cfg_with(self, physics: Any) -> Any: - cfg = types.SimpleNamespace() - cfg.sim = types.SimpleNamespace(physics=physics) - return cfg - - def test_passes_for_newton(self): - cfg = self._cfg_with(NewtonCfg()) - _require_newton_physics(cfg, "Isaac-Test-v0") # no raise - - def test_rejects_physx(self): - cfg = self._cfg_with(PhysxCfg()) - with self.assertRaises(FrontendIncompatibleError) as exc: - _require_newton_physics(cfg, "Isaac-Test-v0") - self.assertIn("presets=newton_mjwarp", str(exc.exception)) - self.assertIn("PhysxCfg", str(exc.exception)) - - def test_rejects_none(self): - cfg = self._cfg_with(None) - with self.assertRaises(FrontendIncompatibleError): - _require_newton_physics(cfg, "Isaac-Test-v0") - - -# ====================================================================== -# Configclass fixtures for the walker / swap tests. -# ====================================================================== +# Fixtures: fake stable/warp symbols and configclass trees. # # These mirror the real cfg shape so :func:`_walk_terms` descends into them # (it only descends into objects with ``__dataclass_fields__``). Term cfgs # use the real :class:`EventTermCfg`/:class:`RewardTermCfg`/:class:`ObservationTermCfg` # so the walker's ``isinstance(ManagerTermBaseCfg)`` discriminator yields them. +# ====================================================================== def _stable_func(env, **params): @@ -256,116 +133,246 @@ def _term(func=None, params: dict | None = None) -> RewardTermCfg: return RewardTermCfg(func=func or _stable_func, weight=1.0, params=params or {}) +@pytest.fixture +def fake_routes(monkeypatch: pytest.MonkeyPatch) -> dict[str, str]: + """Isolate the module-level route registry for a test.""" + routes: dict[str, str] = {} + monkeypatch.setattr(fe, "_WARP_MDP_MODULE_ROUTES", routes) + return routes + + +# ====================================================================== +# Enums +# ====================================================================== + + +def test_frontend_values(): + assert Frontend.TORCH == "torch" + assert Frontend.WARP == "warp" + + +def test_frontend_coercion(): + assert Frontend("torch") is Frontend.TORCH + assert Frontend("warp") is Frontend.WARP + with pytest.raises(ValueError): + Frontend("kit") + + +def test_workflow_values(): + assert Workflow.MANAGER_BASED == "manager_based" + assert Workflow.DIRECT == "direct" + + +# ====================================================================== +# build() +# ====================================================================== + + +def test_warp_manager_build_adapts_cfg_before_construction(): + import isaaclab_experimental.envs as envs + + cfg = object.__new__(ManagerBasedRLEnvCfg) + expected_env = object() + calls: list[tuple[str, Any]] = [] + + def fake_adapt(received_cfg: Any) -> None: + calls.append(("adapt", received_cfg)) + + def fake_env(*, cfg: Any, **kwargs: Any) -> Any: + calls.append(("construct", cfg)) + return expected_env + + with ( + patch.object(fe, "adapt_cfg_for_warp", fake_adapt), + patch.object(envs, "ManagerBasedRLEnvWarp", fake_env), + ): + env = build("warp", cfg, "Isaac-Test") + + assert env is expected_env + assert calls == [("adapt", cfg), ("construct", cfg)] + + +# ====================================================================== +# SceneEntityCfg.from_stable +# ====================================================================== + + +def test_from_stable_copies_minimum_fields(): + stable = StableSceneEntityCfg(name="robot") + warp = WarpSceneEntityCfg.from_stable(stable) + assert isinstance(warp, WarpSceneEntityCfg) + assert warp.name == "robot" + # Warp-only fields stay None until :meth:`resolve` runs. + assert warp.joint_mask is None + assert warp.joint_ids_wp is None + assert warp.body_ids_wp is None + + +def test_from_stable_copies_all_selection_fields(): + stable = StableSceneEntityCfg( + name="robot", + joint_names=["lf_hip"], + joint_ids=[0, 1, 2], + fixed_tendon_names=["tendon_a"], + fixed_tendon_ids=[5], + body_names=["base", "lf_foot"], + body_ids=[0, 4], + object_collection_names=["objs"], + object_collection_ids=[7], + preserve_order=True, + ) + warp = WarpSceneEntityCfg.from_stable(stable) + for field in ( + "name", + "joint_names", + "joint_ids", + "fixed_tendon_names", + "fixed_tendon_ids", + "body_names", + "body_ids", + "object_collection_names", + "object_collection_ids", + "preserve_order", + ): + assert getattr(warp, field) == getattr(stable, field), f"field {field!r} mismatch" + + +# ====================================================================== +# _require_newton_physics +# ====================================================================== + + +def _cfg_with_physics(physics: Any) -> Any: + cfg = types.SimpleNamespace() + cfg.sim = types.SimpleNamespace(physics=physics) + return cfg + + +def test_require_newton_passes_for_newton(): + fe._require_newton_physics(_cfg_with_physics(NewtonCfg()), "Isaac-Test-v0") # no raise + + +def test_require_newton_rejects_physx(): + with pytest.raises(FrontendIncompatibleError) as exc: + fe._require_newton_physics(_cfg_with_physics(PhysxCfg()), "Isaac-Test-v0") + assert "presets=newton_mjwarp" in str(exc.value) + assert "PhysxCfg" in str(exc.value) + + +def test_require_newton_rejects_none(): + with pytest.raises(FrontendIncompatibleError): + fe._require_newton_physics(_cfg_with_physics(None), "Isaac-Test-v0") + + # ====================================================================== # _walk_terms # ====================================================================== -class TestWalkTerms(unittest.TestCase): - def test_yields_each_term_with_its_path(self): - cfg = _CfgFixture( - rewards=_RewardsCfg(r1=_term(), r2=_term()), - events=_EventsCfg(e1=EventTermCfg(func=_stable_func, mode="reset")), - ) - # Configclass instances aren't hashable, so collect paths only. - paths = {".".join(p) for p, _ in _walk_terms(cfg)} - self.assertEqual(paths, {"rewards.r1", "rewards.r2", "events.e1"}) - - def test_descends_into_obs_subgroups(self): - cfg = _CfgFixture( - observations=_ObservationsCfg( - policy=_PolicyObsGroup(o1=ObservationTermCfg(func=_stable_func)), - perception=_ExtraObsGroup(o3=ObservationTermCfg(func=_stable_func)), - ), - ) - paths = {".".join(p) for p, _ in _walk_terms(cfg)} - # Discovery is purely type-driven; no obs group name is hardcoded. - self.assertEqual(paths, {"observations.policy.o1", "observations.perception.o3"}) +def test_walk_terms_yields_each_term_with_its_path(): + cfg = _CfgFixture( + rewards=_RewardsCfg(r1=_term(), r2=_term()), + events=_EventsCfg(e1=EventTermCfg(func=_stable_func, mode="reset")), + ) + # Configclass instances aren't hashable, so collect paths only. + paths = {".".join(p) for p, _ in fe._walk_terms(cfg)} + assert paths == {"rewards.r1", "rewards.r2", "events.e1"} + + +def test_walk_terms_descends_into_obs_subgroups(): + cfg = _CfgFixture( + observations=_ObservationsCfg( + policy=_PolicyObsGroup(o1=ObservationTermCfg(func=_stable_func)), + perception=_ExtraObsGroup(o3=ObservationTermCfg(func=_stable_func)), + ), + ) + paths = {".".join(p) for p, _ in fe._walk_terms(cfg)} + # Discovery is purely type-driven; no obs group name is hardcoded. + assert paths == {"observations.policy.o1", "observations.perception.o3"} + + +def test_walk_terms_stops_at_terms(): + # The walker must not descend into term.params / term.func — yields the term itself. + nested_se_cfg = StableSceneEntityCfg(name="robot") + cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(params={"asset_cfg": nested_se_cfg}))) + terms = list(fe._walk_terms(cfg)) + assert len(terms) == 1 + _, term = terms[0] + assert isinstance(term, RewardTermCfg) + + +def test_walk_terms_skips_non_configclass_attrs(): + # A namespace without __dataclass_fields__ is not descended into. + cfg = types.SimpleNamespace(some_plain_attr="hello") + assert list(fe._walk_terms(cfg)) == [] - def test_stops_at_terms(self): - # The walker must not descend into term.params / term.func — yields the term itself. - nested_se_cfg = StableSceneEntityCfg(name="robot") - cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(params={"asset_cfg": nested_se_cfg}))) - terms = list(_walk_terms(cfg)) - self.assertEqual(len(terms), 1) - _, term = terms[0] - self.assertIsInstance(term, RewardTermCfg) - def test_skips_non_configclass_attrs(self): - # A namespace without __dataclass_fields__ is not descended into. - cfg = types.SimpleNamespace(some_plain_attr="hello") - self.assertEqual(list(_walk_terms(cfg)), []) +def test_walk_terms_skips_none_subtrees(): + cfg = _CfgFixture(rewards=None, events=None) + assert list(fe._walk_terms(cfg)) == [] - def test_skips_none_subtrees(self): - cfg = _CfgFixture(rewards=None, events=None) - self.assertEqual(list(_walk_terms(cfg)), []) + +@configclass +class _GroupWithNestedClass: + """Mirrors the real ``ObservationsCfg.PolicyCfg`` nested-class pattern.""" + + @configclass + class TemplateCfg: + t1: ObservationTermCfg | None = None + + o1: ObservationTermCfg | None = None + + +def test_walk_terms_ignores_nested_class_objects(): + # The walker iterates instance attributes only (mirroring the warp managers' + # ``cfg.__dict__`` consumption), so a nested class object — reachable as an + # attribute on the instance — must never contribute paths. + cfg = _GroupWithNestedClass(o1=ObservationTermCfg(func=_stable_func)) + paths = {".".join(p) for p, _ in fe._walk_terms(cfg)} + assert paths == {"o1"} # ====================================================================== -# _promote_scene_entity_cfgs +# Route registry # ====================================================================== -class TestPromoteSceneEntityCfgs(unittest.TestCase): - def test_promotes_in_params(self): - cfg = _CfgFixture( - rewards=_RewardsCfg( - r1=_term(params={"asset_cfg": StableSceneEntityCfg(name="robot", joint_names=["lf_hip"]), "scale": 1.0}) - ) - ) - _promote_scene_entity_cfgs(cfg) - promoted = cfg.rewards.r1.params["asset_cfg"] - self.assertIsInstance(promoted, WarpSceneEntityCfg) - self.assertEqual(promoted.name, "robot") - self.assertEqual(promoted.joint_names, ["lf_hip"]) - # Non-SceneEntityCfg params are untouched. - self.assertEqual(cfg.rewards.r1.params["scale"], 1.0) - - def test_skips_already_warp(self): - # configclass init deep-copies params, so identity won't hold across - # construction; what we actually want to assert is "no re-promotion": - # the asset_cfg remains a WarpSceneEntityCfg (i.e., wasn't passed - # back through `from_stable`). - warp = WarpSceneEntityCfg(name="robot", joint_names=["lf_hip"]) - cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(params={"asset_cfg": warp}))) - before = cfg.rewards.r1.params["asset_cfg"] - _promote_scene_entity_cfgs(cfg) - after = cfg.rewards.r1.params["asset_cfg"] - self.assertIsInstance(after, WarpSceneEntityCfg) - # The asset_cfg object was not replaced by another from_stable call. - self.assertIs(after, before) - - def test_walks_all_term_groups(self): - cfg = _CfgFixture( - rewards=_RewardsCfg(r1=_term(params={"asset_cfg": StableSceneEntityCfg(name="r")})), - events=_EventsCfg( - e1=EventTermCfg(func=_stable_func, mode="reset", params={"asset_cfg": StableSceneEntityCfg(name="e")}) - ), - observations=_ObservationsCfg( - policy=_PolicyObsGroup( - o1=ObservationTermCfg(func=_stable_func, params={"asset_cfg": StableSceneEntityCfg(name="o-pol")}) - ), - perception=_ExtraObsGroup( - o3=ObservationTermCfg(func=_stable_func, params={"asset_cfg": StableSceneEntityCfg(name="o-per")}) - ), - ), - curriculum=_CurriculumCfg( - c1=EventTermCfg(func=_stable_func, mode="reset", params={"asset_cfg": StableSceneEntityCfg(name="c")}) - ), - ) - _promote_scene_entity_cfgs(cfg) - # Warp-managed groups (rewards, observations incl. sub-groups, events) are promoted. - self.assertIsInstance(cfg.rewards.r1.params["asset_cfg"], WarpSceneEntityCfg) - self.assertIsInstance(cfg.observations.policy.o1.params["asset_cfg"], WarpSceneEntityCfg) - # The perception sub-group is reached even though its attribute name is - # not hardcoded in the framework. - self.assertIsInstance(cfg.observations.perception.o3.params["asset_cfg"], WarpSceneEntityCfg) - # The event manager is warp-first, so its terms are promoted too. - self.assertIsInstance(cfg.events.e1.params["asset_cfg"], WarpSceneEntityCfg) - # Curriculum runs on the stable (torch) manager, so its SceneEntityCfg is - # left untouched — promoting it would hand a warp variant to a torch manager. - self.assertIsInstance(cfg.curriculum.c1.params["asset_cfg"], StableSceneEntityCfg) - self.assertNotIsInstance(cfg.curriculum.c1.params["asset_cfg"], WarpSceneEntityCfg) +def test_register_mdp_route_and_longest_prefix_match(fake_routes: dict[str, str]): + register_mdp_route("isaaclab_tasks.core.locomotion", "warp.locomotion.mdp") + register_mdp_route("isaaclab_tasks.core.locomotion.humanoid", "warp.humanoid.mdp") + + # Longest registered prefix wins. + assert fe._match_route("isaaclab_tasks.core.locomotion.humanoid.mdp.rewards") == "warp.humanoid.mdp" + assert fe._match_route("isaaclab_tasks.core.locomotion.ant.mdp") == "warp.locomotion.mdp" + # Exact package match works too. + assert fe._match_route("isaaclab_tasks.core.locomotion.humanoid") == "warp.humanoid.mdp" + # Prefix matching is per package segment, not per character. + assert fe._match_route("isaaclab_tasks.core.locomotion_extras.mdp") is None + assert fe._match_route("isaaclab.envs.mdp.rewards") is None + + +def test_register_mdp_route_is_idempotent(fake_routes: dict[str, str]): + register_mdp_route("isaaclab_tasks.core.cartpole", "warp.cartpole.mdp") + register_mdp_route("isaaclab_tasks.core.cartpole", "warp.cartpole.mdp") # no raise + assert fake_routes == {"isaaclab_tasks.core.cartpole": "warp.cartpole.mdp"} + + +def test_register_mdp_route_rejects_conflicts(fake_routes: dict[str, str]): + register_mdp_route("isaaclab_tasks.core.cartpole", "warp.cartpole.mdp") + with pytest.raises(ValueError): + register_mdp_route("isaaclab_tasks.core.cartpole", "warp.other.mdp") + + +def test_cfg_route_modules_resolves_cfg_class_module(fake_routes: dict[str, str]): + # Route this test module's own name so the fixture cfg class matches. + register_mdp_route(__name__, "math") + assert fe._cfg_route_modules(_CfgFixture()) == [importlib.import_module("math")] + + +def test_cfg_route_modules_raises_on_broken_target(fake_routes: dict[str, str]): + register_mdp_route(__name__, "definitely_not_an_importable_module") + with pytest.raises(FrontendIncompatibleError): + fe._cfg_route_modules(_CfgFixture()) # ====================================================================== @@ -379,116 +386,149 @@ class _FakeMdpModule: __name__ = "test_fake_warp_mdp" -class TestSwapMdp(unittest.TestCase): - def _patch_modules(self, name_to_symbol: dict[str, Any]) -> _FakeMdpModule: - m = _FakeMdpModule() - for name, sym in name_to_symbol.items(): - setattr(m, name, sym) - return m - - def _patched_warp_mdp_modules(self, modules: list[Any]): - import isaaclab_experimental.envs.frontend as fe - - self._orig = fe._warp_mdp_modules - fe._warp_mdp_modules = lambda task_id: modules # type: ignore[assignment] - - def setUp(self) -> None: - self._orig = None - - def tearDown(self) -> None: - if self._orig is not None: - import isaaclab_experimental.envs.frontend as fe - - fe._warp_mdp_modules = self._orig - - def test_swaps_func_and_class_type(self): - fake = self._patch_modules({"_stable_func": _warp_twin_func, "_StableActionCls": _WarpActionCls}) - self._patched_warp_mdp_modules([fake]) - term_reward = _term(func=_stable_func) - term_action = _term() - term_action.class_type = _StableActionCls # set attr to exercise class_type swap - cfg = _CfgFixture(rewards=_RewardsCfg(r1=term_reward, r2=term_action)) - _swap_mdp(cfg, "Isaac-Test-v0") - self.assertIs(cfg.rewards.r1.func, _warp_twin_func) - self.assertIs(cfg.rewards.r2.class_type, _WarpActionCls) - - def test_missing_twin_raises_with_path_list(self): - fake = self._patch_modules({}) - self._patched_warp_mdp_modules([fake]) - cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(func=_stable_func))) - with self.assertRaises(FrontendIncompatibleError) as exc: - _swap_mdp(cfg, "Isaac-Test-v0") - msg = str(exc.exception) - self.assertIn("rewards.r1.func", msg) - self.assertIn("_stable_func", msg) - # The cfg term wasn't mutated for the missing twin. - self.assertIs(cfg.rewards.r1.func, _stable_func) - - def test_skips_already_warp(self): - fake = self._patch_modules({}) - self._patched_warp_mdp_modules([fake]) - cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(func=_warp_twin_func))) - _swap_mdp(cfg, "Isaac-Test-v0") # no raise - self.assertIs(cfg.rewards.r1.func, _warp_twin_func) +def _fake_mdp_module(name_to_symbol: dict[str, Any]) -> _FakeMdpModule: + module = _FakeMdpModule() + for name, symbol in name_to_symbol.items(): + setattr(module, name, symbol) + return module + + +def _patch_twin_modules(monkeypatch: pytest.MonkeyPatch, modules: list[Any]) -> None: + """Pin twin resolution to ``modules`` and keep unit tests hermetic.""" + monkeypatch.setattr(fe, "_ensure_twin_providers_imported", lambda: None) + monkeypatch.setattr(fe, "_cfg_route_modules", lambda cfg: []) + monkeypatch.setattr(fe, "_twin_modules", lambda symbol_module, cfg_route_modules: modules) + + +def test_swap_mdp_swaps_func_and_class_type(monkeypatch: pytest.MonkeyPatch): + fake = _fake_mdp_module({"_stable_func": _warp_twin_func, "_StableActionCls": _WarpActionCls}) + _patch_twin_modules(monkeypatch, [fake]) + term_reward = _term(func=_stable_func) + term_action = _term() + term_action.class_type = _StableActionCls # set attr to exercise class_type swap + cfg = _CfgFixture(rewards=_RewardsCfg(r1=term_reward, r2=term_action)) + fe._swap_mdp(cfg, "Isaac-Test-v0") + assert cfg.rewards.r1.func is _warp_twin_func + assert cfg.rewards.r2.class_type is _WarpActionCls + + +def test_swap_mdp_missing_twin_raises_with_path_list(monkeypatch: pytest.MonkeyPatch): + _patch_twin_modules(monkeypatch, [_fake_mdp_module({})]) + cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(func=_stable_func))) + with pytest.raises(FrontendIncompatibleError) as exc: + fe._swap_mdp(cfg, "Isaac-Test-v0") + msg = str(exc.value) + assert "rewards.r1.func" in msg + assert "_stable_func" in msg + # The cfg term wasn't mutated for the missing twin. + assert cfg.rewards.r1.func is _stable_func + + +def test_swap_mdp_skips_already_warp(monkeypatch: pytest.MonkeyPatch): + _patch_twin_modules(monkeypatch, [_fake_mdp_module({})]) + cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(func=_warp_twin_func))) + fe._swap_mdp(cfg, "Isaac-Test-v0") # no raise + assert cfg.rewards.r1.func is _warp_twin_func # ====================================================================== -# Twin module routing and resolution +# _promote_scene_entity_cfgs # ====================================================================== -class TestWarpMdpModuleName(unittest.TestCase): - def test_routes_task_submodule_to_current_experimental_layout(self): - self.assertEqual( - _warp_mdp_module_name("isaaclab_tasks.core.cartpole.mdp.rewards"), - "isaaclab_tasks_experimental.manager_based.classic.cartpole.mdp.rewards", +def test_promote_scene_entity_cfgs_promotes_in_params(): + cfg = _CfgFixture( + rewards=_RewardsCfg( + r1=_term(params={"asset_cfg": StableSceneEntityCfg(name="robot", joint_names=["lf_hip"]), "scale": 1.0}) ) + ) + fe._promote_scene_entity_cfgs(cfg) + promoted = cfg.rewards.r1.params["asset_cfg"] + assert isinstance(promoted, WarpSceneEntityCfg) + assert promoted.name == "robot" + assert promoted.joint_names == ["lf_hip"] + # Non-SceneEntityCfg params are untouched. + assert cfg.rewards.r1.params["scale"] == 1.0 + + +def test_promote_scene_entity_cfgs_skips_already_warp(): + # configclass init deep-copies params, so identity won't hold across + # construction; what we actually want to assert is "no re-promotion": + # the asset_cfg remains a WarpSceneEntityCfg (i.e., wasn't passed + # back through `from_stable`). + warp = WarpSceneEntityCfg(name="robot", joint_names=["lf_hip"]) + cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(params={"asset_cfg": warp}))) + before = cfg.rewards.r1.params["asset_cfg"] + fe._promote_scene_entity_cfgs(cfg) + after = cfg.rewards.r1.params["asset_cfg"] + assert isinstance(after, WarpSceneEntityCfg) + # The asset_cfg object was not replaced by another from_stable call. + assert after is before + + +def test_promote_scene_entity_cfgs_walks_all_term_groups(): + cfg = _CfgFixture( + rewards=_RewardsCfg(r1=_term(params={"asset_cfg": StableSceneEntityCfg(name="r")})), + events=_EventsCfg( + e1=EventTermCfg(func=_stable_func, mode="reset", params={"asset_cfg": StableSceneEntityCfg(name="e")}) + ), + observations=_ObservationsCfg( + policy=_PolicyObsGroup( + o1=ObservationTermCfg(func=_stable_func, params={"asset_cfg": StableSceneEntityCfg(name="o-pol")}) + ), + perception=_ExtraObsGroup( + o3=ObservationTermCfg(func=_stable_func, params={"asset_cfg": StableSceneEntityCfg(name="o-per")}) + ), + ), + curriculum=_CurriculumCfg( + c1=EventTermCfg(func=_stable_func, mode="reset", params={"asset_cfg": StableSceneEntityCfg(name="c")}) + ), + ) + fe._promote_scene_entity_cfgs(cfg) + # Warp-managed groups (rewards, observations incl. sub-groups, events) are promoted. + assert isinstance(cfg.rewards.r1.params["asset_cfg"], WarpSceneEntityCfg) + assert isinstance(cfg.observations.policy.o1.params["asset_cfg"], WarpSceneEntityCfg) + # The perception sub-group is reached even though its attribute name is + # not hardcoded in the framework. + assert isinstance(cfg.observations.perception.o3.params["asset_cfg"], WarpSceneEntityCfg) + # The event manager is warp-first, so its terms are promoted too. + assert isinstance(cfg.events.e1.params["asset_cfg"], WarpSceneEntityCfg) + # Curriculum runs on the stable (torch) manager, so its SceneEntityCfg is + # left untouched — promoting it would hand a warp variant to a torch manager. + assert isinstance(cfg.curriculum.c1.params["asset_cfg"], StableSceneEntityCfg) + assert not isinstance(cfg.curriculum.c1.params["asset_cfg"], WarpSceneEntityCfg) - def test_routes_exact_task_mdp_package(self): - self.assertEqual( - _warp_mdp_module_name("isaaclab_tasks.core.velocity.mdp"), - "isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp", - ) - def test_rejects_unregistered_task_module(self): - self.assertIsNone(_warp_mdp_module_name("isaaclab_tasks.core.unknown.mdp")) +# ====================================================================== +# Twin resolution and swap-candidate heuristic +# ====================================================================== + - def test_rejects_non_task_module(self): - self.assertIsNone(_warp_mdp_module_name("isaaclab.envs.mdp.rewards")) +def test_resolve_warp_twin_accepts_warp_origin(): + module = types.SimpleNamespace(foo=_warp_twin_func) + assert fe._resolve_warp_twin("foo", [module]) is _warp_twin_func -class TestResolveWarpTwin(unittest.TestCase): - def test_accepts_warp_origin(self): - m = types.SimpleNamespace() - m.foo = _warp_twin_func - result = _resolve_warp_twin("foo", [m]) - self.assertIs(result, _warp_twin_func) +def test_resolve_warp_twin_rejects_stable_origin(): + module = types.SimpleNamespace(foo=_stable_func) # same name, stable origin + assert fe._resolve_warp_twin("foo", [module]) is None - def test_rejects_stable_origin(self): - m = types.SimpleNamespace() - m.foo = _stable_func # same name, stable origin - self.assertIsNone(_resolve_warp_twin("foo", [m])) - def test_returns_none_when_absent(self): - m = types.SimpleNamespace() - self.assertIsNone(_resolve_warp_twin("missing", [m])) +def test_resolve_warp_twin_returns_none_when_absent(): + assert fe._resolve_warp_twin("missing", [types.SimpleNamespace()]) is None -# ====================================================================== -# Swap candidate heuristic -# ====================================================================== +def test_is_swap_candidate_stable_callable(): + assert fe._is_swap_candidate(_stable_func) -class TestIsSwapCandidate(unittest.TestCase): - def test_stable_callable_is_candidate(self): - self.assertTrue(_is_swap_candidate(_stable_func)) +def test_is_swap_candidate_rejects_warp_callable(): + assert not fe._is_swap_candidate(_warp_twin_func) - def test_warp_callable_is_not(self): - self.assertFalse(_is_swap_candidate(_warp_twin_func)) - def test_non_callable_is_not(self): - self.assertFalse(_is_swap_candidate(42)) - self.assertFalse(_is_swap_candidate("string")) +def test_is_swap_candidate_rejects_non_callables(): + assert not fe._is_swap_candidate(42) + assert not fe._is_swap_candidate("string") # ====================================================================== @@ -497,37 +537,32 @@ def test_non_callable_is_not(self): _DIRECT_TEST_TASKS = { - "_Frontend-Test-Warp-Direct-v0": ("isaaclab_tasks_experimental.fake:DirectEnv", True), - "_Frontend-Test-Stable-Direct-v0": ("isaaclab_tasks.fake:DirectEnv", False), + "_Frontend-Test-Warp-Direct-v0": "isaaclab_tasks_experimental.fake:DirectEnv", + "_Frontend-Test-Stable-Direct-v0": "isaaclab_tasks.fake:DirectEnv", } -class TestAssertDirectWarpRegistration(unittest.TestCase): - @classmethod - def setUpClass(cls): - # Register stub tasks so ``gym.spec`` resolves them. Entry-point strings - # are never invoked (the guard only inspects them), so a fake import - # path is fine. - cls._registered: list[str] = [] - for task_id, (ep, _) in _DIRECT_TEST_TASKS.items(): - try: - gym.register(id=task_id, entry_point=ep, disable_env_checker=True) - cls._registered.append(task_id) - except gym.error.Error: - pass +@pytest.fixture(scope="module", autouse=True) +def _register_direct_stub_tasks(): + # Register stub tasks so ``gym.spec`` resolves them. Entry-point strings + # are never invoked (the guard only inspects them), so a fake import + # path is fine. + for task_id, entry_point in _DIRECT_TEST_TASKS.items(): + with contextlib.suppress(gym.error.Error): + gym.register(id=task_id, entry_point=entry_point, disable_env_checker=True) + yield + - def test_accepts_warp_rooted(self): - _assert_direct_warp_registration("_Frontend-Test-Warp-Direct-v0") # no raise +def test_direct_guard_accepts_warp_rooted(): + fe._assert_direct_warp_registration("_Frontend-Test-Warp-Direct-v0") # no raise - def test_rejects_stable_rooted(self): - with self.assertRaises(FrontendIncompatibleError) as exc: - _assert_direct_warp_registration("_Frontend-Test-Stable-Direct-v0") - self.assertIn("isaaclab_experimental", str(exc.exception)) - def test_rejects_unknown_task(self): - with self.assertRaises(FrontendIncompatibleError): - _assert_direct_warp_registration("Frontend-Test-NotRegistered-v0") +def test_direct_guard_rejects_stable_rooted(): + with pytest.raises(FrontendIncompatibleError) as exc: + fe._assert_direct_warp_registration("_Frontend-Test-Stable-Direct-v0") + assert "isaaclab_experimental" in str(exc.value) -if __name__ == "__main__": - unittest.main() +def test_direct_guard_rejects_unknown_task(): + with pytest.raises(FrontendIncompatibleError): + fe._assert_direct_warp_registration("Frontend-Test-NotRegistered-v0") diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py index af513e675f03..668e758521c4 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py @@ -9,6 +9,15 @@ import gymnasium as gym +from isaaclab_experimental.envs.frontend import register_mdp_route + +# The stable Ant task borrows Humanoid MDP terms, so its warp twins live in the +# experimental humanoid package. +register_mdp_route( + "isaaclab_tasks.core.locomotion.ant", + "isaaclab_tasks_experimental.manager_based.classic.humanoid.mdp", +) + # Reuse agent configs from the stable task package. from isaaclab_tasks.core.locomotion.ant import agents diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py index 9f46551de255..2ab0e3555059 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py @@ -9,6 +9,11 @@ import gymnasium as gym +from isaaclab_experimental.envs.frontend import register_mdp_route + +# Warp twins for the stable cartpole MDP terms live in this package's ``mdp``. +register_mdp_route("isaaclab_tasks.core.cartpole", f"{__name__}.mdp") + # Reuse agent configs from the stable task package. from isaaclab_tasks.core.cartpole import agents diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py index 0bcd020bcd4c..6d2d958b759b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py @@ -9,6 +9,12 @@ import gymnasium as gym +from isaaclab_experimental.envs.frontend import register_mdp_route + +# Warp twins for the stable humanoid MDP terms live in this package's ``mdp``. +# The stable Ant task borrows Humanoid MDP terms, so this route serves it too. +register_mdp_route("isaaclab_tasks.core.locomotion.humanoid", f"{__name__}.mdp") + # Reuse agent configs from the stable task package. from isaaclab_tasks.core.locomotion.humanoid import agents diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py index 0857176a3fc7..9ddf530c28e0 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py @@ -4,3 +4,8 @@ # SPDX-License-Identifier: BSD-3-Clause """Velocity locomotion experimental task registrations (manager-based).""" + +from isaaclab_experimental.envs.frontend import register_mdp_route + +# Warp twins for the stable velocity MDP terms live in this package's ``mdp``. +register_mdp_route("isaaclab_tasks.core.velocity", f"{__name__}.mdp") diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py index fe34199f2321..8589501220cc 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py @@ -4,3 +4,8 @@ # SPDX-License-Identifier: BSD-3-Clause """Reach experimental task registrations (manager-based).""" + +from isaaclab_experimental.envs.frontend import register_mdp_route + +# Warp twins for the stable reach MDP terms live in this package's ``mdp``. +register_mdp_route("isaaclab_tasks.core.reach", f"{__name__}.mdp") From cd45040a1056aa2eb676cf0ad68d8ad73c20d6e4 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 15 Jul 2026 18:12:42 -0700 Subject: [PATCH 26/58] Replace duplicated warp task configs with stable-cfg reuse Every manager-based *-Warp-v0 task duplicated its stable environment configuration file-for-file, so each stable task rewrite had to be mirrored by hand (as happened with the Cartpole rewrite). ManagerBasedRLEnvWarp now adapts its cfg in __init__, so warp env construction accepts stable-derived cfgs directly. On top of that: - Drop the Cartpole, Humanoid, Ant, and Reach-Franka warp registrations and their duplicated cfgs; the stable ids run on warp via --frontend warp presets=newton_mjwarp. Their packages keep only the warp MDP twins and the route registration. - Rewrite the velocity *-Warp-v0 variants as thin subclasses of the stable flat cfgs that only disable the rigid-body material/mass randomization events, which have no warp twins yet. The duplicated velocity base cfg, per-robot cfgs, and unregistered rough cfgs are removed; the variants now select Newton via presets=newton_mjwarp like every other task. - Cover both paths in test_frontend_cfg_conversion.py: the stable ids with full twin coverage adapt cleanly, and every registered warp variant still adapts. - Update the warp environments documentation and the generated environment list accordingly. --- .../newton/warp-environments.rst | 42 ++- docs/source/overview/environments.rst | 20 -- .../changelog.d/warp-manager-bridge.rst | 9 + .../isaaclab_experimental/envs/frontend.py | 7 +- .../envs/manager_based_rl_env_warp.py | 8 + .../test/envs/test_frontend.py | 20 +- .../test/envs/test_frontend_cfg_conversion.py | 102 ++++-- .../changelog.d/warp-manager-bridge.minor.rst | 30 ++ .../changelog.d/warp-manager-bridge.rst | 8 - .../manager_based/classic/ant/__init__.py | 28 +- .../manager_based/classic/ant/ant_env_cfg.py | 197 ------------ .../classic/cartpole/__init__.py | 28 +- .../classic/cartpole/cartpole_env_cfg.py | 199 ------------ .../classic/humanoid/__init__.py | 29 +- .../classic/humanoid/humanoid_env_cfg.py | 232 -------------- .../locomotion/velocity/__init__.py | 29 +- .../velocity/config/a1/flat_env_cfg.py | 59 +--- .../velocity/config/a1/rough_env_cfg.py | 91 ------ .../velocity/config/anymal_b/flat_env_cfg.py | 63 ++-- .../velocity/config/anymal_b/rough_env_cfg.py | 34 -- .../velocity/config/anymal_c/flat_env_cfg.py | 51 +-- .../velocity/config/anymal_c/rough_env_cfg.py | 40 --- .../velocity/config/anymal_d/__init__.py | 25 -- .../velocity/config/anymal_d/flat_env_cfg.py | 47 ++- .../velocity/config/anymal_d/rough_env_cfg.py | 37 --- .../velocity/config/cassie/flat_env_cfg.py | 46 ++- .../velocity/config/cassie/rough_env_cfg.py | 96 ------ .../locomotion/velocity/config/g1/__init__.py | 22 -- .../velocity/config/g1/flat_env_cfg.py | 61 ++-- .../velocity/config/g1/rough_env_cfg.py | 178 ----------- .../velocity/config/go1/flat_env_cfg.py | 47 ++- .../velocity/config/go1/rough_env_cfg.py | 61 ---- .../velocity/config/go2/flat_env_cfg.py | 47 ++- .../velocity/config/go2/rough_env_cfg.py | 60 ---- .../locomotion/velocity/config/h1/__init__.py | 22 -- .../velocity/config/h1/flat_env_cfg.py | 45 +-- .../velocity/config/h1/rough_env_cfg.py | 131 -------- .../locomotion/velocity/velocity_env_cfg.py | 296 ------------------ .../manipulation/reach/__init__.py | 8 +- .../manipulation/reach/config/__init__.py | 4 - .../reach/config/franka/__init__.py | 41 --- .../reach/config/franka/joint_pos_env_cfg.py | 74 ----- .../reach/config/ur_10/__init__.py | 36 --- .../reach/config/ur_10/joint_pos_env_cfg.py | 74 ----- .../manipulation/reach/reach_env_cfg.py | 206 ------------ 45 files changed, 359 insertions(+), 2631 deletions(-) create mode 100644 source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst delete mode 100644 source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.rst delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/cartpole_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/joint_pos_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/joint_pos_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py diff --git a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst index 05c5bc848422..8f682ddcb047 100644 --- a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst +++ b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst @@ -37,33 +37,41 @@ Direct Warp Environments - ``Isaac-Reorient-Cube-Allegro-Direct-Warp-v0`` — Allegro hand cube reorient -Manager-Based Warp Environments -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Manager-Based Warp Execution (``--frontend warp``) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -**Classic** +Manager-based tasks do not need a parallel warp registration: the shared RL +CLI exposes ``--frontend {torch,warp}``, and ``--frontend warp`` adapts the +*stable* task configuration onto the warp runtime at build time (swapping each +MDP term for its warp twin). Select the Newton solver explicitly with +``presets=newton_mjwarp``. Stable tasks with full twin coverage: -- ``Isaac-Cartpole-Warp-v0`` -- ``Isaac-Ant-Warp-v0`` -- ``Isaac-Humanoid-Warp-v0`` +- ``Isaac-Cartpole`` +- ``Isaac-Ant`` +- ``Isaac-Humanoid`` +- ``Isaac-Reach-Franka`` +- ``Isaac-Reach-UR10`` -**Locomotion (Flat)** +A missing twin is a hard error listing the affected terms, so a partially +covered task fails at build time rather than silently changing behavior. + +**Locomotion (Flat) — registered warp variants** + +The velocity tasks still register ``*-Warp-v0`` variants because their stable +configurations contain randomization events without warp twins yet; the +variants reuse the stable flat configurations and only disable those events. +They also require ``presets=newton_mjwarp``. - ``Isaac-Velocity-Flat-AnymalB-Warp-v0`` - ``Isaac-Velocity-Flat-AnymalC-Warp-v0`` - ``Isaac-Velocity-Flat-AnymalD-Warp-v0`` - ``Isaac-Velocity-Flat-Cassie-Warp-v0`` - ``Isaac-Velocity-Flat-G1-Warp-v0`` -- ``Isaac-Velocity-Flat-G1-Warp-v1`` - ``Isaac-Velocity-Flat-H1-Warp-v0`` - ``Isaac-Velocity-Flat-UnitreeA1-Warp-v0`` - ``Isaac-Velocity-Flat-UnitreeGo1-Warp-v0`` - ``Isaac-Velocity-Flat-UnitreeGo2-Warp-v0`` -**Manipulation** - -- ``Isaac-Reach-Franka-Warp-v0`` -- ``Isaac-Reach-UR10-Warp-v0`` - Quick Start ~~~~~~~~~~~ @@ -74,9 +82,13 @@ Quick Start ./isaaclab.sh train --rl_library rsl_rl \ --task Isaac-Cartpole-Direct-Warp-v0 --num_envs 4096 - # Manager-based workflow + # Manager-based workflow: stable task on the warp runtime + ./isaaclab.sh train --rl_library rsl_rl \ + --task Isaac-Cartpole --frontend warp presets=newton_mjwarp --num_envs 4096 + + # Manager-based workflow: registered warp variant ./isaaclab.sh train --rl_library rsl_rl \ - --task Isaac-Velocity-Flat-AnymalC-Warp-v0 --num_envs 4096 + --task Isaac-Velocity-Flat-AnymalD-Warp-v0 presets=newton_mjwarp --num_envs 4096 All RL libraries with warp-compatible wrappers are supported: RSL-RL, RL Games, SKRL, and Stable-Baselines3. diff --git a/docs/source/overview/environments.rst b/docs/source/overview/environments.rst index b82b1bc83e06..8c821bf47bdb 100644 --- a/docs/source/overview/environments.rst +++ b/docs/source/overview/environments.rst @@ -880,11 +880,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Ant-Warp-v0 - - - - Manager Based - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) - - * - IsaacContrib-Assemble-Trocar-G129-Dex3 - IsaacContrib-Assemble-Trocar-G129-Dex3-Eval - Manager Based @@ -940,11 +935,6 @@ inferencing, including reading from an already trained checkpoint and disabling - **skrl** (PPO, BOX_BOX, BOX_DISCRETE, BOX_MULTIDISCRETE, DICT_BOX, DICT_DISCRETE, DICT_MULTIDISCRETE, DISCRETE_BOX, DISCRETE_DISCRETE, DISCRETE_MULTIDISCRETE, MULTIDISCRETE_BOX, MULTIDISCRETE_DISCRETE, MULTIDISCRETE_MULTIDISCRETE, TUPLE_BOX, TUPLE_DISCRETE, TUPLE_MULTIDISCRETE) - | **physics=** ``newton_kamino``, ``newton_mjwarp``, ``ovphysx``, ``physx`` | **presets=** ``box_box``, ``box_discrete``, ``box_multidiscrete``, ``dict_box``, ``dict_discrete``, ``dict_multidiscrete``, ``discrete_box``, ``discrete_discrete``, ``discrete_multidiscrete``, ``multidiscrete_box``, ``multidiscrete_discrete``, ``multidiscrete_multidiscrete``, ``tuple_box``, ``tuple_discrete``, ``tuple_multidiscrete`` - * - Isaac-Cartpole-Warp-v0 - - - - Manager Based - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) - - * - IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav-ROS-Inference - - Manager Based @@ -1060,11 +1050,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Humanoid-Warp-v0 - - - - Manager Based - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) - - * - Isaac-Lift-Cloth-Franka - - Manager Based @@ -1231,11 +1216,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Manager Based - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_mjwarp``, ``physx`` - * - Isaac-Reach-Franka-Warp-v0 - - Isaac-Reach-Franka-Warp-Play-v0 - - Manager Based - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - * - IsaacContrib-Reach-OpenArm - IsaacContrib-Reach-OpenArm-Play - Manager Based diff --git a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst index 4b3996529ba8..ccb21e9b301f 100644 --- a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst +++ b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst @@ -8,6 +8,15 @@ Added :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable` used by ``--frontend=warp`` to adapt stable cfgs onto the warp runtime. +Changed +^^^^^^^ + +* Changed :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` to adapt + its configuration in ``__init__`` via + :func:`~isaaclab_experimental.envs.frontend.adapt_cfg_for_warp`, so + registered warp task variants can derive from stable configurations instead + of duplicating them. + Fixed ^^^^^ diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index 780685198e7c..eefc9327476f 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -172,12 +172,11 @@ def build( return gym.make(task_id, cfg=env_cfg, **construct_kwargs) # Imported lazily so that ``--frontend=torch`` callers don't pay the - # ``isaaclab_experimental.envs`` import cost. Registered ``*-Warp-v0`` tasks - # already provide Warp-native configs; the frontend path adapts stable configs - # immediately before constructing the Warp environment. + # ``isaaclab_experimental.envs`` import cost. The env adapts the cfg itself + # in ``__init__`` (see :func:`adapt_cfg_for_warp`), so stable cfgs and + # warp-native cfgs take the same path. from isaaclab_experimental.envs import ManagerBasedRLEnvWarp - adapt_cfg_for_warp(env_cfg) return ManagerBasedRLEnvWarp(cfg=env_cfg, **construct_kwargs) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index 06ba2db190ae..a90489264a48 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -93,6 +93,14 @@ def __init__(self, cfg: ManagerBasedRLEnvCfg, render_mode: str | None = None, ** render_mode: The render mode for the environment. Defaults to None, which is similar to ``"human"``. """ + # Adapt the cfg for the warp managers (Newton physics check, SceneEntityCfg + # promotion, MDP twin swap). Idempotent: a warp-native cfg passes through + # unchanged, and a stable-derived cfg (``--frontend=warp`` or a registered + # warp task variant subclassing a stable cfg) is adapted in place. + from isaaclab_experimental.envs.frontend import adapt_cfg_for_warp + + adapt_cfg_for_warp(cfg) + # -- counter for curriculum self.common_step_counter = 0 diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py index 03e39287e679..ee374ed1f9ef 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend.py +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -168,28 +168,24 @@ def test_workflow_values(): # ====================================================================== -def test_warp_manager_build_adapts_cfg_before_construction(): +def test_warp_manager_build_constructs_warp_env_with_cfg(): + # Cfg adaptation happens inside ManagerBasedRLEnvWarp.__init__, so build() + # only selects the env class and forwards the cfg and construct kwargs. import isaaclab_experimental.envs as envs cfg = object.__new__(ManagerBasedRLEnvCfg) expected_env = object() - calls: list[tuple[str, Any]] = [] - - def fake_adapt(received_cfg: Any) -> None: - calls.append(("adapt", received_cfg)) + calls: list[tuple[Any, Any]] = [] def fake_env(*, cfg: Any, **kwargs: Any) -> Any: - calls.append(("construct", cfg)) + calls.append((cfg, kwargs)) return expected_env - with ( - patch.object(fe, "adapt_cfg_for_warp", fake_adapt), - patch.object(envs, "ManagerBasedRLEnvWarp", fake_env), - ): - env = build("warp", cfg, "Isaac-Test") + with patch.object(envs, "ManagerBasedRLEnvWarp", fake_env): + env = build("warp", cfg, "Isaac-Test", render_mode="rgb_array") assert env is expected_env - assert calls == [("adapt", cfg), ("construct", cfg)] + assert calls == [(cfg, {"render_mode": "rgb_array"})] # ====================================================================== diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index 909ca8dec08e..8ad8716bfd89 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -3,12 +3,21 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Coverage tests for manager-based Warp task configuration adaptation. +"""Coverage tests for warp task configuration adaptation. -For each registered manager-based Warp task, load its environment configuration, -force Newton physics, and run :func:`adapt_cfg_for_warp` — the same adaptation the -Warp environment runs in its ``__init__``. A dedicated stable Cartpole case also -guards the stable-to-experimental module routing used by ``--frontend warp``. +Two sweeps, both running :func:`adapt_cfg_for_warp` — the same adaptation +:class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` runs in its +``__init__``: + +* Every *stable* task id with declared warp twin coverage: resolve the + ``newton_mjwarp`` preset and adapt the stable cfg, exactly what + ``--frontend warp`` does. +* Every registered manager-based ``*-Warp-v0`` variant (velocity deltas whose + stable cfgs still contain terms without warp twins): load its env cfg, + resolve presets, and adapt. + +A dedicated stable Cartpole case also guards the stable-to-experimental +module routing used by ``--frontend warp``. """ from __future__ import annotations @@ -19,16 +28,32 @@ import isaaclab_tasks_experimental # noqa: F401 import pytest from isaaclab_experimental.envs.frontend import adapt_cfg_for_warp -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_newton.physics import NewtonCfg # Registering the task packages is the whole point — import for side effects. import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils.hydra import resolve_presets + +# Stable task ids whose MDP terms are fully twinned: ``--frontend warp`` adapts +# their cfgs directly, with no parallel ``*-Warp-v0`` registration. Extend this +# list when a new task family gains full twin coverage. +_STABLE_TASKS_WITH_WARP_COVERAGE = [ + "Isaac-Ant", + "Isaac-Cartpole", + "Isaac-Humanoid", + "Isaac-Reach-Franka", + "Isaac-Reach-Franka-Play", + "Isaac-Reach-UR10", + "Isaac-Reach-UR10-Play", +] # Manager-based warp tasks are exactly those whose entry point is the shared warp # env class; direct ``*-Direct-Warp-v0`` tasks register their own env class and # are not cfg-adapted, so they are excluded here. _MANAGER_WARP_ENTRY_POINT = "isaaclab_experimental.envs:ManagerBasedRLEnvWarp" +_WARP_ROOTS = ("isaaclab_experimental", "isaaclab_tasks_experimental") + def _manager_warp_tasks() -> list[tuple[str, str]]: """Return ``(task_id, env_cfg_entry_point)`` for every manager-based warp task.""" @@ -42,36 +67,50 @@ def _manager_warp_tasks() -> list[tuple[str, str]]: return sorted(tasks) -def _instantiate_cfg(cfg_entry_point: str): - """Import ``module:Class`` and instantiate the stable env cfg.""" +def _load_adapted_cfg(cfg_entry_point: str): + """Instantiate an env cfg, resolve the Newton preset, and adapt it for warp.""" module_path, class_name = cfg_entry_point.split(":") - return getattr(importlib.import_module(module_path), class_name)() + cfg = getattr(importlib.import_module(module_path), class_name)() + cfg = resolve_presets(cfg, selected=("newton_mjwarp",)) + assert isinstance(cfg.sim.physics, NewtonCfg), "task does not provide a newton_mjwarp physics preset" + # Raises FrontendIncompatibleError if any warp-managed term lacks a warp twin. + adapt_cfg_for_warp(cfg) + return cfg -_MANAGER_WARP_TASKS = _manager_warp_tasks() +def _cfg_entry_point(task_id: str) -> str: + cfg_entry = (gym.spec(task_id).kwargs or {}).get("env_cfg_entry_point") + assert isinstance(cfg_entry, str), f"{task_id}: no env_cfg_entry_point" + return cfg_entry -def _params(): - return [pytest.param(task_id, cfg_entry_point, id=task_id) for task_id, cfg_entry_point in _MANAGER_WARP_TASKS] +_MANAGER_WARP_TASKS = _manager_warp_tasks() def test_manager_warp_tasks_are_registered(): - """Sanity: the package actually registered manager-based warp tasks.""" - assert _MANAGER_WARP_TASKS, "no manager-based '*-Warp-v0' tasks registered" - - -_WARP_ROOTS = ("isaaclab_experimental", "isaaclab_tasks_experimental") - - -@pytest.mark.parametrize("task_id, cfg_entry_point", _params()) -def test_registered_cfg_adapts_to_warp(task_id: str, cfg_entry_point: str): - """Each registered manager-based Warp cfg adapts without a missing twin.""" - cfg = _instantiate_cfg(cfg_entry_point) - # The warp env requires Newton physics (normally via ``presets=newton_mjwarp``); - # set it directly so the test does not depend on Hydra preset resolution. - cfg.sim.physics = NewtonCfg(solver_cfg=MJWarpSolverCfg(), num_substeps=1) - # Raises FrontendIncompatibleError if any warp-managed term lacks a warp twin. - adapt_cfg_for_warp(cfg) + """Sanity: every expected velocity warp task variant is registered.""" + expected = { + f"Isaac-Velocity-Flat-{robot}-Warp{suffix}-v0" + for robot in ("AnymalB", "AnymalC", "AnymalD", "Cassie", "G1", "H1", "UnitreeA1", "UnitreeGo1", "UnitreeGo2") + for suffix in ("", "-Play") + } + registered = {task_id for task_id, _ in _MANAGER_WARP_TASKS} + assert registered == expected, f"missing: {sorted(expected - registered)}; extra: {sorted(registered - expected)}" + + +@pytest.mark.parametrize("task_id", _STABLE_TASKS_WITH_WARP_COVERAGE, ids=_STABLE_TASKS_WITH_WARP_COVERAGE) +def test_stable_task_cfg_adapts_to_warp(task_id: str): + """Each covered stable task adapts without a missing twin (the --frontend warp path).""" + _load_adapted_cfg(_cfg_entry_point(task_id)) + + +@pytest.mark.parametrize( + "task_id, cfg_entry_point", + [pytest.param(task_id, cfg_entry, id=task_id) for task_id, cfg_entry in _MANAGER_WARP_TASKS], +) +def test_registered_warp_variant_cfg_adapts(task_id: str, cfg_entry_point: str): + """Each registered ``*-Warp-v0`` variant cfg adapts without a missing twin.""" + cfg = _load_adapted_cfg(cfg_entry_point) # Action terms carry a ``class_type`` (not a ``func``) and live on a base that # is not a ManagerTermBaseCfg; guard that the adapter still swaps them to the @@ -91,12 +130,7 @@ def test_stable_cartpole_cfg_adapts_to_current_warp_module_layout(): """The stable Cartpole cfg resolves task-specific twins in the current package layout.""" from isaaclab_experimental.managers.action_manager import ActionTerm - from isaaclab_tasks.core.cartpole.cartpole_manager_env_cfg import CartpoleEnvCfg - - cfg = CartpoleEnvCfg() - cfg.sim.physics = NewtonCfg(solver_cfg=MJWarpSolverCfg(), num_substeps=1) - - adapt_cfg_for_warp(cfg) + cfg = _load_adapted_cfg(_cfg_entry_point("Isaac-Cartpole")) assert cfg.rewards.pole_pos.func.__module__.startswith( "isaaclab_tasks_experimental.manager_based.classic.cartpole.mdp" diff --git a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst new file mode 100644 index 000000000000..dfb8c6270757 --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst @@ -0,0 +1,30 @@ +Changed +^^^^^^^ + +* **Breaking:** Changed the manager-based velocity ``*-Warp-v0`` task variants + to reuse the stable flat configurations, disabling only the randomization + events that have no warp twins yet. The variants now require selecting the + Newton solver on the CLI via ``presets=newton_mjwarp`` instead of + hard-coding it in the configuration. + +Removed +^^^^^^^ + +* **Breaking:** Removed the duplicated manager-based warp task registrations + ``Isaac-Cartpole-Warp-v0``, ``Isaac-Humanoid-Warp-v0``, ``Isaac-Ant-Warp-v0``, + ``Isaac-Reach-Franka-Warp-v0``, and ``Isaac-Reach-Franka-Warp-Play-v0`` + together with their duplicated environment configurations. Run the stable + task ids with ``--frontend warp`` and ``presets=newton_mjwarp`` instead, + e.g. ``--task Isaac-Cartpole --frontend warp presets=newton_mjwarp``. +* Removed the unregistered rough velocity warp configurations; rough-terrain + warp tasks remain unsupported until :class:`~isaaclab.terrains.TerrainImporter` + gains Warp APIs. + +Fixed +^^^^^ + +* Fixed the direct Warp Cartpole task to match the stable task's observations, + reset ranges, termination condition, reward scaling, and scene configuration. +* Fixed the Warp Cartpole ``survival_success_rate`` twin to report the + ``Metrics/success_rate`` value on-device instead of silently dropping the + metric. diff --git a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.rst deleted file mode 100644 index 4ac20a99187a..000000000000 --- a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.rst +++ /dev/null @@ -1,8 +0,0 @@ -Fixed -^^^^^ - -* Fixed the direct Warp Cartpole task to match the stable task's observations, - reset ranges, termination condition, reward scaling, and scene configuration. -* Fixed the Warp Cartpole ``survival_success_rate`` twin to report the - ``Metrics/success_rate`` value on-device instead of silently dropping the - metric. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py index 668e758521c4..0263ee66c216 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py @@ -3,11 +3,11 @@ # # SPDX-License-Identifier: BSD-3-Clause -""" -Ant locomotion environment (experimental manager-based entry point). -""" +"""Warp MDP route for the stable Ant task. -import gymnasium as gym +There is no separate task registration: run the stable ``Isaac-Ant`` task with +``--frontend warp`` and ``presets=newton_mjwarp``. +""" from isaaclab_experimental.envs.frontend import register_mdp_route @@ -17,23 +17,3 @@ "isaaclab_tasks.core.locomotion.ant", "isaaclab_tasks_experimental.manager_based.classic.humanoid.mdp", ) - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.locomotion.ant import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Ant-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.ant_env_cfg:AntEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AntPPORunnerCfg", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_manager_ppo_cfg.yaml", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_manager_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{agents.__name__}:sb3_manager_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py deleted file mode 100644 index 66a9131011f3..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py +++ /dev/null @@ -1,197 +0,0 @@ -# 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 - -# Ant reuses humanoid's experimental MDP (mirrors stable pattern). -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.classic.humanoid.mdp as mdp - -## -# Pre-defined configs -## -from isaaclab_assets.robots.ant import ANT_CFG # isort: skip - - -@configclass -class MySceneCfg(InteractiveSceneCfg): - """Configuration for the terrain scene with an ant robot.""" - - # terrain - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="plane", - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg( - friction_combine_mode="average", - restitution_combine_mode="average", - static_friction=1.0, - dynamic_friction=1.0, - restitution=0.0, - ), - debug_vis=False, - ) - - # robot - robot = ANT_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - - # lights - light = AssetBaseCfg( - prim_path="/World/light", - spawn=sim_utils.DistantLightCfg(color=(0.75, 0.75, 0.75), intensity=3000.0), - ) - - -## -# MDP settings -## - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - joint_effort = mdp.JointEffortActionCfg(asset_name="robot", joint_names=[".*"], scale=7.5) - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - @configclass - class PolicyCfg(ObsGroup): - """Observations for the policy.""" - - base_height = ObsTerm(func=mdp.base_pos_z) - base_lin_vel = ObsTerm(func=mdp.base_lin_vel) - base_ang_vel = ObsTerm(func=mdp.base_ang_vel) - base_yaw_roll = ObsTerm(func=mdp.base_yaw_roll) - base_angle_to_target = ObsTerm(func=mdp.base_angle_to_target, params={"target_pos": (1000.0, 0.0, 0.0)}) - base_up_proj = ObsTerm(func=mdp.base_up_proj) - base_heading_proj = ObsTerm(func=mdp.base_heading_proj, params={"target_pos": (1000.0, 0.0, 0.0)}) - joint_pos_norm = ObsTerm(func=mdp.joint_pos_limit_normalized) - joint_vel_rel = ObsTerm(func=mdp.joint_vel_rel, scale=0.2) - actions = ObsTerm(func=mdp.last_action) - - def __post_init__(self): - self.enable_corruption = False - self.concatenate_terms = True - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - reset_base = EventTerm( - func=mdp.reset_root_state_uniform, - mode="reset", - params={"pose_range": {}, "velocity_range": {}}, - ) - - reset_robot_joints = EventTerm( - func=mdp.reset_joints_by_offset, - mode="reset", - params={ - "position_range": (-0.2, 0.2), - "velocity_range": (-0.1, 0.1), - }, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # (1) Reward for moving forward - progress = RewTerm(func=mdp.progress_reward, weight=1.0, params={"target_pos": (1000.0, 0.0, 0.0)}) - # (2) Stay alive bonus - alive = RewTerm(func=mdp.is_alive, weight=0.5) - # (3) Reward for non-upright posture - upright = RewTerm(func=mdp.upright_posture_bonus, weight=0.1, params={"threshold": 0.93}) - # (4) Reward for moving in the right direction - move_to_target = RewTerm( - func=mdp.move_to_target_bonus, weight=0.5, params={"threshold": 0.8, "target_pos": (1000.0, 0.0, 0.0)} - ) - # (5) Penalty for large action commands - action_l2 = RewTerm(func=mdp.action_l2, weight=-0.005) - # (6) Penalty for energy consumption - energy = RewTerm(func=mdp.power_consumption, weight=-0.05, params={"gear_ratio": {".*": 15.0}}) - # (7) Penalty for reaching close to joint limits - joint_pos_limits = RewTerm( - func=mdp.joint_pos_limits_penalty_ratio, weight=-0.1, params={"threshold": 0.99, "gear_ratio": {".*": 15.0}} - ) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - # (1) Terminate if the episode length is exceeded - time_out = DoneTerm(func=mdp.time_out, time_out=True) - # (2) Terminate if the robot falls - torso_height = DoneTerm(func=mdp.root_height_below_minimum, params={"minimum_height": 0.31}) - - -## -# Environment configuration -## - - -@configclass -class AntEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the MuJoCo-style Ant walking environment.""" - - # Simulation settings - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=38, - nconmax=15, - ls_iterations=10, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - # Scene settings - scene: MySceneCfg = MySceneCfg(num_envs=4096, env_spacing=5.0, clone_in_fabric=True) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() - - def __post_init__(self): - """Post initialization.""" - # general settings - self.decimation = 2 - self.episode_length_s = 16.0 - # simulation settings - self.sim.dt = 1 / 120.0 - self.sim.render_interval = self.decimation - # default friction material - self.sim.physics_material.static_friction = 1.0 - self.sim.physics_material.dynamic_friction = 1.0 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py index 2ab0e3555059..e527f251fb5c 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py @@ -3,33 +3,13 @@ # # SPDX-License-Identifier: BSD-3-Clause -""" -Cartpole balancing environment (experimental manager-based entry point). -""" +"""Warp MDP twins for the stable Cartpole task. -import gymnasium as gym +There is no separate task registration: run the stable ``Isaac-Cartpole`` task +with ``--frontend warp`` and ``presets=newton_mjwarp``. +""" from isaaclab_experimental.envs.frontend import register_mdp_route # Warp twins for the stable cartpole MDP terms live in this package's ``mdp``. register_mdp_route("isaaclab_tasks.core.cartpole", f"{__name__}.mdp") - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.cartpole import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Cartpole-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.cartpole_env_cfg:CartpoleEnvCfg", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_manager_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_manager_ppo_cfg:CartpolePPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_manager_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{agents.__name__}:sb3_manager_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/cartpole_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/cartpole_env_cfg.py deleted file mode 100644 index ae187a593a16..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/cartpole_env_cfg.py +++ /dev/null @@ -1,199 +0,0 @@ -# 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 - -import math - -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg, AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.classic.cartpole.mdp as mdp - -## -# Pre-defined configs -## -from isaaclab_assets.robots.cartpole import CARTPOLE_CFG # isort:skip - - -## -# Scene definition -## - - -@configclass -class CartpoleSceneCfg(InteractiveSceneCfg): - """Configuration for a cart-pole scene.""" - - # ground plane - # ground = AssetBaseCfg( - # prim_path="/World/ground", - # spawn=sim_utils.GroundPlaneCfg(size=(100.0, 100.0)), - # ) - - # cartpole - robot: ArticulationCfg = CARTPOLE_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - - # lights - dome_light = AssetBaseCfg( - prim_path="/World/DomeLight", - spawn=sim_utils.DomeLightCfg(color=(0.9, 0.9, 0.9), intensity=500.0), - ) - - -## -# MDP settings -## - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - joint_effort = mdp.JointEffortActionCfg(asset_name="robot", joint_names=["slider_to_cart"], scale=100.0) - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - @configclass - class PolicyCfg(ObsGroup): - """Observations for policy group.""" - - # observation terms (order preserved) - joint_pos_rel = ObsTerm(func=mdp.joint_pos_rel) - joint_vel_rel = ObsTerm(func=mdp.joint_vel_rel) - - def __post_init__(self) -> None: - self.enable_corruption = False - self.concatenate_terms = True - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - # reset - reset_cart_position = EventTerm( - func=mdp.reset_joints_by_offset, - mode="reset", - params={ - "asset_cfg": SceneEntityCfg("robot", joint_names=["slider_to_cart"]), - "position_range": (-1.0, 1.0), - "velocity_range": (-0.5, 0.5), - }, - ) - - reset_pole_position = EventTerm( - func=mdp.reset_joints_by_offset, - mode="reset", - params={ - "asset_cfg": SceneEntityCfg("robot", joint_names=["cart_to_pole"]), - "position_range": (-0.25 * math.pi, 0.25 * math.pi), - "velocity_range": (-0.25 * math.pi, 0.25 * math.pi), - }, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # (1) Constant running reward - alive = RewTerm(func=mdp.is_alive, weight=1.0) - # (2) Failure penalty - terminating = RewTerm(func=mdp.is_terminated, weight=-2.0) - # (3) Primary task: keep pole upright - pole_pos = RewTerm( - func=mdp.joint_pos_target_l2, - weight=-1.0, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["cart_to_pole"]), "target": 0.0}, - ) - # (4) Shaping tasks: lower cart velocity - cart_vel = RewTerm( - func=mdp.joint_vel_l1, - weight=-0.01, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["slider_to_cart"])}, - ) - # (5) Shaping tasks: lower pole angular velocity - pole_vel = RewTerm( - func=mdp.joint_vel_l1, - weight=-0.005, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["cart_to_pole"])}, - ) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - # (1) Time out - time_out = DoneTerm(func=mdp.time_out, time_out=True) - # (2) Cart out of bounds - cart_out_of_bounds = DoneTerm( - func=mdp.joint_pos_out_of_manual_limit, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["slider_to_cart"]), "bounds": (-3.0, 3.0)}, - ) - - -## -# Environment configuration -## - - -@configclass -class CartpoleEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the cartpole environment.""" - - # Scene settings - scene: CartpoleSceneCfg = CartpoleSceneCfg(num_envs=4096, env_spacing=4.0, clone_in_fabric=True) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - events: EventCfg = EventCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - # Simulation settings - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=5, - nconmax=3, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - use_cuda_graph=True, - ) - ) - - # Post initialization - def __post_init__(self) -> None: - """Post initialization.""" - # general settings - self.decimation = 2 - self.episode_length_s = 5 - # viewer settings - self.viewer.eye = (8.0, 0.0, 5.0) - # simulation settings - self.sim.dt = 1 / 120 - self.sim.render_interval = self.decimation diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py index 6d2d958b759b..0d31a47ece94 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py @@ -3,34 +3,13 @@ # # SPDX-License-Identifier: BSD-3-Clause -""" -Humanoid locomotion environment (experimental manager-based entry point). -""" +"""Warp MDP twins for the stable Humanoid task. -import gymnasium as gym +There is no separate task registration: run the stable ``Isaac-Humanoid`` task +with ``--frontend warp`` and ``presets=newton_mjwarp``. +""" from isaaclab_experimental.envs.frontend import register_mdp_route # Warp twins for the stable humanoid MDP terms live in this package's ``mdp``. -# The stable Ant task borrows Humanoid MDP terms, so this route serves it too. register_mdp_route("isaaclab_tasks.core.locomotion.humanoid", f"{__name__}.mdp") - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.locomotion.humanoid import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Humanoid-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.humanoid_env_cfg:HumanoidEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:HumanoidPPORunnerCfg", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_manager_ppo_cfg.yaml", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_manager_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{agents.__name__}:sb3_manager_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py deleted file mode 100644 index 9de689ffb98e..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py +++ /dev/null @@ -1,232 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.classic.humanoid.mdp as mdp - -from isaaclab_assets.robots.humanoid import HUMANOID_CFG # isort:skip - - -## -# Scene definition -## - - -@configclass -class MySceneCfg(InteractiveSceneCfg): - """Configuration for the terrain scene with a humanoid robot.""" - - # terrain - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="plane", - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg(static_friction=1.0, dynamic_friction=1.0, restitution=0.0), - debug_vis=False, - ) - - # robot - robot = HUMANOID_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - - # lights - light = AssetBaseCfg( - prim_path="/World/light", - spawn=sim_utils.DistantLightCfg(color=(0.75, 0.75, 0.75), intensity=3000.0), - ) - - -## -# MDP settings -## - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - joint_effort = mdp.JointEffortActionCfg( - asset_name="robot", - joint_names=[".*"], - scale={ - ".*_waist.*": 67.5, - ".*_upper_arm.*": 67.5, - "pelvis": 67.5, - ".*_lower_arm": 45.0, - ".*_thigh:0": 45.0, - ".*_thigh:1": 135.0, - ".*_thigh:2": 45.0, - ".*_shin": 90.0, - ".*_foot.*": 22.5, - }, - ) - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - @configclass - class PolicyCfg(ObsGroup): - """Observations for the policy.""" - - base_height = ObsTerm(func=mdp.base_pos_z) - base_lin_vel = ObsTerm(func=mdp.base_lin_vel) - base_ang_vel = ObsTerm(func=mdp.base_ang_vel, scale=0.25) - base_yaw_roll = ObsTerm(func=mdp.base_yaw_roll) - base_angle_to_target = ObsTerm(func=mdp.base_angle_to_target, params={"target_pos": (1000.0, 0.0, 0.0)}) - base_up_proj = ObsTerm(func=mdp.base_up_proj) - base_heading_proj = ObsTerm(func=mdp.base_heading_proj, params={"target_pos": (1000.0, 0.0, 0.0)}) - joint_pos_norm = ObsTerm(func=mdp.joint_pos_limit_normalized) - joint_vel_rel = ObsTerm(func=mdp.joint_vel_rel, scale=0.1) - actions = ObsTerm(func=mdp.last_action) - - def __post_init__(self): - self.enable_corruption = False - self.concatenate_terms = True - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - reset_base = EventTerm( - func=mdp.reset_root_state_uniform, - mode="reset", - params={"pose_range": {}, "velocity_range": {}}, - ) - - reset_robot_joints = EventTerm( - func=mdp.reset_joints_by_offset, - mode="reset", - params={ - "position_range": (-0.2, 0.2), - "velocity_range": (-0.1, 0.1), - }, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # (1) Reward for moving forward - progress = RewTerm(func=mdp.progress_reward, weight=1.0, params={"target_pos": (1000.0, 0.0, 0.0)}) - # (2) Stay alive bonus - alive = RewTerm(func=mdp.is_alive, weight=2.0) - # (3) Reward for non-upright posture - upright = RewTerm(func=mdp.upright_posture_bonus, weight=0.1, params={"threshold": 0.93}) - # (4) Reward for moving in the right direction - move_to_target = RewTerm( - func=mdp.move_to_target_bonus, weight=0.5, params={"threshold": 0.8, "target_pos": (1000.0, 0.0, 0.0)} - ) - # (5) Penalty for large action commands - action_l2 = RewTerm(func=mdp.action_l2, weight=-0.01) - # (6) Penalty for energy consumption - energy = RewTerm( - func=mdp.power_consumption, - weight=-0.005, - params={ - "gear_ratio": { - ".*_waist.*": 67.5, - ".*_upper_arm.*": 67.5, - "pelvis": 67.5, - ".*_lower_arm": 45.0, - ".*_thigh:0": 45.0, - ".*_thigh:1": 135.0, - ".*_thigh:2": 45.0, - ".*_shin": 90.0, - ".*_foot.*": 22.5, - } - }, - ) - # (7) Penalty for reaching close to joint limits - joint_pos_limits = RewTerm( - func=mdp.joint_pos_limits_penalty_ratio, - weight=-0.25, - params={ - "threshold": 0.98, - "gear_ratio": { - ".*_waist.*": 67.5, - ".*_upper_arm.*": 67.5, - "pelvis": 67.5, - ".*_lower_arm": 45.0, - ".*_thigh:0": 45.0, - ".*_thigh:1": 135.0, - ".*_thigh:2": 45.0, - ".*_shin": 90.0, - ".*_foot.*": 22.5, - }, - }, - ) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - # (1) Terminate if the episode length is exceeded - time_out = DoneTerm(func=mdp.time_out, time_out=True) - # (2) Terminate if the robot falls - torso_height = DoneTerm(func=mdp.root_height_below_minimum, params={"minimum_height": 0.8}) - - -@configclass -class HumanoidEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the MuJoCo-style Humanoid walking environment.""" - - # Scene settings - scene: MySceneCfg = MySceneCfg(num_envs=4096, env_spacing=5.0, clone_in_fabric=True) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() - - def __post_init__(self): - """Post initialization.""" - # general settings - self.decimation = 2 - self.episode_length_s = 16.0 - # simulation settings - self.sim: SimulationCfg = SimulationCfg( - dt=1 / 120.0, - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=80, - nconmax=25, - ls_iterations=15, - cone="pyramidal", - update_data_interval=2, - impratio=1, - integrator="implicitfast", - ), - num_substeps=2, - debug_mode=False, - ), - ) - # self.sim.dt = 1 / 120.0 - self.sim.render_interval = self.decimation - # default friction material - self.sim.physics_material.static_friction = 1.0 - self.sim.physics_material.dynamic_friction = 1.0 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py index 9ddf530c28e0..436c1da089b6 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py @@ -3,9 +3,36 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Velocity locomotion experimental task registrations (manager-based).""" +"""Velocity locomotion experimental task registrations (manager-based). + +The per-robot ``*-Warp-v0`` variants reuse the stable flat cfgs and only +disable the randomization events that have no warp twins yet (see +:func:`disable_unsupported_randomization_events`). Once those twins exist, +the variants can be dropped in favor of ``--frontend warp`` on the stable +task ids. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING from isaaclab_experimental.envs.frontend import register_mdp_route +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnvCfg + # Warp twins for the stable velocity MDP terms live in this package's ``mdp``. register_mdp_route("isaaclab_tasks.core.velocity", f"{__name__}.mdp") + + +def disable_unsupported_randomization_events(cfg: ManagerBasedRLEnvCfg) -> None: + """Disable stable randomization events that have no warp twins yet. + + The warp event manager invokes term functions with a Warp env mask, and no + warp twins exist yet for the rigid-body material/mass randomization events. + A stable term on a warp manager is a hard error at adaptation time, so the + warp task variants disable these terms until twins are available. + """ + for name in ("physics_material", "add_base_mass"): + if getattr(cfg.events, name, None) is not None: + setattr(cfg.events, name, None) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py index f213286d0cfb..d58290799424 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py @@ -3,58 +3,33 @@ # # SPDX-License-Identifier: BSD-3-Clause -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +"""Warp variants of the stable Unitree A1 flat velocity task configuration.""" -from isaaclab.sim import SimulationCfg from isaaclab.utils.configclass import configclass -from .rough_env_cfg import UnitreeA1RoughEnvCfg +from isaaclab_tasks.contrib.velocity.config.a1.flat_env_cfg import UnitreeA1FlatEnvCfg as _StableUnitreeA1FlatEnvCfg +from isaaclab_tasks.contrib.velocity.config.a1.flat_env_cfg import ( + UnitreeA1FlatEnvCfg_PLAY as _StableUnitreeA1FlatEnvCfg_PLAY, +) + +from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( + disable_unsupported_randomization_events, +) @configclass -class UnitreeA1FlatEnvCfg(UnitreeA1RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=60, - nconmax=30, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) +class UnitreeA1FlatEnvCfg(_StableUnitreeA1FlatEnvCfg): + """Stable flat cfg with the randomization events lacking warp twins disabled.""" def __post_init__(self): - # post init of parent super().__post_init__() + disable_unsupported_randomization_events(self) - # override rewards - self.rewards.flat_orientation_l2.weight = -2.5 - self.rewards.feet_air_time.weight = 0.25 - - # change terrain to flat - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - # no height scan - # self.scene.height_scanner = None - # self.observations.policy.height_scan = None - # no terrain curriculum - self.curriculum.terrain_levels = None +@configclass +class UnitreeA1FlatEnvCfg_PLAY(_StableUnitreeA1FlatEnvCfg_PLAY): + """Play variant of :class:`UnitreeA1FlatEnvCfg` for the warp runtime.""" -class UnitreeA1FlatEnvCfg_PLAY(UnitreeA1FlatEnvCfg): - def __post_init__(self) -> None: - # post init of parent + def __post_init__(self): super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing event - self.events.base_external_force_torque = None - self.events.push_robot = None + disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/rough_env_cfg.py deleted file mode 100644 index ee0f236043a2..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/rough_env_cfg.py +++ /dev/null @@ -1,91 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm - -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import ( - LocomotionVelocityRoughEnvCfg, - TerminationsCfg, -) - -from isaaclab_assets.robots.unitree import UNITREE_A1_CFG # isort: skip - - -class TerminationsCfg_A1(TerminationsCfg): - base_too_low = DoneTerm(func=mdp.root_height_below_minimum, params={"minimum_height": 0.2}) - - -@configclass -class UnitreeA1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - terminations: TerminationsCfg_A1 = TerminationsCfg_A1() - - def __post_init__(self): - # post init of parent - super().__post_init__() - - self.scene.robot = UNITREE_A1_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01 - - # reduce action scale - self.actions.joint_pos.scale = 0.25 - - # event - self.events.push_robot = None - # TODO: TEMPORARILY DISABLED - adding this causes NaNs in the simulation - # self.events.add_base_mass.params["mass_distribution_params"] = (-1.0, 3.0) - # self.events.add_base_mass.params["asset_cfg"].body_names = "trunk" - self.events.base_external_force_torque.params["asset_cfg"].body_names = "trunk" - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - - # rewards - self.rewards.feet_air_time.params["sensor_cfg"].body_names = ".*_foot" - self.rewards.feet_air_time.weight = 0.01 - self.rewards.undesired_contacts.params["sensor_cfg"].body_names = ".*thigh" - self.rewards.dof_torques_l2.weight = -0.0002 - self.rewards.track_lin_vel_xy_exp.weight = 1.5 - self.rewards.track_ang_vel_z_exp.weight = 0.75 - self.rewards.dof_acc_l2.weight = -2.5e-7 - self.terminations.base_contact.params["sensor_cfg"].body_names = "trunk" - - -@configclass -class UnitreeA1RoughEnvCfg_PLAY(UnitreeA1RoughEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # spawn the robot randomly in the grid (instead of their terrain levels) - self.scene.terrain.max_init_terrain_level = None - # reduce the number of terrains to save memory - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing event - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py index a5c825bf2e98..900d87e64227 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py @@ -3,58 +3,33 @@ # # SPDX-License-Identifier: BSD-3-Clause -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +"""Warp variants of the stable AnymalB flat velocity task configuration.""" -from isaaclab.sim import SimulationCfg from isaaclab.utils.configclass import configclass -from .rough_env_cfg import AnymalBRoughEnvCfg +from isaaclab_tasks.contrib.velocity.config.anymal_b.flat_env_cfg import AnymalBFlatEnvCfg as _StableAnymalBFlatEnvCfg +from isaaclab_tasks.contrib.velocity.config.anymal_b.flat_env_cfg import ( + AnymalBFlatEnvCfg_PLAY as _StableAnymalBFlatEnvCfg_PLAY, +) + +from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( + disable_unsupported_randomization_events, +) @configclass -class AnymalBFlatEnvCfg(AnymalBRoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=75, - nconmax=15, - cone="elliptic", - impratio=100, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) +class AnymalBFlatEnvCfg(_StableAnymalBFlatEnvCfg): + """Stable flat cfg with the randomization events lacking warp twins disabled.""" def __post_init__(self): - # post init of parent super().__post_init__() + disable_unsupported_randomization_events(self) - # override rewards - self.rewards.flat_orientation_l2.weight = -5.0 - self.rewards.dof_torques_l2.weight = -2.5e-5 - self.rewards.feet_air_time.weight = 0.5 - # change terrain to flat - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - # no height scan - self.scene.height_scanner = None - self.observations.policy.height_scan = None - # no terrain curriculum - self.curriculum.terrain_levels = None - - -class AnymalBFlatEnvCfg_PLAY(AnymalBFlatEnvCfg): - def __post_init__(self) -> None: - # post init of parent - super().__post_init__() - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing event - self.events.base_external_force_torque = None - self.events.push_robot = None +@configclass +class AnymalBFlatEnvCfg_PLAY(_StableAnymalBFlatEnvCfg_PLAY): + """Play variant of :class:`AnymalBFlatEnvCfg` for the warp runtime.""" + + def __post_init__(self): + super().__post_init__() + disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py deleted file mode 100644 index 3a7d80af470f..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py +++ /dev/null @@ -1,34 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -from isaaclab_assets import ANYMAL_B_CFG # isort: skip - - -@configclass -class AnymalBRoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.robot = ANYMAL_B_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.manager_call_max_mode = {"Scene": 1} - - -@configclass -class AnymalBRoughEnvCfg_PLAY(AnymalBRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py index 3cc348a97675..edb8ffd7a90e 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py @@ -3,50 +3,33 @@ # # SPDX-License-Identifier: BSD-3-Clause -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +"""Warp variants of the stable AnymalC flat velocity task configuration.""" -from isaaclab.sim import SimulationCfg from isaaclab.utils.configclass import configclass -from .rough_env_cfg import AnymalCRoughEnvCfg +from isaaclab_tasks.contrib.velocity.config.anymal_c.flat_env_cfg import AnymalCFlatEnvCfg as _StableAnymalCFlatEnvCfg +from isaaclab_tasks.contrib.velocity.config.anymal_c.flat_env_cfg import ( + AnymalCFlatEnvCfg_PLAY as _StableAnymalCFlatEnvCfg_PLAY, +) + +from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( + disable_unsupported_randomization_events, +) @configclass -class AnymalCFlatEnvCfg(AnymalCRoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=120, - nconmax=15, - cone="elliptic", - impratio=100, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) +class AnymalCFlatEnvCfg(_StableAnymalCFlatEnvCfg): + """Stable flat cfg with the randomization events lacking warp twins disabled.""" def __post_init__(self): - # post init of parent super().__post_init__() + disable_unsupported_randomization_events(self) - # override rewards - self.rewards.flat_orientation_l2.weight = -5.0 - self.rewards.dof_torques_l2.weight = -2.5e-5 - self.rewards.feet_air_time.weight = 0.5 - # change terrain to flat - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - # no terrain curriculum - self.curriculum.terrain_levels = None +@configclass +class AnymalCFlatEnvCfg_PLAY(_StableAnymalCFlatEnvCfg_PLAY): + """Play variant of :class:`AnymalCFlatEnvCfg` for the warp runtime.""" -class AnymalCFlatEnvCfg_PLAY(AnymalCFlatEnvCfg): - def __post_init__(self) -> None: + def __post_init__(self): super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None + disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py deleted file mode 100644 index 1d2f2676c64a..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py +++ /dev/null @@ -1,40 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -## -# Pre-defined configs -## -from isaaclab_assets.robots.anymal import ANYMAL_C_CFG # isort: skip - - -@configclass -class AnymalCRoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - # switch robot to anymal-c - self.scene.robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.scene.robot.actuators["legs"].armature = 0.01 - self.manager_call_max_mode = {"Scene": 1} - - -@configclass -class AnymalCRoughEnvCfg_PLAY(AnymalCRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/__init__.py index 3722de03ed4e..e54ed0b4d6ca 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/__init__.py @@ -12,31 +12,6 @@ # Register Gym environments. ## -# Rough env disabled: requires isaaclab_physx which is not yet available on dev/newton. -# The package exists on upstream/develop (commit 308400f1d35) but has not been merged. -# Re-enable once dev/newton picks up isaaclab_physx. -# gym.register( -# id="Isaac-Velocity-Rough-AnymalD-Warp-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:AnymalDRoughEnvCfg", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDRoughPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", -# }, -# ) - -# gym.register( -# id="Isaac-Velocity-Rough-AnymalD-Warp-Play-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:AnymalDRoughEnvCfg_PLAY", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDRoughPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", -# }, -# ) - gym.register( id="Isaac-Velocity-Flat-AnymalD-Warp-v0", entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py index 702443e815e2..ca010dffb13c 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py @@ -3,44 +3,33 @@ # # SPDX-License-Identifier: BSD-3-Clause -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +"""Warp variants of the stable AnymalD flat velocity task configuration.""" -from isaaclab.sim import SimulationCfg from isaaclab.utils.configclass import configclass -from .rough_env_cfg import AnymalDRoughEnvCfg +from isaaclab_tasks.core.velocity.config.anymal_d.flat_env_cfg import AnymalDFlatEnvCfg as _StableAnymalDFlatEnvCfg +from isaaclab_tasks.core.velocity.config.anymal_d.flat_env_cfg import ( + AnymalDFlatEnvCfg_PLAY as _StableAnymalDFlatEnvCfg_PLAY, +) + +from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( + disable_unsupported_randomization_events, +) @configclass -class AnymalDFlatEnvCfg(AnymalDRoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=60, - nconmax=25, - cone="elliptic", - impratio=100.0, - ), - num_substeps=1, - debug_mode=False, - ) - ) +class AnymalDFlatEnvCfg(_StableAnymalDFlatEnvCfg): + """Stable flat cfg with the randomization events lacking warp twins disabled.""" def __post_init__(self): super().__post_init__() - self.rewards.flat_orientation_l2.weight = -5.0 - self.rewards.dof_torques_l2.weight = -2.5e-5 - self.rewards.feet_air_time.weight = 0.5 - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None + disable_unsupported_randomization_events(self) + +@configclass +class AnymalDFlatEnvCfg_PLAY(_StableAnymalDFlatEnvCfg_PLAY): + """Play variant of :class:`AnymalDFlatEnvCfg` for the warp runtime.""" -class AnymalDFlatEnvCfg_PLAY(AnymalDFlatEnvCfg): - def __post_init__(self) -> None: + def __post_init__(self): super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None + disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py deleted file mode 100644 index 71f49edd2e68..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py +++ /dev/null @@ -1,37 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -## -# Pre-defined configs -## -from isaaclab_assets.robots.anymal import ANYMAL_D_CFG # isort: skip - - -@configclass -class AnymalDRoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.robot = ANYMAL_D_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.manager_call_max_mode = {"Scene": 1} - - -@configclass -class AnymalDRoughEnvCfg_PLAY(AnymalDRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/flat_env_cfg.py index 818c7aa30a1e..4b42afa8a3a5 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/flat_env_cfg.py @@ -3,43 +3,33 @@ # # SPDX-License-Identifier: BSD-3-Clause -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +"""Warp variants of the stable Cassie flat velocity task configuration.""" -from isaaclab.sim import SimulationCfg from isaaclab.utils.configclass import configclass -from .rough_env_cfg import CassieRoughEnvCfg +from isaaclab_tasks.core.velocity.config.cassie.flat_env_cfg import CassieFlatEnvCfg as _StableCassieFlatEnvCfg +from isaaclab_tasks.core.velocity.config.cassie.flat_env_cfg import ( + CassieFlatEnvCfg_PLAY as _StableCassieFlatEnvCfg_PLAY, +) + +from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( + disable_unsupported_randomization_events, +) @configclass -class CassieFlatEnvCfg(CassieRoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=52, - nconmax=15, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) +class CassieFlatEnvCfg(_StableCassieFlatEnvCfg): + """Stable flat cfg with the randomization events lacking warp twins disabled.""" def __post_init__(self): super().__post_init__() - self.rewards.flat_orientation_l2.weight = -2.5 - self.rewards.feet_air_time.weight = 5.0 - self.rewards.joint_deviation_hip.params["asset_cfg"].joint_names = ["hip_rotation_.*"] - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None + disable_unsupported_randomization_events(self) + +@configclass +class CassieFlatEnvCfg_PLAY(_StableCassieFlatEnvCfg_PLAY): + """Play variant of :class:`CassieFlatEnvCfg` for the warp runtime.""" -class CassieFlatEnvCfg_PLAY(CassieFlatEnvCfg): - def __post_init__(self) -> None: + def __post_init__(self): super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False + disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/rough_env_cfg.py deleted file mode 100644 index 2fc44ea36cdc..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/rough_env_cfg.py +++ /dev/null @@ -1,96 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg - -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import ( - LocomotionVelocityRoughEnvCfg, - RewardsCfg, -) - -from isaaclab_assets.robots.cassie import CASSIE_CFG # isort: skip - - -@configclass -class CassieRewardsCfg(RewardsCfg): - termination_penalty = RewTerm(func=mdp.is_terminated, weight=-200.0) - feet_air_time = RewTerm( - func=mdp.feet_air_time_positive_biped, - weight=2.5, - params={ - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*toe"), - "command_name": "base_velocity", - "threshold": 0.3, - }, - ) - joint_deviation_hip = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["hip_abduction_.*", "hip_rotation_.*"])}, - ) - joint_deviation_toes = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["toe_joint_.*"])}, - ) - dof_pos_limits = RewTerm( - func=mdp.joint_pos_limits, - weight=-1.0, - params={"asset_cfg": SceneEntityCfg("robot", joint_names="toe_joint_.*")}, - ) - - -@configclass -class CassieRoughEnvCfg(LocomotionVelocityRoughEnvCfg): - rewards: CassieRewardsCfg = CassieRewardsCfg() - - def __post_init__(self): - super().__post_init__() - self.scene.robot = CASSIE_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.actions.joint_pos.scale = 0.5 - self.events.push_robot = None - # TODO: TEMPORARILY DISABLED - adding this causes NaNs in the simulation - # self.events.add_base_mass = None - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.base_external_force_torque.params["asset_cfg"].body_names = [".*pelvis"] - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - self.terminations.base_contact.params["sensor_cfg"].body_names = [".*pelvis"] - self.rewards.undesired_contacts = None - self.rewards.dof_torques_l2.weight = -5.0e-6 - self.rewards.track_lin_vel_xy_exp.weight = 2.0 - self.rewards.track_ang_vel_z_exp.weight = 1.0 - self.rewards.action_rate_l2.weight *= 1.5 - self.rewards.dof_acc_l2.weight *= 1.5 - - -@configclass -class CassieRoughEnvCfg_PLAY(CassieRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/__init__.py index f49e103f431a..47d6909f9356 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/__init__.py @@ -12,28 +12,6 @@ # Register Gym environments. ## -# gym.register( -# id="Isaac-Velocity-Rough-G1-Warp-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:G1RoughEnvCfg", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1RoughPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", -# }, -# ) - - -# gym.register( -# id="Isaac-Velocity-Rough-G1-Warp-Play-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:G1RoughEnvCfg_PLAY", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1RoughPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", -# }, -# ) gym.register( id="Isaac-Velocity-Flat-G1-Warp-v0", diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py index e99dd95ff82f..9c2011fe13c8 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py @@ -3,56 +3,31 @@ # # SPDX-License-Identifier: BSD-3-Clause -from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +"""Warp variants of the stable G1 flat velocity task configuration.""" -from isaaclab.sim import SimulationCfg from isaaclab.utils.configclass import configclass -from .rough_env_cfg import G1RoughEnvCfg +from isaaclab_tasks.core.velocity.config.g1.flat_env_cfg import G1FlatEnvCfg as _StableG1FlatEnvCfg +from isaaclab_tasks.core.velocity.config.g1.flat_env_cfg import G1FlatEnvCfg_PLAY as _StableG1FlatEnvCfg_PLAY + +from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( + disable_unsupported_randomization_events, +) @configclass -class G1FlatEnvCfg(G1RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=95, - nconmax=10, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) +class G1FlatEnvCfg(_StableG1FlatEnvCfg): + """Stable flat cfg with the randomization events lacking warp twins disabled.""" def __post_init__(self): super().__post_init__() - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None - self.rewards.track_ang_vel_z_exp.weight = 1.0 - self.rewards.lin_vel_z_l2.weight = -0.2 - self.rewards.action_rate_l2.weight = -0.005 - self.rewards.dof_acc_l2.weight = -1.0e-7 - self.rewards.feet_air_time.weight = 0.75 - self.rewards.feet_air_time.params["threshold"] = 0.4 - self.rewards.dof_torques_l2.weight = -2.0e-6 - self.rewards.dof_torques_l2.params["asset_cfg"] = SceneEntityCfg( - "robot", joint_names=[".*_hip_.*", ".*_knee_joint"] - ) - self.commands.base_velocity.ranges.lin_vel_x = (0.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (-0.5, 0.5) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - - -class G1FlatEnvCfg_PLAY(G1FlatEnvCfg): - def __post_init__(self) -> None: + disable_unsupported_randomization_events(self) + + +@configclass +class G1FlatEnvCfg_PLAY(_StableG1FlatEnvCfg_PLAY): + """Play variant of :class:`G1FlatEnvCfg` for the warp runtime.""" + + def __post_init__(self): super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None + disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py deleted file mode 100644 index 7f9d12d0f201..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py +++ /dev/null @@ -1,178 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg - -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import ( - LocomotionVelocityRoughEnvCfg, - RewardsCfg, -) - -## -# Pre-defined configs -## -from isaaclab_assets import G1_MINIMAL_CFG # isort: skip - - -@configclass -class G1Rewards(RewardsCfg): - """Reward terms for the MDP.""" - - termination_penalty = RewTerm(func=mdp.is_terminated, weight=-200.0) - track_lin_vel_xy_exp = RewTerm( - func=mdp.track_lin_vel_xy_yaw_frame_exp, - weight=1.0, - params={"command_name": "base_velocity", "std": 0.5}, - ) - track_ang_vel_z_exp = RewTerm( - func=mdp.track_ang_vel_z_world_exp, weight=2.0, params={"command_name": "base_velocity", "std": 0.5} - ) - feet_air_time = RewTerm( - func=mdp.feet_air_time_positive_biped, - weight=0.25, - params={ - "command_name": "base_velocity", - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*_ankle_roll_link"), - "threshold": 0.4, - }, - ) - feet_slide = RewTerm( - func=mdp.feet_slide, - weight=-0.1, - params={ - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*_ankle_roll_link"), - "asset_cfg": SceneEntityCfg("robot", body_names=".*_ankle_roll_link"), - }, - ) - dof_pos_limits = RewTerm( - func=mdp.joint_pos_limits, - weight=-1.0, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_ankle_pitch_joint", ".*_ankle_roll_joint"])}, - ) - joint_deviation_hip = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.1, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_hip_yaw_joint", ".*_hip_roll_joint"])}, - ) - joint_deviation_arms = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.1, - params={ - "asset_cfg": SceneEntityCfg( - "robot", - joint_names=[ - ".*_shoulder_pitch_joint", - ".*_shoulder_roll_joint", - ".*_shoulder_yaw_joint", - ".*_elbow_pitch_joint", - ".*_elbow_roll_joint", - ], - ) - }, - ) - joint_deviation_fingers = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.05, - params={ - "asset_cfg": SceneEntityCfg( - "robot", - joint_names=[ - ".*_five_joint", - ".*_three_joint", - ".*_six_joint", - ".*_four_joint", - ".*_zero_joint", - ".*_one_joint", - ".*_two_joint", - ], - ) - }, - ) - joint_deviation_torso = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.1, - params={"asset_cfg": SceneEntityCfg("robot", joint_names="torso_joint")}, - ) - - -@configclass -class G1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - rewards: G1Rewards = G1Rewards() - - def __post_init__(self): - super().__post_init__() - self.scene.robot = G1_MINIMAL_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.events.push_robot = None - # TODO: TEMPORARILY DISABLED - adding this causes NaNs in the simulation - # self.events.add_base_mass = None - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.base_external_force_torque.params["asset_cfg"].body_names = ["torso_link"] - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - - # Rewards - self.rewards.lin_vel_z_l2.weight = 0.0 - self.rewards.undesired_contacts = None - self.rewards.flat_orientation_l2.weight = -1.0 - self.rewards.action_rate_l2.weight = -0.005 - self.rewards.dof_acc_l2.weight = -1.25e-7 - self.rewards.dof_acc_l2.params["asset_cfg"] = SceneEntityCfg( - "robot", joint_names=[".*_hip_.*", ".*_knee_joint"] - ) - self.rewards.dof_torques_l2.weight = -1.5e-7 - self.rewards.dof_torques_l2.params["asset_cfg"] = SceneEntityCfg( - "robot", joint_names=[".*_hip_.*", ".*_knee_joint", ".*_ankle_.*"] - ) - - # Commands - self.commands.base_velocity.ranges.lin_vel_x = (0.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (-0.0, 0.0) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - - # terminations - self.terminations.base_contact.params["sensor_cfg"].body_names = "torso_link" - - -@configclass -class G1RoughEnvCfg_PLAY(G1RoughEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.episode_length_s = 40.0 - # spawn the robot randomly in the grid (instead of their terrain levels) - self.scene.terrain.max_init_terrain_level = None - # reduce the number of terrains to save memory - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - - self.commands.base_velocity.ranges.lin_vel_x = (1.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - self.commands.base_velocity.ranges.heading = (0.0, 0.0) - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py index dd4012882075..c54355196dab 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py @@ -3,44 +3,33 @@ # # SPDX-License-Identifier: BSD-3-Clause -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +"""Warp variants of the stable Unitree Go1 flat velocity task configuration.""" -from isaaclab.sim import SimulationCfg from isaaclab.utils.configclass import configclass -from .rough_env_cfg import UnitreeGo1RoughEnvCfg +from isaaclab_tasks.contrib.velocity.config.go1.flat_env_cfg import UnitreeGo1FlatEnvCfg as _StableUnitreeGo1FlatEnvCfg +from isaaclab_tasks.contrib.velocity.config.go1.flat_env_cfg import ( + UnitreeGo1FlatEnvCfg_PLAY as _StableUnitreeGo1FlatEnvCfg_PLAY, +) + +from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( + disable_unsupported_randomization_events, +) @configclass -class UnitreeGo1FlatEnvCfg(UnitreeGo1RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=60, - nconmax=25, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) +class UnitreeGo1FlatEnvCfg(_StableUnitreeGo1FlatEnvCfg): + """Stable flat cfg with the randomization events lacking warp twins disabled.""" def __post_init__(self): super().__post_init__() - self.rewards.flat_orientation_l2.weight = -2.5 - self.rewards.feet_air_time.weight = 0.25 - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None + disable_unsupported_randomization_events(self) + +@configclass +class UnitreeGo1FlatEnvCfg_PLAY(_StableUnitreeGo1FlatEnvCfg_PLAY): + """Play variant of :class:`UnitreeGo1FlatEnvCfg` for the warp runtime.""" -class UnitreeGo1FlatEnvCfg_PLAY(UnitreeGo1FlatEnvCfg): - def __post_init__(self) -> None: + def __post_init__(self): super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None + disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py deleted file mode 100644 index cb3a602394a6..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py +++ /dev/null @@ -1,61 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -from isaaclab_assets.robots.unitree import UNITREE_GO1_CFG # isort: skip - - -@configclass -class UnitreeGo1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.robot = UNITREE_GO1_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.manager_call_max_mode = {"Scene": 1} - self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01 - self.actions.joint_pos.scale = 0.25 - self.events.push_robot = None - self.events.base_external_force_torque.params["asset_cfg"].body_names = "trunk" - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - self.rewards.feet_air_time.params["sensor_cfg"].body_names = ".*_foot" - self.rewards.feet_air_time.weight = 0.01 - self.rewards.undesired_contacts = None - self.rewards.dof_torques_l2.weight = -0.0002 - self.rewards.track_lin_vel_xy_exp.weight = 1.5 - self.rewards.track_ang_vel_z_exp.weight = 0.75 - self.rewards.dof_acc_l2.weight = -2.5e-7 - self.terminations.base_contact.params["sensor_cfg"].body_names = "trunk" - - -@configclass -class UnitreeGo1RoughEnvCfg_PLAY(UnitreeGo1RoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py index 349cc450ca9f..a54d94dabcc9 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py @@ -3,44 +3,33 @@ # # SPDX-License-Identifier: BSD-3-Clause -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +"""Warp variants of the stable Unitree Go2 flat velocity task configuration.""" -from isaaclab.sim import SimulationCfg from isaaclab.utils.configclass import configclass -from .rough_env_cfg import UnitreeGo2RoughEnvCfg +from isaaclab_tasks.core.velocity.config.go2.flat_env_cfg import UnitreeGo2FlatEnvCfg as _StableUnitreeGo2FlatEnvCfg +from isaaclab_tasks.core.velocity.config.go2.flat_env_cfg import ( + UnitreeGo2FlatEnvCfg_PLAY as _StableUnitreeGo2FlatEnvCfg_PLAY, +) + +from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( + disable_unsupported_randomization_events, +) @configclass -class UnitreeGo2FlatEnvCfg(UnitreeGo2RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=65, - nconmax=35, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) +class UnitreeGo2FlatEnvCfg(_StableUnitreeGo2FlatEnvCfg): + """Stable flat cfg with the randomization events lacking warp twins disabled.""" def __post_init__(self): super().__post_init__() - self.rewards.flat_orientation_l2.weight = -2.5 - self.rewards.feet_air_time.weight = 0.25 - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None + disable_unsupported_randomization_events(self) + +@configclass +class UnitreeGo2FlatEnvCfg_PLAY(_StableUnitreeGo2FlatEnvCfg_PLAY): + """Play variant of :class:`UnitreeGo2FlatEnvCfg` for the warp runtime.""" -class UnitreeGo2FlatEnvCfg_PLAY(UnitreeGo2FlatEnvCfg): - def __post_init__(self) -> None: + def __post_init__(self): super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None + disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/rough_env_cfg.py deleted file mode 100644 index a25748da9885..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/rough_env_cfg.py +++ /dev/null @@ -1,60 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -from isaaclab_assets.robots.unitree import UNITREE_GO2_CFG # isort: skip - - -@configclass -class UnitreeGo2RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.robot = UNITREE_GO2_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01 - self.actions.joint_pos.scale = 0.25 - self.events.push_robot = None - self.events.base_external_force_torque.params["asset_cfg"].body_names = "base" - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - self.rewards.feet_air_time.params["sensor_cfg"].body_names = ".*_foot" - self.rewards.feet_air_time.weight = 0.01 - self.rewards.undesired_contacts = None - self.rewards.dof_torques_l2.weight = -0.0002 - self.rewards.track_lin_vel_xy_exp.weight = 1.5 - self.rewards.track_ang_vel_z_exp.weight = 0.75 - self.rewards.dof_acc_l2.weight = -2.5e-7 - self.terminations.base_contact.params["sensor_cfg"].body_names = "base" - - -@configclass -class UnitreeGo2RoughEnvCfg_PLAY(UnitreeGo2RoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/__init__.py index c731af7ca986..07deee5ee458 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/__init__.py @@ -12,28 +12,6 @@ # Register Gym environments. ## -# gym.register( -# id="Isaac-Velocity-Rough-H1-Warp-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:H1RoughEnvCfg", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1RoughPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", -# }, -# ) - -# gym.register( -# id="Isaac-Velocity-Rough-H1-Warp-Play-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:H1RoughEnvCfg_PLAY", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1RoughPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", -# }, -# ) - gym.register( id="Isaac-Velocity-Flat-H1-Warp-v0", entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py index a7021343464a..088d5786b223 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py @@ -3,44 +3,31 @@ # # SPDX-License-Identifier: BSD-3-Clause -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +"""Warp variants of the stable H1 flat velocity task configuration.""" -from isaaclab.sim import SimulationCfg from isaaclab.utils.configclass import configclass -from .rough_env_cfg import H1RoughEnvCfg +from isaaclab_tasks.core.velocity.config.h1.flat_env_cfg import H1FlatEnvCfg as _StableH1FlatEnvCfg +from isaaclab_tasks.core.velocity.config.h1.flat_env_cfg import H1FlatEnvCfg_PLAY as _StableH1FlatEnvCfg_PLAY + +from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( + disable_unsupported_randomization_events, +) @configclass -class H1FlatEnvCfg(H1RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=65, - nconmax=15, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) +class H1FlatEnvCfg(_StableH1FlatEnvCfg): + """Stable flat cfg with the randomization events lacking warp twins disabled.""" def __post_init__(self): super().__post_init__() - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None - self.rewards.feet_air_time.weight = 1.0 - self.rewards.feet_air_time.params["threshold"] = 0.6 + disable_unsupported_randomization_events(self) + +@configclass +class H1FlatEnvCfg_PLAY(_StableH1FlatEnvCfg_PLAY): + """Play variant of :class:`H1FlatEnvCfg` for the warp runtime.""" -class H1FlatEnvCfg_PLAY(H1FlatEnvCfg): - def __post_init__(self) -> None: + def __post_init__(self): super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None + disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py deleted file mode 100644 index 1b65e018c9e4..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py +++ /dev/null @@ -1,131 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg - -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import ( - LocomotionVelocityRoughEnvCfg, - RewardsCfg, -) - -## -# Pre-defined configs -## -from isaaclab_assets import H1_MINIMAL_CFG # isort: skip - - -@configclass -class H1Rewards(RewardsCfg): - """Reward terms for the MDP.""" - - termination_penalty = RewTerm(func=mdp.is_terminated, weight=-200.0) - lin_vel_z_l2 = None - track_lin_vel_xy_exp = RewTerm( - func=mdp.track_lin_vel_xy_yaw_frame_exp, - weight=1.0, - params={"command_name": "base_velocity", "std": 0.5}, - ) - track_ang_vel_z_exp = RewTerm( - func=mdp.track_ang_vel_z_world_exp, weight=1.0, params={"command_name": "base_velocity", "std": 0.5} - ) - feet_air_time = RewTerm( - func=mdp.feet_air_time_positive_biped, - weight=0.25, - params={ - "command_name": "base_velocity", - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*ankle_link"), - "threshold": 0.4, - }, - ) - feet_slide = RewTerm( - func=mdp.feet_slide, - weight=-0.25, - params={ - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*ankle_link"), - "asset_cfg": SceneEntityCfg("robot", body_names=".*ankle_link"), - }, - ) - dof_pos_limits = RewTerm( - func=mdp.joint_pos_limits, weight=-1.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*_ankle")} - ) - joint_deviation_hip = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_hip_yaw", ".*_hip_roll"])}, - ) - joint_deviation_arms = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_shoulder_.*", ".*_elbow"])}, - ) - joint_deviation_torso = RewTerm( - func=mdp.joint_deviation_l1, weight=-0.1, params={"asset_cfg": SceneEntityCfg("robot", joint_names="torso")} - ) - - -@configclass -class H1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - rewards: H1Rewards = H1Rewards() - - def __post_init__(self): - super().__post_init__() - self.scene.robot = H1_MINIMAL_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.events.push_robot = None - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.base_external_force_torque.params["asset_cfg"].body_names = [".*torso_link"] - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - self.rewards.undesired_contacts = None - self.rewards.flat_orientation_l2.weight = -1.0 - self.rewards.dof_torques_l2.weight = 0.0 - self.rewards.action_rate_l2.weight = -0.005 - self.rewards.dof_acc_l2.weight = -1.25e-7 - self.commands.base_velocity.ranges.lin_vel_x = (0.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - self.terminations.base_contact.params["sensor_cfg"].body_names = ".*torso_link" - - -@configclass -class H1RoughEnvCfg_PLAY(H1RoughEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.episode_length_s = 40.0 - # spawn the robot randomly in the grid (instead of their terrain levels) - self.scene.terrain.max_init_terrain_level = None - # reduce the number of terrains to save memory - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - - self.commands.base_velocity.ranges.lin_vel_x = (1.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - self.commands.base_velocity.ranges.heading = (0.0, 0.0) - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py deleted file mode 100644 index 52d53c3db268..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py +++ /dev/null @@ -1,296 +0,0 @@ -# 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 - -import math -from dataclasses import MISSING - -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg, AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import CurriculumTermCfg as CurrTerm -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sensors import ContactSensorCfg -from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR -from isaaclab.utils.configclass import configclass -from isaaclab.utils.noise import UniformNoiseCfg as Unoise - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp - -## -# Pre-defined configs -## -from isaaclab.terrains.config.rough import ROUGH_TERRAINS_CFG # isort: skip - - -## -# Scene definition -## - - -@configclass -class MySceneCfg(InteractiveSceneCfg): - """Configuration for the terrain scene with a legged robot.""" - - # ground terrain - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="generator", - terrain_generator=ROUGH_TERRAINS_CFG, - max_init_terrain_level=5, - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg( - friction_combine_mode="multiply", - restitution_combine_mode="multiply", - static_friction=1.0, - dynamic_friction=1.0, - ), - visual_material=sim_utils.MdlFileCfg( - mdl_path=f"{ISAACLAB_NUCLEUS_DIR}/Materials/TilesMarbleSpiderWhiteBrickBondHoned/TilesMarbleSpiderWhiteBrickBondHoned.mdl", - project_uvw=True, - texture_scale=(0.25, 0.25), - ), - debug_vis=False, - ) - # robots - robot: ArticulationCfg = MISSING - # sensors - contact_forces = ContactSensorCfg( - prim_path="{ENV_REGEX_NS}/Robot/.*", - filter_prim_paths_expr=[], - history_length=3, - track_air_time=True, - ) - # lights - sky_light = AssetBaseCfg( - prim_path="/World/skyLight", - spawn=sim_utils.DomeLightCfg( - intensity=750.0, - texture_file=f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr", - ), - ) - - -## -# MDP settings -## - - -@configclass -class CommandsCfg: - """Command specifications for the MDP.""" - - base_velocity = mdp.UniformVelocityCommandCfg( - asset_name="robot", - resampling_time_range=(10.0, 10.0), - rel_standing_envs=0.02, - rel_heading_envs=1.0, - heading_command=True, - heading_control_stiffness=0.5, - debug_vis=True, - ranges=mdp.UniformVelocityCommandCfg.Ranges( - lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.0, 1.0), heading=(-math.pi, math.pi) - ), - ) - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - joint_pos = mdp.JointPositionActionCfg(asset_name="robot", joint_names=[".*"], scale=0.5, use_default_offset=True) - - -@configclass -class PolicyCfg(ObsGroup): - """Observations for policy group.""" - - # observation terms (order preserved) - base_lin_vel = ObsTerm(func=mdp.base_lin_vel, noise=Unoise(n_min=-0.1, n_max=0.1)) - base_ang_vel = ObsTerm(func=mdp.base_ang_vel, noise=Unoise(n_min=-0.2, n_max=0.2)) - projected_gravity = ObsTerm( - func=mdp.projected_gravity, - noise=Unoise(n_min=-0.05, n_max=0.05), - ) - velocity_commands = ObsTerm(func=mdp.generated_commands, params={"command_name": "base_velocity"}) - joint_pos = ObsTerm(func=mdp.joint_pos_rel, noise=Unoise(n_min=-0.01, n_max=0.01)) - joint_vel = ObsTerm(func=mdp.joint_vel_rel, noise=Unoise(n_min=-1.5, n_max=1.5)) - actions = ObsTerm(func=mdp.last_action) - - def __post_init__(self): - self.enable_corruption = True - self.concatenate_terms = True - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - # FIXME(warp-migration): COM randomization in exp manager-based locomotion currently causes - # NaNs and is temporarily disabled. - # base_com = EventTerm( - # func=mdp.randomize_rigid_body_com, - # mode="startup", - # params={ - # "asset_cfg": SceneEntityCfg("robot", body_names="base"), - # "com_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.01, 0.01)}, - # }, - # ) - base_com = None - - # reset - base_external_force_torque = EventTerm( - func=mdp.apply_external_force_torque, - mode="reset", - params={ - "asset_cfg": SceneEntityCfg("robot", body_names="base"), - "force_range": (0.0, 0.0), - "torque_range": (-0.0, 0.0), - }, - ) - - reset_base = EventTerm( - func=mdp.reset_root_state_uniform, - mode="reset", - params={ - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (-0.5, 0.5), - "y": (-0.5, 0.5), - "z": (-0.5, 0.5), - "roll": (-0.5, 0.5), - "pitch": (-0.5, 0.5), - "yaw": (-0.5, 0.5), - }, - }, - ) - - reset_robot_joints = EventTerm( - func=mdp.reset_joints_by_scale, - mode="reset", - params={ - "position_range": (0.5, 1.5), - "velocity_range": (0.0, 0.0), - }, - ) - - # interval - push_robot = EventTerm( - func=mdp.push_by_setting_velocity, - mode="interval", - interval_range_s=(10.0, 15.0), - params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}}, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # -- task - track_lin_vel_xy_exp = RewTerm( - func=mdp.track_lin_vel_xy_exp, weight=1.0, params={"command_name": "base_velocity", "std": math.sqrt(0.25)} - ) - track_ang_vel_z_exp = RewTerm( - func=mdp.track_ang_vel_z_exp, weight=0.5, params={"command_name": "base_velocity", "std": math.sqrt(0.25)} - ) - # -- penalties - lin_vel_z_l2 = RewTerm(func=mdp.lin_vel_z_l2, weight=-2.0) - ang_vel_xy_l2 = RewTerm(func=mdp.ang_vel_xy_l2, weight=-0.05) - dof_torques_l2 = RewTerm(func=mdp.joint_torques_l2, weight=-1.0e-5) - dof_acc_l2 = RewTerm(func=mdp.joint_acc_l2, weight=-2.5e-7) - action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=-0.01) - feet_air_time = RewTerm( - func=mdp.feet_air_time, - weight=0.125, - params={ - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*FOOT"), - "command_name": "base_velocity", - "threshold": 0.5, - }, - ) - undesired_contacts = RewTerm( - func=mdp.undesired_contacts, - weight=-1.0, - params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*THIGH"), "threshold": 1.0}, - ) - # -- optional penalties - flat_orientation_l2 = RewTerm(func=mdp.flat_orientation_l2, weight=0.0) - dof_pos_limits = RewTerm(func=mdp.joint_pos_limits, weight=0.0) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - time_out = DoneTerm(func=mdp.time_out, time_out=True) - base_contact = DoneTerm( - func=mdp.illegal_contact, - params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names="base"), "threshold": 1.0}, - ) - - -@configclass -class CurriculumCfg: - """Curriculum terms for the MDP.""" - - terrain_levels = CurrTerm(func=mdp.terrain_levels_vel) - - -## -# Environment configuration -## - - -@configclass -class LocomotionVelocityRoughEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the locomotion velocity-tracking environment.""" - - # Scene settings - scene: MySceneCfg = MySceneCfg(num_envs=4096, env_spacing=2.5, replicate_physics=True) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - commands: CommandsCfg = CommandsCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() - curriculum: CurriculumCfg = CurriculumCfg() - - def __post_init__(self): - """Post initialization.""" - # general settings - self.decimation = 4 - self.episode_length_s = 20.0 - # simulation settings - self.sim.dt = 1.0 / 200.0 - self.sim.render_interval = self.decimation - self.sim.physics_material = self.scene.terrain.physics_material - # update sensor update periods - if self.scene.contact_forces is not None: - self.scene.contact_forces.update_period = self.sim.dt - # check if terrain levels curriculum is enabled - if getattr(self.curriculum, "terrain_levels", None) is not None: - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.curriculum = True - else: - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.curriculum = False diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py index 8589501220cc..09abb3d8aafc 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py @@ -3,7 +3,13 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Reach experimental task registrations (manager-based).""" +"""Warp MDP twins for the stable Reach tasks. + +There is no separate task registration: run the stable joint-position reach +tasks (``Isaac-Reach-Franka``, ``Isaac-Reach-UR10``) with ``--frontend warp`` +and ``presets=newton_mjwarp``. The IK and OSC variants require controller +action twins that do not exist yet. +""" from isaaclab_experimental.envs.frontend import register_mdp_route diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/__init__.py deleted file mode 100644 index 460a30569089..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# 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 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/__init__.py deleted file mode 100644 index 9f01c7b2e729..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/__init__.py +++ /dev/null @@ -1,41 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.reach.config.franka import agents - -## -# Register Gym environments. -## - -## -# Joint Position Control -## - -gym.register( - id="Isaac-Reach-Franka-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaReachEnvCfg", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Reach-Franka-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaReachEnvCfg_PLAY", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/joint_pos_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/joint_pos_env_cfg.py deleted file mode 100644 index d6d36afd15f5..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/joint_pos_env_cfg.py +++ /dev/null @@ -1,74 +0,0 @@ -# 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 - -import math - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.manipulation.reach.mdp as mdp -from isaaclab_tasks_experimental.manager_based.manipulation.reach.reach_env_cfg import ReachEnvCfg - -## -# Pre-defined configs -## -from isaaclab_assets import FRANKA_PANDA_CFG # isort: skip - - -## -# Environment configuration -## - - -@configclass -class FrankaReachEnvCfg(ReachEnvCfg): - sim: SimulationCfg = SimulationCfg( - dt=1 / 60, - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=50, - nconmax=20, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ), - ) - - def __post_init__(self): - # post init of parent - super().__post_init__() - - # switch robot to franka - self.scene.robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - # override rewards - self.rewards.end_effector_position_tracking.params["asset_cfg"].body_names = ["panda_hand"] - self.rewards.end_effector_position_tracking_fine_grained.params["asset_cfg"].body_names = ["panda_hand"] - self.rewards.end_effector_orientation_tracking.params["asset_cfg"].body_names = ["panda_hand"] - - # override actions - self.actions.arm_action = mdp.JointPositionActionCfg( - asset_name="robot", joint_names=["panda_joint.*"], scale=0.5, use_default_offset=True - ) - # override command generator body - # end-effector is along z-direction - self.commands.ee_pose.body_name = "panda_hand" - self.commands.ee_pose.ranges.pitch = (math.pi, math.pi) - - -@configclass -class FrankaReachEnvCfg_PLAY(FrankaReachEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # disable randomization for play - self.observations.policy.enable_corruption = False diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/__init__.py deleted file mode 100644 index cdff71bf6523..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/__init__.py +++ /dev/null @@ -1,36 +0,0 @@ -# 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 - -# UR10 env disabled: USD asset has composition errors (broken asset file). -# Fails on both torch baseline and warp with: -# RuntimeError: USD stage has composition errors while loading provided stage -# Re-enable once the UR10 USD asset is fixed. - -# import gymnasium as gym -# from isaaclab_tasks.core.reach.config.ur_10 import agents - -# gym.register( -# id="Isaac-Reach-UR10-Warp-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:UR10ReachEnvCfg", -# "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UR10ReachPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", -# }, -# ) - -# gym.register( -# id="Isaac-Reach-UR10-Warp-Play-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:UR10ReachEnvCfg_PLAY", -# "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UR10ReachPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", -# }, -# ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/joint_pos_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/joint_pos_env_cfg.py deleted file mode 100644 index 42104e23b1b2..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/joint_pos_env_cfg.py +++ /dev/null @@ -1,74 +0,0 @@ -# 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 - -import math - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.manipulation.reach.mdp as mdp -from isaaclab_tasks_experimental.manager_based.manipulation.reach.reach_env_cfg import ReachEnvCfg - -## -# Pre-defined configs -## -from isaaclab_assets import UR10_CFG # isort: skip - - -## -# Environment configuration -## - - -@configclass -class UR10ReachEnvCfg(ReachEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=50, - nconmax=20, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - # post init of parent - super().__post_init__() - - # switch robot to ur10 - self.scene.robot = UR10_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - # override events - self.events.reset_robot_joints.params["position_range"] = (0.75, 1.25) - # override rewards - self.rewards.end_effector_position_tracking.params["asset_cfg"].body_names = ["ee_link"] - self.rewards.end_effector_position_tracking_fine_grained.params["asset_cfg"].body_names = ["ee_link"] - self.rewards.end_effector_orientation_tracking.params["asset_cfg"].body_names = ["ee_link"] - # override actions - self.actions.arm_action = mdp.JointPositionActionCfg( - asset_name="robot", joint_names=[".*"], scale=0.5, use_default_offset=True - ) - # override command generator body - # end-effector is along x-direction - self.commands.ee_pose.body_name = "ee_link" - self.commands.ee_pose.ranges.pitch = (math.pi / 2, math.pi / 2) - - -@configclass -class UR10ReachEnvCfg_PLAY(UR10ReachEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # disable randomization for play - self.observations.policy.enable_corruption = False diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py deleted file mode 100644 index e83a27e44db5..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py +++ /dev/null @@ -1,206 +0,0 @@ -# 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 - -from dataclasses import MISSING - -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg, AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import ActionTermCfg as ActionTerm -from isaaclab.managers import CurriculumTermCfg as CurrTerm -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.utils.configclass import configclass -from isaaclab.utils.noise import UniformNoiseCfg as Unoise - -import isaaclab_tasks_experimental.manager_based.manipulation.reach.mdp as mdp - -## -# Scene definition -## - - -@configclass -class ReachSceneCfg(InteractiveSceneCfg): - """Configuration for the scene with a robotic arm.""" - - # world - ground = AssetBaseCfg( - prim_path="/World/ground", - spawn=sim_utils.GroundPlaneCfg(), - init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, -1.05)), - ) - - # robots - robot: ArticulationCfg = MISSING - - # lights - light = AssetBaseCfg( - prim_path="/World/light", - spawn=sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=2500.0), - ) - - -## -# MDP settings -## - - -@configclass -class CommandsCfg: - """Command terms for the MDP.""" - - ee_pose = mdp.UniformPoseCommandCfg( - asset_name="robot", - body_name=MISSING, - resampling_time_range=(4.0, 4.0), - debug_vis=True, - ranges=mdp.UniformPoseCommandCfg.Ranges( - pos_x=(0.35, 0.65), - pos_y=(-0.2, 0.2), - pos_z=(0.15, 0.5), - roll=(0.0, 0.0), - pitch=MISSING, # depends on end-effector axis - yaw=(-3.14, 3.14), - ), - ) - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - arm_action: ActionTerm = MISSING - gripper_action: ActionTerm | None = None - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - @configclass - class PolicyCfg(ObsGroup): - """Observations for policy group.""" - - # observation terms (order preserved) - joint_pos = ObsTerm(func=mdp.joint_pos_rel, noise=Unoise(n_min=-0.01, n_max=0.01)) - joint_vel = ObsTerm(func=mdp.joint_vel_rel, noise=Unoise(n_min=-0.01, n_max=0.01)) - pose_command = ObsTerm(func=mdp.generated_commands, params={"command_name": "ee_pose"}) - actions = ObsTerm(func=mdp.last_action) - - def __post_init__(self): - self.enable_corruption = True - self.concatenate_terms = True - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - reset_base = EventTerm( - func=mdp.reset_root_state_uniform, - mode="reset", - params={"pose_range": {}, "velocity_range": {}}, - ) - - reset_robot_joints = EventTerm( - func=mdp.reset_joints_by_scale, - mode="reset", - params={ - "position_range": (0.5, 1.5), - "velocity_range": (0.0, 0.0), - }, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # task terms - end_effector_position_tracking = RewTerm( - func=mdp.position_command_error, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", body_names=MISSING), "command_name": "ee_pose"}, - ) - end_effector_position_tracking_fine_grained = RewTerm( - func=mdp.position_command_error_tanh, - weight=0.1, - params={"asset_cfg": SceneEntityCfg("robot", body_names=MISSING), "std": 0.1, "command_name": "ee_pose"}, - ) - end_effector_orientation_tracking = RewTerm( - func=mdp.orientation_command_error, - weight=-0.1, - params={"asset_cfg": SceneEntityCfg("robot", body_names=MISSING), "command_name": "ee_pose"}, - ) - - # action penalty - action_rate = RewTerm(func=mdp.action_rate_l2, weight=-0.0001) - joint_vel = RewTerm( - func=mdp.joint_vel_l2, - weight=-0.0001, - params={"asset_cfg": SceneEntityCfg("robot")}, - ) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - time_out = DoneTerm(func=mdp.time_out, time_out=True) - - -@configclass -class CurriculumCfg: - """Curriculum terms for the MDP.""" - - action_rate = CurrTerm( - func=mdp.modify_reward_weight, params={"term_name": "action_rate", "weight": -0.005, "num_steps": 4500} - ) - - joint_vel = CurrTerm( - func=mdp.modify_reward_weight, params={"term_name": "joint_vel", "weight": -0.001, "num_steps": 4500} - ) - - -## -# Environment configuration -## - - -@configclass -class ReachEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the reach end-effector pose tracking environment.""" - - # Scene settings - scene: ReachSceneCfg = ReachSceneCfg(num_envs=4096, env_spacing=2.5) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - commands: CommandsCfg = CommandsCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() - curriculum: CurriculumCfg = CurriculumCfg() - - def __post_init__(self): - """Post initialization.""" - # general settings - self.decimation = 2 - self.sim.render_interval = self.decimation - self.episode_length_s = 12.0 - self.viewer.eye = (3.5, 3.5, 3.5) - # simulation settings - self.sim.dt = 1.0 / 60.0 From 910baeeee66573342e4f9fd2681b98d083070943 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 15 Jul 2026 18:32:44 -0700 Subject: [PATCH 27/58] Share stable cfgs with direct warp envs via warp_entry_point Direct warp tasks duplicated their stable task configuration in a parallel *-Direct-Warp-v0 registration, so every stable cfg change had to be mirrored by hand. A stable direct registration can now declare its warp implementation with a warp_entry_point kwarg (mirroring env_cfg_entry_point); --frontend warp constructs that class with the stable cfg and swaps nothing else. Declared for Isaac-Cartpole-Direct, Isaac-Ant-Direct, and Isaac-Humanoid-Direct, whose warp cfgs matched the stable ones field-for-field; their *-Direct-Warp-v0 registrations and duplicated cfgs are removed and the env classes annotate the stable cfg types. The Allegro reorient warp task keeps its own registration: its cube is modeled as an articulation, unlike the stable task's rigid object, so the configurations genuinely differ. --- .../newton/warp-environments.rst | 19 +++- docs/source/overview/environments.rst | 15 --- .../changelog.d/warp-manager-bridge.rst | 6 + .../isaaclab_experimental/envs/frontend.py | 52 +++++++-- .../test/envs/test_frontend.py | 39 +++++-- .../test/envs/test_frontend_cfg_conversion.py | 35 +++++- .../changelog.d/warp-manager-bridge.rst | 7 ++ .../isaaclab_tasks/core/cartpole/__init__.py | 1 + .../core/locomotion/ant/__init__.py | 1 + .../core/locomotion/humanoid/__init__.py | 1 + .../changelog.d/warp-manager-bridge.minor.rst | 7 ++ .../direct/ant/__init__.py | 28 +---- .../direct/ant/ant_env_warp.py | 8 +- .../direct/ant/ant_env_warp_cfg.py | 82 -------------- .../direct/cartpole/__init__.py | 29 +---- .../direct/cartpole/cartpole_warp_env.py | 6 +- .../direct/cartpole/cartpole_warp_env_cfg.py | 74 ------------ .../direct/humanoid/__init__.py | 28 +---- .../direct/humanoid/humanoid_warp_env.py | 8 +- .../direct/humanoid/humanoid_warp_env_cfg.py | 105 ------------------ 20 files changed, 175 insertions(+), 376 deletions(-) create mode 100644 source/isaaclab_tasks/changelog.d/warp-manager-bridge.rst delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env_cfg.py diff --git a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst index 8f682ddcb047..f1c1aee930ad 100644 --- a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst +++ b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst @@ -31,9 +31,18 @@ Available Environments Direct Warp Environments ^^^^^^^^^^^^^^^^^^^^^^^^ -- ``Isaac-Cartpole-Direct-Warp-v0`` — Cartpole balance -- ``Isaac-Ant-Direct-Warp-v0`` — Ant locomotion -- ``Isaac-Humanoid-Direct-Warp-v0`` — Humanoid locomotion +Direct tasks share the stable task configuration: the stable registration +declares its warp implementation through the ``warp_entry_point`` kwarg +(mirroring ``env_cfg_entry_point``), and ``--frontend warp`` swaps only the +environment class. Stable tasks with a declared warp implementation: + +- ``Isaac-Cartpole-Direct`` — Cartpole balance +- ``Isaac-Ant-Direct`` — Ant locomotion +- ``Isaac-Humanoid-Direct`` — Humanoid locomotion + +Warp-native direct task with its own registration and configuration (the cube +is modeled differently than in the stable task): + - ``Isaac-Reorient-Cube-Allegro-Direct-Warp-v0`` — Allegro hand cube reorient @@ -78,9 +87,9 @@ Quick Start .. code-block:: bash - # Direct workflow + # Direct workflow: stable task, warp env implementation ./isaaclab.sh train --rl_library rsl_rl \ - --task Isaac-Cartpole-Direct-Warp-v0 --num_envs 4096 + --task Isaac-Cartpole-Direct --frontend warp presets=newton_mjwarp --num_envs 4096 # Manager-based workflow: stable task on the warp runtime ./isaaclab.sh train --rl_library rsl_rl \ diff --git a/docs/source/overview/environments.rst b/docs/source/overview/environments.rst index 8c821bf47bdb..5568f77e2820 100644 --- a/docs/source/overview/environments.rst +++ b/docs/source/overview/environments.rst @@ -875,11 +875,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``ovphysx``, ``physx`` - * - Isaac-Ant-Direct-Warp-v0 - - - - Direct - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - * - IsaacContrib-Assemble-Trocar-G129-Dex3 - IsaacContrib-Assemble-Trocar-G129-Dex3-Eval - Manager Based @@ -924,11 +919,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``ovphysx``, ``physx`` - * - Isaac-Cartpole-Direct-Warp-v0 - - - - Direct - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) - - * - IsaacContrib-Cartpole-Showcase-Direct - - Direct @@ -1045,11 +1035,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_mjwarp``, ``ovphysx``, ``physx`` - * - Isaac-Humanoid-Direct-Warp-v0 - - - - Direct - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Lift-Cloth-Franka - - Manager Based diff --git a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst index ccb21e9b301f..9cb34541bee0 100644 --- a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst +++ b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst @@ -7,6 +7,12 @@ Added * Added :mod:`isaaclab_experimental.envs.frontend` runtime selector and :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable` used by ``--frontend=warp`` to adapt stable cfgs onto the warp runtime. +* Added :func:`isaaclab_experimental.envs.frontend.register_mdp_route` so task + packages declare where their warp MDP twins live; the frontend holds no + per-task table. +* Added ``warp_entry_point`` registration support: a stable direct task may + declare its warp environment class, which ``--frontend=warp`` constructs + with the stable configuration. Changed ^^^^^^^ diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index eefc9327476f..72e892aebd60 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -166,10 +166,16 @@ def build( workflow = _detect_workflow(env_cfg) if workflow is Workflow.DIRECT: - # Direct workflows aren't adapted — they must already be registered - # under the warp packages (e.g. ``Isaac-Cartpole-Direct-Warp``). - _assert_direct_warp_registration(task_id) - return gym.make(task_id, cfg=env_cfg, **construct_kwargs) + # Direct workflows aren't cfg-adapted: a hand-written warp env class + # implements the task. Preferred path: the stable registration declares + # it via ``warp_entry_point`` and shares the stable cfg. Fallback: the + # task itself is registered under the warp packages (warp-native cfg). + env_class = _resolve_direct_warp_class(task_id) + if env_class is None: + _assert_direct_warp_registration(task_id) + return gym.make(task_id, cfg=env_cfg, **construct_kwargs) + _require_newton_physics(env_cfg, type(env_cfg).__name__) + return env_class(cfg=env_cfg, **construct_kwargs) # Imported lazily so that ``--frontend=torch`` callers don't pay the # ``isaaclab_experimental.envs`` import cost. The env adapts the cfg itself @@ -349,19 +355,49 @@ def _detect_workflow(cfg: Any) -> Workflow: ) -def _assert_direct_warp_registration(task_id: str) -> None: - """For direct workflows, the task must be pre-registered under the warp packages.""" +def _resolve_direct_warp_class(task_id: str) -> type | None: + """Return the warp env class a stable direct task declares, if any. + + Stable direct registrations may carry a ``warp_entry_point`` kwarg of the + form ``"module.path:ClassName"`` (mirroring ``env_cfg_entry_point``) naming + the warp implementation of the same task. The class is constructed with the + *stable* cfg, so the task needs no parallel warp registration or duplicated + configuration. + """ try: spec = gym.spec(task_id) except gym.error.NameNotFound as exc: raise FrontendIncompatibleError(f"--frontend=warp: task {task_id!r} is not registered with gymnasium.") from exc + target = (spec.kwargs or {}).get("warp_entry_point") + if not isinstance(target, str): + return None + module_name, _, class_name = target.partition(":") + try: + module = importlib.import_module(module_name) + except ImportError as exc: + raise FrontendIncompatibleError( + f"--frontend=warp: task {task_id!r} declares warp_entry_point {target!r}, but it failed to import: {exc}" + ) from exc + env_class = getattr(module, class_name, None) + if env_class is None: + raise FrontendIncompatibleError( + f"--frontend=warp: task {task_id!r} declares warp_entry_point {target!r}," + f" but {class_name!r} is not defined in {module_name!r}." + ) + return env_class + + +def _assert_direct_warp_registration(task_id: str) -> None: + """For direct workflows without a ``warp_entry_point``, the task must be warp-registered.""" + spec = gym.spec(task_id) ep = spec.entry_point module = ep if isinstance(ep, str) else (getattr(ep, "__module__", "") or "") if not module.startswith(_WARP_ROOT_PREFIXES): raise FrontendIncompatibleError( f"--frontend=warp on direct task {task_id!r}: entry_point {ep!r}" - f" is not under {list(_WARP_ROOT_PREFIXES)}. Direct tasks must be" - f" registered as a warp env class (e.g. *-Direct-Warp-v0)." + f" is not under {list(_WARP_ROOT_PREFIXES)}. Direct tasks must either" + f" declare a `warp_entry_point` in their stable registration or be" + f" registered as a warp env class." ) diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py index ee374ed1f9ef..f112a6ddf562 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend.py +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -533,8 +533,11 @@ def test_is_swap_candidate_rejects_non_callables(): _DIRECT_TEST_TASKS = { - "_Frontend-Test-Warp-Direct-v0": "isaaclab_tasks_experimental.fake:DirectEnv", - "_Frontend-Test-Stable-Direct-v0": "isaaclab_tasks.fake:DirectEnv", + "_Frontend-Test-Warp-Direct-v0": ("isaaclab_tasks_experimental.fake:DirectEnv", None), + "_Frontend-Test-Stable-Direct-v0": ("isaaclab_tasks.fake:DirectEnv", None), + "_Frontend-Test-Declared-Direct-v0": ("isaaclab_tasks.fake:DirectEnv", "types:SimpleNamespace"), + "_Frontend-Test-Broken-Direct-v0": ("isaaclab_tasks.fake:DirectEnv", "definitely_not_a_module:Env"), + "_Frontend-Test-Missing-Class-Direct-v0": ("isaaclab_tasks.fake:DirectEnv", "types:NoSuchClass"), } @@ -542,10 +545,12 @@ def test_is_swap_candidate_rejects_non_callables(): def _register_direct_stub_tasks(): # Register stub tasks so ``gym.spec`` resolves them. Entry-point strings # are never invoked (the guard only inspects them), so a fake import - # path is fine. - for task_id, entry_point in _DIRECT_TEST_TASKS.items(): + # path is fine. ``warp_entry_point`` targets are imported, so the valid + # stub points at a real stdlib symbol. + for task_id, (entry_point, warp_entry_point) in _DIRECT_TEST_TASKS.items(): + kwargs = {"warp_entry_point": warp_entry_point} if warp_entry_point else {} with contextlib.suppress(gym.error.Error): - gym.register(id=task_id, entry_point=entry_point, disable_env_checker=True) + gym.register(id=task_id, entry_point=entry_point, disable_env_checker=True, kwargs=kwargs) yield @@ -559,6 +564,26 @@ def test_direct_guard_rejects_stable_rooted(): assert "isaaclab_experimental" in str(exc.value) -def test_direct_guard_rejects_unknown_task(): +def test_resolve_direct_warp_class_returns_declared_class(): + import types as types_module + + assert fe._resolve_direct_warp_class("_Frontend-Test-Declared-Direct-v0") is types_module.SimpleNamespace + + +def test_resolve_direct_warp_class_returns_none_without_declaration(): + assert fe._resolve_direct_warp_class("_Frontend-Test-Stable-Direct-v0") is None + + +def test_resolve_direct_warp_class_rejects_broken_module(): + with pytest.raises(FrontendIncompatibleError): + fe._resolve_direct_warp_class("_Frontend-Test-Broken-Direct-v0") + + +def test_resolve_direct_warp_class_rejects_missing_class(): + with pytest.raises(FrontendIncompatibleError): + fe._resolve_direct_warp_class("_Frontend-Test-Missing-Class-Direct-v0") + + +def test_resolve_direct_warp_class_rejects_unknown_task(): with pytest.raises(FrontendIncompatibleError): - fe._assert_direct_warp_registration("Frontend-Test-NotRegistered-v0") + fe._resolve_direct_warp_class("Frontend-Test-NotRegistered-v0") diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index 8ad8716bfd89..5509ec25a6f6 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -48,8 +48,8 @@ ] # Manager-based warp tasks are exactly those whose entry point is the shared warp -# env class; direct ``*-Direct-Warp-v0`` tasks register their own env class and -# are not cfg-adapted, so they are excluded here. +# env class; direct tasks provide their own env class (via ``warp_entry_point`` +# or a warp-native registration) and are not cfg-adapted, so they are excluded here. _MANAGER_WARP_ENTRY_POINT = "isaaclab_experimental.envs:ManagerBasedRLEnvWarp" _WARP_ROOTS = ("isaaclab_experimental", "isaaclab_tasks_experimental") @@ -126,6 +126,37 @@ def test_registered_warp_variant_cfg_adapts(task_id: str, cfg_entry_point: str): ) +def _direct_tasks_with_warp_entry_point() -> list[tuple[str, str]]: + """Return ``(task_id, warp_entry_point)`` for stable direct tasks declaring a warp env.""" + tasks: list[tuple[str, str]] = [] + for task_id, spec in gym.registry.items(): + target = (spec.kwargs or {}).get("warp_entry_point") + if isinstance(target, str): + tasks.append((task_id, target)) + return sorted(tasks) + + +_DIRECT_WARP_TASKS = _direct_tasks_with_warp_entry_point() + + +def test_direct_tasks_declare_expected_warp_entry_points(): + """Sanity: the stable direct tasks with warp implementations declare them.""" + declared = {task_id for task_id, _ in _DIRECT_WARP_TASKS} + assert {"Isaac-Cartpole-Direct", "Isaac-Ant-Direct", "Isaac-Humanoid-Direct"} <= declared + + +@pytest.mark.parametrize( + "task_id, warp_entry_point", + [pytest.param(task_id, target, id=task_id) for task_id, target in _DIRECT_WARP_TASKS], +) +def test_declared_direct_warp_entry_point_resolves(task_id: str, warp_entry_point: str): + """Each declared ``warp_entry_point`` imports and names a warp-rooted class.""" + module_path, class_name = warp_entry_point.split(":") + env_class = getattr(importlib.import_module(module_path), class_name) + assert isinstance(env_class, type), f"{task_id}: {warp_entry_point} is not a class" + assert env_class.__module__.startswith(_WARP_ROOTS) + + def test_stable_cartpole_cfg_adapts_to_current_warp_module_layout(): """The stable Cartpole cfg resolves task-specific twins in the current package layout.""" from isaaclab_experimental.managers.action_manager import ActionTerm diff --git a/source/isaaclab_tasks/changelog.d/warp-manager-bridge.rst b/source/isaaclab_tasks/changelog.d/warp-manager-bridge.rst new file mode 100644 index 000000000000..f34db06907a5 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/warp-manager-bridge.rst @@ -0,0 +1,7 @@ +Added +^^^^^ + +* Added ``warp_entry_point`` declarations to the ``Isaac-Cartpole-Direct``, + ``Isaac-Ant-Direct``, and ``Isaac-Humanoid-Direct`` registrations so + ``--frontend warp`` constructs the Warp implementation of these tasks with + the stable task configuration. diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py index bdb55030b3b2..45c6ca0572fa 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py @@ -25,6 +25,7 @@ disable_env_checker=True, kwargs={ "env_cfg_entry_point": f"{__name__}.cartpole_direct_env_cfg:CartpoleEnvCfg", + "warp_entry_point": "isaaclab_tasks_experimental.direct.cartpole.cartpole_warp_env:CartpoleWarpEnv", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_direct_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpoleDirectPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_direct_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/__init__.py index 4adb37ca8eb6..ae86a5269bde 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/__init__.py @@ -25,6 +25,7 @@ disable_env_checker=True, kwargs={ "env_cfg_entry_point": f"{__name__}.ant_direct_env_cfg:AntEnvCfg", + "warp_entry_point": "isaaclab_tasks_experimental.direct.ant.ant_env_warp:AntWarpEnv", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_direct_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AntDirectPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_direct_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/__init__.py index d13b7697595a..0387402a7957 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/__init__.py @@ -25,6 +25,7 @@ disable_env_checker=True, kwargs={ "env_cfg_entry_point": f"{__name__}.humanoid_direct_env_cfg:HumanoidEnvCfg", + "warp_entry_point": "isaaclab_tasks_experimental.direct.humanoid.humanoid_warp_env:HumanoidWarpEnv", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_direct_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:HumanoidDirectPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_direct_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst index dfb8c6270757..782b6f25ef9a 100644 --- a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst +++ b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst @@ -16,6 +16,13 @@ Removed together with their duplicated environment configurations. Run the stable task ids with ``--frontend warp`` and ``presets=newton_mjwarp`` instead, e.g. ``--task Isaac-Cartpole --frontend warp presets=newton_mjwarp``. +* **Breaking:** Removed the duplicated direct warp task registrations + ``Isaac-Cartpole-Direct-Warp-v0``, ``Isaac-Ant-Direct-Warp-v0``, and + ``Isaac-Humanoid-Direct-Warp-v0`` together with their duplicated + environment configurations; the stable registrations declare the warp + environment classes via ``warp_entry_point``. Run the stable task ids with + ``--frontend warp`` and ``presets=newton_mjwarp`` instead, e.g. + ``--task Isaac-Cartpole-Direct --frontend warp presets=newton_mjwarp``. * Removed the unregistered rough velocity warp configurations; rough-terrain warp tasks remain unsupported until :class:`~isaaclab.terrains.TerrainImporter` gains Warp APIs. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/__init__.py index 33d5ac926387..fd466b5394a3 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/__init__.py @@ -3,26 +3,10 @@ # # SPDX-License-Identifier: BSD-3-Clause -""" -Ant locomotion environment. -""" - -import gymnasium as gym +"""Warp implementation of the stable direct Ant task. -## -# Register Gym environments. -## - -stable_agents = "isaaclab_tasks.core.locomotion.ant.agents" - -gym.register( - id="Isaac-Ant-Direct-Warp-v0", - entry_point=f"{__name__}.ant_env_warp:AntWarpEnv", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.ant_env_warp_cfg:AntWarpEnvCfg", - "rl_games_cfg_entry_point": f"{stable_agents}:rl_games_direct_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{stable_agents}.rsl_rl_ppo_cfg:AntDirectPPORunnerCfg", - "skrl_cfg_entry_point": f"{stable_agents}:skrl_direct_ppo_cfg.yaml", - }, -) +There is no separate task registration: the stable ``Isaac-Ant-Direct`` +registration declares :class:`~isaaclab_tasks_experimental.direct.ant.ant_env_warp.AntWarpEnv` +as its ``warp_entry_point``; run it with ``--frontend warp`` and +``presets=newton_mjwarp``. +""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp.py index 2bd7390039b7..17ce3c2ee788 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp.py @@ -5,13 +5,13 @@ from __future__ import annotations -from isaaclab_tasks_experimental.direct.locomotion.locomotion_env_warp import LocomotionWarpEnv +from isaaclab_tasks.core.locomotion.ant.ant_direct_env_cfg import AntEnvCfg -from .ant_env_warp_cfg import AntWarpEnvCfg +from isaaclab_tasks_experimental.direct.locomotion.locomotion_env_warp import LocomotionWarpEnv class AntWarpEnv(LocomotionWarpEnv): - cfg: AntWarpEnvCfg + cfg: AntEnvCfg - def __init__(self, cfg: AntWarpEnvCfg, render_mode: str | None = None, **kwargs): + def __init__(self, cfg: AntEnvCfg, render_mode: str | None = None, **kwargs): super().__init__(cfg, render_mode, **kwargs) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp_cfg.py deleted file mode 100644 index 6e8b7fb81a95..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp_cfg.py +++ /dev/null @@ -1,82 +0,0 @@ -# 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 - -from __future__ import annotations - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg -from isaaclab.envs import DirectRLEnvCfg -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils.configclass import configclass - -from isaaclab_assets import ANT_CFG - - -@configclass -class AntWarpEnvCfg(DirectRLEnvCfg): - # env - episode_length_s = 15.0 - decimation = 2 - action_scale = 0.5 - action_space = 8 - observation_space = 36 - state_space = 0 - - solver_cfg = MJWarpSolverCfg( - njmax=45, - nconmax=25, - cone="pyramidal", - integrator="implicitfast", - impratio=1, - ) - newton_cfg = NewtonCfg( - solver_cfg=solver_cfg, - num_substeps=1, - debug_mode=False, - use_cuda_graph=True, - ) - - # simulation - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation, physics=newton_cfg) - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="plane", - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg( - friction_combine_mode="average", - restitution_combine_mode="average", - static_friction=1.0, - dynamic_friction=1.0, - restitution=0.0, - ), - debug_vis=False, - ) - - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg( - num_envs=4096, env_spacing=4.0, replicate_physics=True, clone_in_fabric=True - ) - - # robot - robot: ArticulationCfg = ANT_CFG.replace(prim_path="/World/envs/env_.*/Robot") - joint_gears: list = [15, 15, 15, 15, 15, 15, 15, 15] - - heading_weight: float = 0.5 - up_weight: float = 0.1 - - energy_cost_scale: float = 0.05 - actions_cost_scale: float = 0.005 - alive_reward_scale: float = 0.5 - dof_vel_scale: float = 0.2 - - death_cost: float = -2.0 - termination_height: float = 0.31 - - angular_velocity_scale: float = 1.0 - contact_force_scale: float = 0.1 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py index a7fc9492fff6..8dea234be2f8 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py @@ -3,27 +3,10 @@ # # SPDX-License-Identifier: BSD-3-Clause -""" -Cartpole balancing environment. -""" - -import gymnasium as gym +"""Warp implementation of the stable direct Cartpole task. -## -# Register Gym environments. -## - -stable_agents = "isaaclab_tasks.core.cartpole.agents" - -gym.register( - id="Isaac-Cartpole-Direct-Warp-v0", - entry_point=f"{__name__}.cartpole_warp_env:CartpoleWarpEnv", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.cartpole_warp_env_cfg:CartpoleWarpEnvCfg", - "rl_games_cfg_entry_point": f"{stable_agents}:rl_games_direct_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{stable_agents}.rsl_rl_ppo_cfg:CartpoleDirectPPORunnerCfg", - "skrl_cfg_entry_point": f"{stable_agents}:skrl_direct_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{stable_agents}:sb3_ppo_cfg.yaml", - }, -) +There is no separate task registration: the stable ``Isaac-Cartpole-Direct`` +registration declares :class:`~isaaclab_tasks_experimental.direct.cartpole.cartpole_warp_env.CartpoleWarpEnv` +as its ``warp_entry_point``; run it with ``--frontend warp`` and +``presets=newton_mjwarp``. +""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py index 3e8bee2c4070..8ebae3c986b7 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py @@ -17,7 +17,7 @@ from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane if TYPE_CHECKING: - from .cartpole_warp_env_cfg import CartpoleWarpEnvCfg + from isaaclab_tasks.core.cartpole.cartpole_direct_env_cfg import CartpoleEnvCfg @wp.kernel @@ -184,9 +184,9 @@ def initialize_state( class CartpoleWarpEnv(DirectRLEnvWarp): - cfg: CartpoleWarpEnvCfg + cfg: CartpoleEnvCfg - def __init__(self, cfg: CartpoleWarpEnvCfg, render_mode: str | None = None, **kwargs) -> None: + def __init__(self, cfg: CartpoleEnvCfg, render_mode: str | None = None, **kwargs) -> None: super().__init__(cfg, render_mode, **kwargs) # Get the indices (develop API: find_joints returns (indices, names)) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env_cfg.py deleted file mode 100644 index a11299635cfa..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env_cfg.py +++ /dev/null @@ -1,74 +0,0 @@ -# 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 - -from __future__ import annotations - -import math - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.assets import ArticulationCfg -from isaaclab.envs import DirectRLEnvCfg, ViewerCfg -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from isaaclab_assets.robots.cartpole import CARTPOLE_CFG - - -@configclass -class CartpoleWarpEnvCfg(DirectRLEnvCfg): - # env - decimation = 2 - episode_length_s = 5.0 - action_scale = 100.0 # [N] - action_space = 1 - observation_space = 4 - state_space = 0 - - # viewer - viewer: ViewerCfg = ViewerCfg(eye=(8.0, 0.0, 5.0)) - - solver_cfg = MJWarpSolverCfg( - njmax=5, - nconmax=3, - cone="pyramidal", - integrator="implicitfast", - impratio=1, - ) - - newton_cfg = NewtonCfg( - solver_cfg=solver_cfg, - num_substeps=1, - debug_mode=False, - use_cuda_graph=True, - ) - - # simulation - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation, physics=newton_cfg) - - # robot - robot_cfg: ArticulationCfg = CARTPOLE_CFG.replace(prim_path="/World/envs/env_.*/Robot") - cart_dof_name = "slider_to_cart" - pole_dof_name = "cart_to_pole" - - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg( - num_envs=4096, env_spacing=4.0, replicate_physics=True, clone_in_fabric=True - ) - - # reset - max_cart_pos = 3.0 # the cart is reset if it exceeds that position [m] - initial_cart_position_range = (-1.0, 1.0) # [m] - initial_cart_velocity_range = (-0.5, 0.5) # [m/s] - initial_pole_angle_range = (-0.25 * math.pi, 0.25 * math.pi) # [rad] - initial_pole_velocity_range = (-0.25 * math.pi, 0.25 * math.pi) # [rad/s] - - # reward scales - rew_scale_alive = 1.0 - rew_scale_terminated = -2.0 - rew_scale_pole_pos = -1.0 - rew_scale_cart_vel = -0.01 - rew_scale_pole_vel = -0.005 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/__init__.py index c38d7dd955e5..663e41a7c283 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/__init__.py @@ -3,26 +3,10 @@ # # SPDX-License-Identifier: BSD-3-Clause -""" -Humanoid locomotion environment. -""" - -import gymnasium as gym +"""Warp implementation of the stable direct Humanoid task. -## -# Register Gym environments. -## - -stable_agents = "isaaclab_tasks.core.locomotion.humanoid.agents" - -gym.register( - id="Isaac-Humanoid-Direct-Warp-v0", - entry_point=f"{__name__}.humanoid_warp_env:HumanoidWarpEnv", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.humanoid_warp_env_cfg:HumanoidWarpEnvCfg", - "rl_games_cfg_entry_point": f"{stable_agents}:rl_games_direct_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{stable_agents}.rsl_rl_ppo_cfg:HumanoidDirectPPORunnerCfg", - "skrl_cfg_entry_point": f"{stable_agents}:skrl_direct_ppo_cfg.yaml", - }, -) +There is no separate task registration: the stable ``Isaac-Humanoid-Direct`` +registration declares :class:`~isaaclab_tasks_experimental.direct.humanoid.humanoid_warp_env.HumanoidWarpEnv` +as its ``warp_entry_point``; run it with ``--frontend warp`` and +``presets=newton_mjwarp``. +""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env.py index 2fd0ac8433e2..113d761b711c 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env.py @@ -5,13 +5,13 @@ from __future__ import annotations -from isaaclab_tasks_experimental.direct.locomotion.locomotion_env_warp import LocomotionWarpEnv +from isaaclab_tasks.core.locomotion.humanoid.humanoid_direct_env_cfg import HumanoidEnvCfg -from .humanoid_warp_env_cfg import HumanoidWarpEnvCfg +from isaaclab_tasks_experimental.direct.locomotion.locomotion_env_warp import LocomotionWarpEnv class HumanoidWarpEnv(LocomotionWarpEnv): - cfg: HumanoidWarpEnvCfg + cfg: HumanoidEnvCfg - def __init__(self, cfg: HumanoidWarpEnvCfg, render_mode: str | None = None, **kwargs): + def __init__(self, cfg: HumanoidEnvCfg, render_mode: str | None = None, **kwargs): super().__init__(cfg, render_mode, **kwargs) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env_cfg.py deleted file mode 100644 index d7fa473eda38..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env_cfg.py +++ /dev/null @@ -1,105 +0,0 @@ -# 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 - -from __future__ import annotations - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg -from isaaclab.envs import DirectRLEnvCfg -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils.configclass import configclass - -from isaaclab_assets import HUMANOID_CFG - - -@configclass -class HumanoidWarpEnvCfg(DirectRLEnvCfg): - # env - episode_length_s = 15.0 - decimation = 2 - action_scale = 1.0 - action_space = 21 - observation_space = 75 - state_space = 0 - - solver_cfg = MJWarpSolverCfg( - njmax=80, - nconmax=25, - cone="pyramidal", - integrator="implicitfast", - update_data_interval=2, - impratio=1, - ) - newton_cfg = NewtonCfg( - solver_cfg=solver_cfg, - num_substeps=2, - debug_mode=False, - use_cuda_graph=True, - ) - - # simulation - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation, physics=newton_cfg) - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="plane", - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg( - friction_combine_mode="average", - restitution_combine_mode="average", - static_friction=1.0, - dynamic_friction=1.0, - restitution=0.0, - ), - debug_vis=False, - ) - - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg( - num_envs=4096, env_spacing=4.0, replicate_physics=True, clone_in_fabric=True - ) - - # robot - robot: ArticulationCfg = HUMANOID_CFG.replace(prim_path="/World/envs/env_.*/Robot") - joint_gears: list = [ - 67.5000, # left_upper_arm - 67.5000, # left_upper_arm - 45.0000, # left_lower_arm - 67.5000, # lower_waist - 67.5000, # lower_waist - 67.5000, # pelvis - 45.0000, # left_thigh: x - 135.0000, # left_thigh: y - 45.0000, # left_thigh: z - 90.0000, # left_knee - 22.5, # left_foot - 22.5, # left_foot - 45.0000, # right_thigh: x - 135.0000, # right_thigh: y - 45.0000, # right_thigh: z - 90.0000, # right_knee - 22.5, # right_foot - 22.5, # right_foot - 67.5000, # right_upper_arm - 67.5000, # right_upper_arm - 45.0000, # right_lower_arm - ] - - heading_weight: float = 0.5 - up_weight: float = 0.1 - - energy_cost_scale: float = 0.05 - actions_cost_scale: float = 0.01 - alive_reward_scale: float = 2.0 - dof_vel_scale: float = 0.1 - - death_cost: float = -1.0 - termination_height: float = 0.8 - - angular_velocity_scale: float = 0.25 - contact_force_scale: float = 0.01 From 9e2d4c81050625a15c663b10b0edbdfe96e9683b Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 15 Jul 2026 23:59:39 -0700 Subject: [PATCH 28/58] Tidy frontend CLI helper and warp env cleanups Extract the --frontend flag into add_frontend_args so other entrypoints (e.g. play) can reuse it without pulling in all training arguments; add_common_train_args keeps calling it so every train and benchmark CLI registers the flag. Also derive SceneEntityCfg.from_stable from the stable dataclass fields instead of a hand-maintained list, drop single-use aliases in the wrench observation twin, and trim render comments to one line. --- scripts/reinforcement_learning/common.py | 37 +++++++++++++------ .../envs/direct_rl_env_warp.py | 8 +--- .../envs/manager_based_rl_env_warp.py | 4 +- .../envs/mdp/observations.py | 6 +-- .../managers/scene_entity_cfg.py | 21 ++--------- 5 files changed, 33 insertions(+), 43 deletions(-) diff --git a/scripts/reinforcement_learning/common.py b/scripts/reinforcement_learning/common.py index a36b5ac9266b..efebd09747fe 100644 --- a/scripts/reinforcement_learning/common.py +++ b/scripts/reinforcement_learning/common.py @@ -233,6 +233,30 @@ def resolve_play_checkpoint(checkpoint: str | None, framework: str, task: str) - return path +def add_frontend_args(parser: argparse.ArgumentParser) -> None: + """Add the environment-runtime selector argument. + + The flag is always registered so the CLI surface does not depend on + optional packages; the warp runtime itself is imported only when selected + (see :func:`create_isaaclab_env`). + + Args: + parser: The parser to add the argument to. + """ + parser.add_argument( + "--frontend", + type=str, + choices=["torch", "warp"], + default="torch", + help=( + "Runtime that constructs the environment. 'torch' uses the registered stable environment via" + " gym.make. 'warp' (experimental) adapts a manager-based task config onto the Warp runtime, or" + " dispatches a direct task to its registered Warp environment; requires isaaclab_experimental" + " and `presets=newton_mjwarp`." + ), + ) + + def add_common_train_args( parser: argparse.ArgumentParser, *, @@ -257,18 +281,7 @@ def add_common_train_args( ) 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( - "--frontend", - type=str, - choices=["torch", "warp"], - default="torch", - help=( - "Runtime that constructs the environment. 'torch' uses the registered stable environment via" - " gym.make. 'warp' (experimental) adapts a manager-based task config onto the Warp runtime, or" - " dispatches a direct task to its registered Warp environment; requires isaaclab_experimental" - " and `presets=newton_mjwarp`." - ), - ) + add_frontend_args(parser) if include_agent: parser.add_argument("--agent", type=str, default=agent_default, help=agent_help) parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment") diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py index ae57af4bb55f..f1c015c98515 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py @@ -170,9 +170,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs # Initialize when a Kit viewport exists. ViewportCameraController uses omni.kit # (renderer camera); skip in kitless Newton-only runs where no Kit app is running. - # Mirrors stable :class:`isaaclab.envs.DirectRLEnv` (PR #5103 changed the visualizer API). - has_visualizers = self.sim.has_active_visualizers() - if (self.sim.has_gui or has_visualizers) and has_kit(): + if (self.sim.has_gui or self.sim.has_active_visualizers()) and has_kit(): self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) else: self.viewport_camera_controller = None @@ -567,9 +565,7 @@ def render(self, recompute: bool = False) -> np.ndarray | None: if self.render_mode == "human" or self.render_mode is None: return None elif self.render_mode == "rgb_array": - # Rendering is possible if we have GUI or offscreen rendering enabled — mirror - # stable (PR #4646 replaced /isaaclab/has_gui and /isaaclab/render/offscreen - # settings with cached SimulationContext properties). + # rendering requires a GUI or offscreen rendering (mirrors the stable env) if not (self.sim.has_gui or self.sim.has_offscreen_render): render_mode_name = "NO_GUI_OR_RENDERING" raise RuntimeError( diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index a90489264a48..502c9a132042 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -416,9 +416,7 @@ def render(self, recompute: bool = False) -> np.ndarray | None: if self.render_mode == "human" or self.render_mode is None: return None elif self.render_mode == "rgb_array": - # check that if any render could have happened — mirror stable - # (PR #4646 replaced /isaaclab/has_gui and /isaaclab/render/offscreen settings - # with cached SimulationContext properties). + # rendering requires a GUI or offscreen rendering (mirrors the stable env) if not (self.sim.has_gui or self.sim.has_offscreen_render): raise RuntimeError( f"Cannot render '{self.render_mode}' when the simulation render mode does not support" diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py index e36ddd189f27..4bc49cc1e2ae 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py @@ -403,13 +403,11 @@ def body_incoming_wrench(env: ManagerBasedEnv, out, sensor_cfg: SceneEntityCfg) into ``out`` (shape ``(num_envs, 6 * num_selected_bodies)``). """ sensor = env.scene.sensors[sensor_cfg.name] - force = sensor.data.force - torque = sensor.data.torque - if force is None or torque is None: + if sensor.data.force is None or sensor.data.torque is None: raise RuntimeError("Joint wrench sensor data is not initialized. Call sim.reset() before reading observations.") wp.launch( kernel=_body_incoming_wrench_kernel, dim=env.num_envs, - inputs=[force.warp, torque.warp, sensor_cfg.body_ids_wp, out], + inputs=[sensor.data.force.warp, sensor.data.torque.warp, sensor_cfg.body_ids_wp, out], device=env.device, ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py b/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py index 1052a7ec25ed..f6d30489eeca 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py @@ -42,25 +42,10 @@ class SceneEntityCfg(_SceneEntityCfg): def from_stable(cls, stable: _SceneEntityCfg) -> SceneEntityCfg: """Build a warp scene-entity cfg from a stable one. - Copies every selection field declared on the stable cfg (name, joint - names/ids, body names/ids, fixed-tendon names/ids, object-collection - names/ids, ``preserve_order``); warp-specific fields stay ``None`` - and are filled by :meth:`resolve` at scene build time. Used by the - warp frontend to promote :class:`SceneEntityCfg` instances embedded - in term ``params`` without resorting to a ``__class__`` reassign. + Copies every field declared on the stable cfg; the warp-specific fields + stay ``None`` and are filled by :meth:`resolve` at scene build time. """ - return cls( - name=stable.name, - joint_names=stable.joint_names, - joint_ids=stable.joint_ids, - fixed_tendon_names=stable.fixed_tendon_names, - fixed_tendon_ids=stable.fixed_tendon_ids, - body_names=stable.body_names, - body_ids=stable.body_ids, - object_collection_names=stable.object_collection_names, - object_collection_ids=stable.object_collection_ids, - preserve_order=stable.preserve_order, - ) + return cls(**{name: getattr(stable, name) for name in _SceneEntityCfg.__dataclass_fields__}) def resolve(self, scene: InteractiveScene): # run the stable resolution first (fills joint_ids/body_ids from names/regex) From 132112c7dfcac9a863b9f08904180763ca1474c0 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 16 Jul 2026 00:07:43 -0700 Subject: [PATCH 29/58] Mirror the stable core task layout in the experimental package MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The experimental task package still used the manager_based/direct split that the stable package left behind when it consolidated each task under core/. That mismatch was the only reason the warp MDP routing needed non-obvious mappings. Move the warp task packages to the same layout — core/cartpole, core/locomotion/{ant,humanoid}, core/velocity, core/reach, and core/reorient for the Allegro task — with each task package holding both its manager-based MDP twins and its direct warp env class, exactly like its stable counterpart. Routes and warp_entry_point declarations now read as the mechanical mirror they are. --- .../test/envs/test_frontend_cfg_conversion.py | 8 ++------ .../isaaclab_tasks/core/cartpole/__init__.py | 2 +- .../core/locomotion/ant/__init__.py | 2 +- .../core/locomotion/humanoid/__init__.py | 2 +- .../changelog.d/warp-manager-bridge.minor.rst | 4 ++++ .../{direct => core}/__init__.py | 6 +----- .../classic => core}/cartpole/__init__.py | 8 +++++--- .../cartpole/cartpole_warp_env.py | 0 .../classic => core}/cartpole/mdp/__init__.py | 0 .../classic => core}/cartpole/mdp/__init__.pyi | 0 .../classic => core}/cartpole/mdp/rewards.py | 0 .../locomotion}/__init__.py | 2 +- .../locomotion}/ant/__init__.py | 11 +++++++---- .../locomotion}/ant/ant_env_warp.py | 2 +- .../core/locomotion/humanoid/__init__.py | 18 ++++++++++++++++++ .../locomotion}/humanoid/humanoid_warp_env.py | 2 +- .../locomotion}/humanoid/mdp/__init__.py | 0 .../locomotion}/humanoid/mdp/__init__.pyi | 0 .../locomotion}/humanoid/mdp/observations.py | 0 .../locomotion}/humanoid/mdp/rewards.py | 0 .../locomotion/locomotion_env_warp.py | 0 .../manipulation => core}/reach/__init__.py | 0 .../reach/mdp/__init__.py | 0 .../reach/mdp/__init__.pyi | 0 .../manipulation => core}/reach/mdp/rewards.py | 0 .../locomotion => core/reorient}/__init__.py | 2 ++ .../reorient/config}/__init__.py | 2 ++ .../reorient/config}/allegro_hand/__init__.py | 2 +- .../allegro_hand/allegro_hand_warp_env_cfg.py | 0 .../reorient}/inhand_manipulation_warp_env.py | 4 +++- .../locomotion => core}/velocity/__init__.py | 0 .../velocity/config/__init__.py | 0 .../velocity/config/a1/__init__.py | 0 .../velocity/config/a1/flat_env_cfg.py | 2 +- .../velocity/config/anymal_b/__init__.py | 0 .../velocity/config/anymal_b/flat_env_cfg.py | 2 +- .../velocity/config/anymal_c/__init__.py | 0 .../velocity/config/anymal_c/flat_env_cfg.py | 2 +- .../velocity/config/anymal_d/__init__.py | 0 .../velocity/config/anymal_d/flat_env_cfg.py | 2 +- .../velocity/config/cassie/__init__.py | 0 .../velocity/config/cassie/flat_env_cfg.py | 2 +- .../velocity/config/g1/__init__.py | 0 .../velocity/config/g1/flat_env_cfg.py | 2 +- .../velocity/config/go1/__init__.py | 0 .../velocity/config/go1/flat_env_cfg.py | 2 +- .../velocity/config/go2/__init__.py | 0 .../velocity/config/go2/flat_env_cfg.py | 2 +- .../velocity/config/h1/__init__.py | 0 .../velocity/config/h1/flat_env_cfg.py | 2 +- .../velocity/mdp/__init__.py | 0 .../velocity/mdp/__init__.pyi | 0 .../velocity/mdp/curriculums.py | 0 .../velocity/mdp/rewards.py | 0 .../velocity/mdp/terminations.py | 0 .../direct/ant/__init__.py | 12 ------------ .../direct/cartpole/__init__.py | 12 ------------ .../direct/humanoid/__init__.py | 12 ------------ .../manager_based/__init__.py | 10 ---------- .../manager_based/classic/__init__.py | 12 ------------ .../manager_based/classic/humanoid/__init__.py | 15 --------------- .../manager_based/locomotion/__init__.py | 6 ------ 62 files changed, 60 insertions(+), 114 deletions(-) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct => core}/__init__.py (68%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic => core}/cartpole/__init__.py (52%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct => core}/cartpole/cartpole_warp_env.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic => core}/cartpole/mdp/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic => core}/cartpole/mdp/__init__.pyi (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic => core}/cartpole/mdp/rewards.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/manipulation => core/locomotion}/__init__.py (79%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic => core/locomotion}/ant/__init__.py (50%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct => core/locomotion}/ant/ant_env_warp.py (83%) create mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/__init__.py rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct => core/locomotion}/humanoid/humanoid_warp_env.py (84%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic => core/locomotion}/humanoid/mdp/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic => core/locomotion}/humanoid/mdp/__init__.pyi (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic => core/locomotion}/humanoid/mdp/observations.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/classic => core/locomotion}/humanoid/mdp/rewards.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct => core}/locomotion/locomotion_env_warp.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/manipulation => core}/reach/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/manipulation => core}/reach/mdp/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/manipulation => core}/reach/mdp/__init__.pyi (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/manipulation => core}/reach/mdp/rewards.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct/locomotion => core/reorient}/__init__.py (75%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct/inhand_manipulation => core/reorient/config}/__init__.py (77%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct => core/reorient/config}/allegro_hand/__init__.py (92%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct => core/reorient/config}/allegro_hand/allegro_hand_warp_env_cfg.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{direct/inhand_manipulation => core/reorient}/inhand_manipulation_warp_env.py (99%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/a1/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/a1/flat_env_cfg.py (93%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/anymal_b/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/anymal_b/flat_env_cfg.py (93%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/anymal_c/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/anymal_c/flat_env_cfg.py (93%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/anymal_d/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/anymal_d/flat_env_cfg.py (93%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/cassie/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/cassie/flat_env_cfg.py (93%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/g1/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/g1/flat_env_cfg.py (93%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/go1/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/go1/flat_env_cfg.py (93%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/go2/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/go2/flat_env_cfg.py (93%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/h1/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/config/h1/flat_env_cfg.py (93%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/mdp/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/mdp/__init__.pyi (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/mdp/curriculums.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/mdp/rewards.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/{manager_based/locomotion => core}/velocity/mdp/terminations.py (100%) delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/__init__.py diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index 5509ec25a6f6..4fc8d100675b 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -163,10 +163,6 @@ def test_stable_cartpole_cfg_adapts_to_current_warp_module_layout(): cfg = _load_adapted_cfg(_cfg_entry_point("Isaac-Cartpole")) - assert cfg.rewards.pole_pos.func.__module__.startswith( - "isaaclab_tasks_experimental.manager_based.classic.cartpole.mdp" - ) - assert cfg.rewards.success_rate.func.__module__.startswith( - "isaaclab_tasks_experimental.manager_based.classic.cartpole.mdp" - ) + assert cfg.rewards.pole_pos.func.__module__.startswith("isaaclab_tasks_experimental.core.cartpole.mdp") + assert cfg.rewards.success_rate.func.__module__.startswith("isaaclab_tasks_experimental.core.cartpole.mdp") assert issubclass(cfg.actions.joint_effort.class_type, ActionTerm) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py index 45c6ca0572fa..cc8bff65c70a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py @@ -25,7 +25,7 @@ disable_env_checker=True, kwargs={ "env_cfg_entry_point": f"{__name__}.cartpole_direct_env_cfg:CartpoleEnvCfg", - "warp_entry_point": "isaaclab_tasks_experimental.direct.cartpole.cartpole_warp_env:CartpoleWarpEnv", + "warp_entry_point": "isaaclab_tasks_experimental.core.cartpole.cartpole_warp_env:CartpoleWarpEnv", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_direct_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpoleDirectPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_direct_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/__init__.py index ae86a5269bde..1e506e5ef4ae 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/__init__.py @@ -25,7 +25,7 @@ disable_env_checker=True, kwargs={ "env_cfg_entry_point": f"{__name__}.ant_direct_env_cfg:AntEnvCfg", - "warp_entry_point": "isaaclab_tasks_experimental.direct.ant.ant_env_warp:AntWarpEnv", + "warp_entry_point": "isaaclab_tasks_experimental.core.locomotion.ant.ant_env_warp:AntWarpEnv", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_direct_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AntDirectPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_direct_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/__init__.py index 0387402a7957..e3cee56589b2 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/__init__.py @@ -25,7 +25,7 @@ disable_env_checker=True, kwargs={ "env_cfg_entry_point": f"{__name__}.humanoid_direct_env_cfg:HumanoidEnvCfg", - "warp_entry_point": "isaaclab_tasks_experimental.direct.humanoid.humanoid_warp_env:HumanoidWarpEnv", + "warp_entry_point": "isaaclab_tasks_experimental.core.locomotion.humanoid.humanoid_warp_env:HumanoidWarpEnv", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_direct_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:HumanoidDirectPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_direct_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst index 782b6f25ef9a..425d89c1651b 100644 --- a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst +++ b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst @@ -6,6 +6,10 @@ Changed events that have no warp twins yet. The variants now require selecting the Newton solver on the CLI via ``presets=newton_mjwarp`` instead of hard-coding it in the configuration. +* **Breaking:** Changed the package layout to mirror :mod:`isaaclab_tasks.core` + (``isaaclab_tasks_experimental.core.``), replacing the previous + ``manager_based``/``direct`` split. Update imports of the old paths to the + ``core`` equivalents (e.g. ``isaaclab_tasks_experimental.core.cartpole.mdp``). Removed ^^^^^^^ diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/__init__.py similarity index 68% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/__init__.py index 3e2b7945ebde..bf495309161f 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/__init__.py @@ -3,8 +3,4 @@ # # SPDX-License-Identifier: BSD-3-Clause -""" -Direct workflow environments. -""" - -import gymnasium as gym +"""Warp-first task implementations mirroring the :mod:`isaaclab_tasks.core` layout.""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py similarity index 52% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py index e527f251fb5c..9a0eaa256a77 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py @@ -3,10 +3,12 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Warp MDP twins for the stable Cartpole task. +"""Warp implementations for the stable Cartpole tasks. -There is no separate task registration: run the stable ``Isaac-Cartpole`` task -with ``--frontend warp`` and ``presets=newton_mjwarp``. +Holds the MDP twins for the manager-based ``Isaac-Cartpole`` task and the +direct :class:`~isaaclab_tasks_experimental.core.cartpole.cartpole_warp_env.CartpoleWarpEnv` +declared by ``Isaac-Cartpole-Direct``. There is no separate task registration: +run the stable ids with ``--frontend warp`` and ``presets=newton_mjwarp``. """ from isaaclab_experimental.envs.frontend import register_mdp_route diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.pyi similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/__init__.pyi rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.pyi diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/__init__.py similarity index 79% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/__init__.py index 6cd56351b6e5..e37ca86717b5 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/__init__.py @@ -3,4 +3,4 @@ # # SPDX-License-Identifier: BSD-3-Clause -from .reach import * +"""Warp-first locomotion task implementations.""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/__init__.py similarity index 50% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/__init__.py index 0263ee66c216..277fea940e58 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/__init__.py @@ -3,10 +3,13 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Warp MDP route for the stable Ant task. +"""Warp implementations for the stable Ant tasks. -There is no separate task registration: run the stable ``Isaac-Ant`` task with -``--frontend warp`` and ``presets=newton_mjwarp``. +Holds the direct +:class:`~isaaclab_tasks_experimental.core.locomotion.ant.ant_env_warp.AntWarpEnv` +declared by ``Isaac-Ant-Direct``; the manager-based MDP twins are borrowed from +the Humanoid package. There is no separate task registration: run the stable +ids with ``--frontend warp`` and ``presets=newton_mjwarp``. """ from isaaclab_experimental.envs.frontend import register_mdp_route @@ -15,5 +18,5 @@ # experimental humanoid package. register_mdp_route( "isaaclab_tasks.core.locomotion.ant", - "isaaclab_tasks_experimental.manager_based.classic.humanoid.mdp", + "isaaclab_tasks_experimental.core.locomotion.humanoid.mdp", ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/ant_env_warp.py similarity index 83% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/ant_env_warp.py index 17ce3c2ee788..7d8957e4d9c7 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/ant_env_warp.py @@ -7,7 +7,7 @@ from isaaclab_tasks.core.locomotion.ant.ant_direct_env_cfg import AntEnvCfg -from isaaclab_tasks_experimental.direct.locomotion.locomotion_env_warp import LocomotionWarpEnv +from isaaclab_tasks_experimental.core.locomotion.locomotion_env_warp import LocomotionWarpEnv class AntWarpEnv(LocomotionWarpEnv): diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/__init__.py new file mode 100644 index 000000000000..7f05dee5608e --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/__init__.py @@ -0,0 +1,18 @@ +# 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 + +"""Warp implementations for the stable Humanoid tasks. + +Holds the MDP twins for the manager-based ``Isaac-Humanoid`` task (also +serving the Ant task, which borrows them) and the direct +:class:`~isaaclab_tasks_experimental.core.locomotion.humanoid.humanoid_warp_env.HumanoidWarpEnv` +declared by ``Isaac-Humanoid-Direct``. There is no separate task registration: +run the stable ids with ``--frontend warp`` and ``presets=newton_mjwarp``. +""" + +from isaaclab_experimental.envs.frontend import register_mdp_route + +# Warp twins for the stable humanoid MDP terms live in this package's ``mdp``. +register_mdp_route("isaaclab_tasks.core.locomotion.humanoid", f"{__name__}.mdp") diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/humanoid_warp_env.py similarity index 84% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/humanoid_warp_env.py index 113d761b711c..2a843510717a 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/humanoid_warp_env.py @@ -7,7 +7,7 @@ from isaaclab_tasks.core.locomotion.humanoid.humanoid_direct_env_cfg import HumanoidEnvCfg -from isaaclab_tasks_experimental.direct.locomotion.locomotion_env_warp import LocomotionWarpEnv +from isaaclab_tasks_experimental.core.locomotion.locomotion_env_warp import LocomotionWarpEnv class HumanoidWarpEnv(LocomotionWarpEnv): diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/__init__.pyi similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/__init__.pyi rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/__init__.pyi diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/observations.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/observations.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/observations.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/observations.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/rewards.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/rewards.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.pyi similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.pyi rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.pyi diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/__init__.py similarity index 75% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/__init__.py index 460a30569089..48bb97b31113 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/__init__.py @@ -2,3 +2,5 @@ # All rights reserved. # # SPDX-License-Identifier: BSD-3-Clause + +"""Warp-first in-hand reorientation task implementations.""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/__init__.py similarity index 77% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/__init__.py index 460a30569089..771b55fc3605 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/__init__.py @@ -2,3 +2,5 @@ # All rights reserved. # # SPDX-License-Identifier: BSD-3-Clause + +"""Warp task registrations for reorientation robots.""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/__init__.py similarity index 92% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/__init__.py index e7fc129c633f..13ab818b132d 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/__init__.py @@ -13,7 +13,7 @@ # Register Gym environments. ## -inhand_task_entry = "isaaclab_tasks_experimental.direct.inhand_manipulation" +inhand_task_entry = "isaaclab_tasks_experimental.core.reorient" stable_agents = "isaaclab_tasks.core.reorient.config.allegro_hand.agents" gym.register( diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/allegro_hand_warp_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/allegro_hand_warp_env_cfg.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/allegro_hand_warp_env_cfg.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/allegro_hand_warp_env_cfg.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/inhand_manipulation_warp_env.py similarity index 99% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/inhand_manipulation_warp_env.py index 2082d470ff9b..ae8a8b30222c 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/inhand_manipulation_warp_env.py @@ -20,7 +20,9 @@ from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane if TYPE_CHECKING: - from isaaclab_tasks_experimental.direct.allegro_hand.allegro_hand_warp_env_cfg import AllegroHandWarpEnvCfg + from isaaclab_tasks_experimental.core.reorient.config.allegro_hand.allegro_hand_warp_env_cfg import ( + AllegroHandWarpEnvCfg, + ) @wp.kernel diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/flat_env_cfg.py similarity index 93% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/flat_env_cfg.py index d58290799424..c9948d38a278 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/flat_env_cfg.py @@ -12,7 +12,7 @@ UnitreeA1FlatEnvCfg_PLAY as _StableUnitreeA1FlatEnvCfg_PLAY, ) -from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( +from isaaclab_tasks_experimental.core.velocity import ( disable_unsupported_randomization_events, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/flat_env_cfg.py similarity index 93% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/flat_env_cfg.py index 900d87e64227..0c01243aa7d0 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/flat_env_cfg.py @@ -12,7 +12,7 @@ AnymalBFlatEnvCfg_PLAY as _StableAnymalBFlatEnvCfg_PLAY, ) -from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( +from isaaclab_tasks_experimental.core.velocity import ( disable_unsupported_randomization_events, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/flat_env_cfg.py similarity index 93% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/flat_env_cfg.py index edb8ffd7a90e..7777d8274e6a 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/flat_env_cfg.py @@ -12,7 +12,7 @@ AnymalCFlatEnvCfg_PLAY as _StableAnymalCFlatEnvCfg_PLAY, ) -from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( +from isaaclab_tasks_experimental.core.velocity import ( disable_unsupported_randomization_events, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/flat_env_cfg.py similarity index 93% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/flat_env_cfg.py index ca010dffb13c..ccd874878e2f 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/flat_env_cfg.py @@ -12,7 +12,7 @@ AnymalDFlatEnvCfg_PLAY as _StableAnymalDFlatEnvCfg_PLAY, ) -from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( +from isaaclab_tasks_experimental.core.velocity import ( disable_unsupported_randomization_events, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/flat_env_cfg.py similarity index 93% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/flat_env_cfg.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/flat_env_cfg.py index 4b42afa8a3a5..37696377d675 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/flat_env_cfg.py @@ -12,7 +12,7 @@ CassieFlatEnvCfg_PLAY as _StableCassieFlatEnvCfg_PLAY, ) -from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( +from isaaclab_tasks_experimental.core.velocity import ( disable_unsupported_randomization_events, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/flat_env_cfg.py similarity index 93% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/flat_env_cfg.py index 9c2011fe13c8..e11424804874 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/flat_env_cfg.py @@ -10,7 +10,7 @@ from isaaclab_tasks.core.velocity.config.g1.flat_env_cfg import G1FlatEnvCfg as _StableG1FlatEnvCfg from isaaclab_tasks.core.velocity.config.g1.flat_env_cfg import G1FlatEnvCfg_PLAY as _StableG1FlatEnvCfg_PLAY -from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( +from isaaclab_tasks_experimental.core.velocity import ( disable_unsupported_randomization_events, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/flat_env_cfg.py similarity index 93% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/flat_env_cfg.py index c54355196dab..bbe78df0bd1b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/flat_env_cfg.py @@ -12,7 +12,7 @@ UnitreeGo1FlatEnvCfg_PLAY as _StableUnitreeGo1FlatEnvCfg_PLAY, ) -from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( +from isaaclab_tasks_experimental.core.velocity import ( disable_unsupported_randomization_events, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/flat_env_cfg.py similarity index 93% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/flat_env_cfg.py index a54d94dabcc9..87eb63e6e2fa 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/flat_env_cfg.py @@ -12,7 +12,7 @@ UnitreeGo2FlatEnvCfg_PLAY as _StableUnitreeGo2FlatEnvCfg_PLAY, ) -from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( +from isaaclab_tasks_experimental.core.velocity import ( disable_unsupported_randomization_events, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/flat_env_cfg.py similarity index 93% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/flat_env_cfg.py index 088d5786b223..6c7ca5701261 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/flat_env_cfg.py @@ -10,7 +10,7 @@ from isaaclab_tasks.core.velocity.config.h1.flat_env_cfg import H1FlatEnvCfg as _StableH1FlatEnvCfg from isaaclab_tasks.core.velocity.config.h1.flat_env_cfg import H1FlatEnvCfg_PLAY as _StableH1FlatEnvCfg_PLAY -from isaaclab_tasks_experimental.manager_based.locomotion.velocity import ( +from isaaclab_tasks_experimental.core.velocity import ( disable_unsupported_randomization_events, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.pyi rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/terminations.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/terminations.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/__init__.py deleted file mode 100644 index fd466b5394a3..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# 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 - -"""Warp implementation of the stable direct Ant task. - -There is no separate task registration: the stable ``Isaac-Ant-Direct`` -registration declares :class:`~isaaclab_tasks_experimental.direct.ant.ant_env_warp.AntWarpEnv` -as its ``warp_entry_point``; run it with ``--frontend warp`` and -``presets=newton_mjwarp``. -""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py deleted file mode 100644 index 8dea234be2f8..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# 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 - -"""Warp implementation of the stable direct Cartpole task. - -There is no separate task registration: the stable ``Isaac-Cartpole-Direct`` -registration declares :class:`~isaaclab_tasks_experimental.direct.cartpole.cartpole_warp_env.CartpoleWarpEnv` -as its ``warp_entry_point``; run it with ``--frontend warp`` and -``presets=newton_mjwarp``. -""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/__init__.py deleted file mode 100644 index 663e41a7c283..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# 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 - -"""Warp implementation of the stable direct Humanoid task. - -There is no separate task registration: the stable ``Isaac-Humanoid-Direct`` -registration declares :class:`~isaaclab_tasks_experimental.direct.humanoid.humanoid_warp_env.HumanoidWarpEnv` -as its ``warp_entry_point``; run it with ``--frontend warp`` and -``presets=newton_mjwarp``. -""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/__init__.py deleted file mode 100644 index 7f23883e6332..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# 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 - -"""Experimental registrations for manager-based tasks. - -We intentionally only register new Gym IDs pointing at experimental entry points. -Task definitions (configs/mdp) remain in `isaaclab_tasks` to avoid duplication. -""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/__init__.py deleted file mode 100644 index 79c13e2aa8f4..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# 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 - -"""Classic environments for control. - -These environments are based on the MuJoCo environments provided by OpenAI. - -Reference: - https://github.com/openai/gym/tree/master/gym/envs/mujoco -""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py deleted file mode 100644 index 0d31a47ece94..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -# 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 - -"""Warp MDP twins for the stable Humanoid task. - -There is no separate task registration: run the stable ``Isaac-Humanoid`` task -with ``--frontend warp`` and ``presets=newton_mjwarp``. -""" - -from isaaclab_experimental.envs.frontend import register_mdp_route - -# Warp twins for the stable humanoid MDP terms live in this package's ``mdp``. -register_mdp_route("isaaclab_tasks.core.locomotion.humanoid", f"{__name__}.mdp") diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/__init__.py deleted file mode 100644 index 0660d38f0658..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# 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 - -"""Locomotion experimental task registrations (manager-based).""" From 5129070748eeca32dc9a727e411f9639e5c5f38b Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 16 Jul 2026 00:08:41 -0700 Subject: [PATCH 30/58] Encapsulate the warp frontend machinery in a class MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Group the warp-side machinery — route registry, cfg adaptation steps, twin resolution, and both env build paths — under a WarpFrontend class instead of a flat set of module functions. The module-level build(), adapt_cfg_for_warp(), and register_mdp_route() API is unchanged for callers; the class gives the machinery one named owner, keeps the route registry as explicit class state, and leaves room for future frontends to subclass or replace it. --- .../isaaclab_experimental/envs/frontend.py | 956 ++++++++++-------- .../test/envs/test_frontend.py | 85 +- 2 files changed, 554 insertions(+), 487 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index 72e892aebd60..ca5174c479f5 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -5,17 +5,20 @@ """Runtime selector for IsaacLab tasks (``--frontend {torch,warp}``). -Two responsibilities only: - -1. Decide which runtime constructs the env (torch via ``gym.make``, or warp via - :class:`ManagerBasedRLEnvWarp`). -2. For the warp runtime, adapt a stable manager-based cfg in place so the - warp managers can consume it (Newton physics + warp-twin term funcs and - action classes + warp variant of :class:`SceneEntityCfg`). - -The adaptation is a single function (:func:`_adapt_cfg_for_warp`); there is no -plugin framework. All warp MDP twins (functions *and* classes) must exist for -the chosen task — a missing twin is a hard failure, not a silent drop. +:func:`build` decides which runtime constructs the env: torch via ``gym.make``, +or warp via :class:`WarpFrontend`, which encapsulates everything the warp +runtime needs from a stable task definition: + +* Manager-based tasks: adapt the stable cfg in place (Newton physics check, + :class:`SceneEntityCfg` promotion, MDP twin swap) and construct + :class:`ManagerBasedRLEnvWarp`. All warp MDP twins (functions *and* classes) + must exist for the chosen task — a missing twin is a hard failure, not a + silent drop. +* Direct tasks: construct the warp env class the stable registration declares + via its ``warp_entry_point`` kwarg, sharing the stable cfg. + +Task packages declare where their warp MDP twins live with +:func:`register_mdp_route`; the frontend holds no per-task table. """ from __future__ import annotations @@ -25,7 +28,7 @@ from collections.abc import Iterable, Iterator from enum import StrEnum from types import ModuleType -from typing import Any +from typing import Any, ClassVar import gymnasium as gym @@ -38,6 +41,7 @@ __all__ = [ "Frontend", "FrontendIncompatibleError", + "WarpFrontend", "Workflow", "adapt_cfg_for_warp", "build", @@ -69,74 +73,10 @@ class FrontendIncompatibleError(RuntimeError): # --------------------------------------------------------------------------- -# Public entry point +# Public entry points # --------------------------------------------------------------------------- -# Prefixes used to decide whether a registered task's entry-point lives under -# the warp packages. Used by the direct-workflow guard. -_WARP_ROOT_PREFIXES: tuple[str, ...] = ("isaaclab_experimental", "isaaclab_tasks_experimental") - -# Stable package prefixes mapped to the warp MDP module that twins their terms. -# The stable task package is organized by public task domain while the -# experimental package separates manager-based workflows by family, so the -# module trees are not one-to-one. The registry is populated by the experimental -# task packages themselves via :func:`register_mdp_route` at import time, so -# adding a new task family never requires editing this module. -_WARP_MDP_MODULE_ROUTES: dict[str, str] = {} - -# Root package that provides task-specific warp MDP twins. Imported (lazily, at -# swap time) so its task packages run their register_mdp_route() calls even when -# the caller only imported the stable task package. -_TWIN_PROVIDER_PACKAGE = "isaaclab_tasks_experimental" - - -def register_mdp_route(stable_package: str, warp_mdp_module: str) -> None: - """Register the warp MDP module that twins a stable package's terms. - - Experimental task packages call this at import time to declare where the - warp twins of a stable task's MDP terms live. :func:`_swap_mdp` consults - the registry to resolve task-specific twins for ``--frontend=warp``. Twin - lookup is by symbol name on ``warp_mdp_module`` itself, mirroring how a - stable cfg consumes its ``mdp`` namespace — so the module should re-export - everything the task needs (task-specific twins plus the generic forwards). - - Two kinds of keys are useful: - - * A *task package* (e.g. ``"isaaclab_tasks.core.reach"``) routes every cfg - class defined under it. This also covers terms whose stable functions are - defined in core/shared packages but overridden by task-specific twins. - * An *MDP package* (e.g. ``"isaaclab_tasks.core.locomotion.humanoid.mdp"``) - routes symbols *defined* under it, which serves other tasks that borrow - those terms. - - Args: - stable_package: Stable package prefix to route (longest match wins). - warp_mdp_module: Warp MDP module providing the twins, e.g. - ``"isaaclab_tasks_experimental.manager_based.classic.cartpole.mdp"``. - - Raises: - ValueError: If ``stable_package`` is already routed to a different module. - """ - existing = _WARP_MDP_MODULE_ROUTES.get(stable_package) - if existing is not None and existing != warp_mdp_module: - raise ValueError( - f"MDP route conflict for {stable_package!r}: already routed to {existing!r}, got {warp_mdp_module!r}." - ) - _WARP_MDP_MODULE_ROUTES[stable_package] = warp_mdp_module - - -# Top-level cfg groups whose managers run warp-first. Only terms under these are -# adapted (SceneEntityCfg promotion + MDP twin swap). The event manager is -# warp-first too — it invokes term funcs with a Warp env-mask, so a stable event -# func (which expects torch ``env_ids``) breaks at runtime; its funcs must be -# swapped to warp twins. The curriculum, recorder and command managers run on the -# stable (torch) implementation, so their terms are left untouched. A stable term -# left on a warp manager would break, so a missing twin in these groups is a hard -# error; a stable term on a stable manager is correct, so those groups are skipped. -_WARP_MANAGED_GROUPS: frozenset[str] = frozenset({"observations", "rewards", "terminations", "actions", "events"}) - - def build( frontend: Frontend | str, env_cfg: Any, @@ -146,10 +86,10 @@ def build( """Construct the env on the selected runtime. Args: - frontend: ``"torch"`` (default IsaacLab path) or ``"warp"`` (warp managers - + :class:`ManagerBasedRLEnvWarp` for manager-based tasks). + frontend: ``"torch"`` (default IsaacLab path) or ``"warp"`` + (:class:`WarpFrontend`). env_cfg: Stable env cfg. Mutated in place when ``frontend == "warp"``. - task_id: Gym registration id, e.g. ``"Isaac-Cartpole-v0"``. + task_id: Gym registration id, e.g. ``"Isaac-Cartpole"``. **construct_kwargs: Forwarded to the env constructor (``render_mode``, …). Returns: @@ -157,400 +97,524 @@ def build( Raises: FrontendIncompatibleError: If the warp runtime can't run the task - (wrong physics, missing MDP twins, direct task not registered - as a warp env). + (wrong physics, missing MDP twins, direct task without a warp + implementation). """ frontend = Frontend(frontend) if frontend is Frontend.TORCH: return gym.make(task_id, cfg=env_cfg, **construct_kwargs) + return WarpFrontend.build_env(env_cfg, task_id, **construct_kwargs) - workflow = _detect_workflow(env_cfg) - if workflow is Workflow.DIRECT: - # Direct workflows aren't cfg-adapted: a hand-written warp env class - # implements the task. Preferred path: the stable registration declares - # it via ``warp_entry_point`` and shares the stable cfg. Fallback: the - # task itself is registered under the warp packages (warp-native cfg). - env_class = _resolve_direct_warp_class(task_id) - if env_class is None: - _assert_direct_warp_registration(task_id) - return gym.make(task_id, cfg=env_cfg, **construct_kwargs) - _require_newton_physics(env_cfg, type(env_cfg).__name__) - return env_class(cfg=env_cfg, **construct_kwargs) - # Imported lazily so that ``--frontend=torch`` callers don't pay the - # ``isaaclab_experimental.envs`` import cost. The env adapts the cfg itself - # in ``__init__`` (see :func:`adapt_cfg_for_warp`), so stable cfgs and - # warp-native cfgs take the same path. - from isaaclab_experimental.envs import ManagerBasedRLEnvWarp +def adapt_cfg_for_warp(cfg: Any) -> None: + """Mutate a stable manager-based ``cfg`` in place so warp managers can consume it. - return ManagerBasedRLEnvWarp(cfg=env_cfg, **construct_kwargs) + Functional alias for :meth:`WarpFrontend.adapt_cfg`, used by + :class:`ManagerBasedRLEnvWarp` in its ``__init__``. + """ + WarpFrontend.adapt_cfg(cfg) # --------------------------------------------------------------------------- -# Cfg adaptation (warp only) +# Warp frontend # --------------------------------------------------------------------------- -def adapt_cfg_for_warp(cfg: Any) -> None: - """Mutate a stable manager-based ``cfg`` in place so warp managers can consume it. +class WarpFrontend: + """Builds environments on the warp runtime from stable task definitions. - Called by :func:`build` for ``--frontend=warp`` on a stable manager-based - task. Idempotent: re-running on an already-adapted cfg is a no-op (the steps - below skip Warp-origin symbols and already-promoted entities). - - Three steps, each independently testable: - - 1. :func:`_require_newton_physics` — hard check that ``cfg.sim.physics`` is - :class:`~isaaclab_newton.physics.NewtonCfg`. The user is responsible for - selecting the Newton variant of the task's :class:`PresetCfg` via - ``presets=newton_mjwarp``; we don't auto-inject. - 2. :func:`_promote_scene_entity_cfgs` — replace stable - :class:`~isaaclab.managers.SceneEntityCfg` instances under each term's - ``params`` with the warp variant (which adds warp-cached ``joint_mask``, - ``joint_ids_wp``, ``body_ids_wp`` fields). - 3. :func:`_swap_mdp` — for every MDP term found anywhere in the cfg tree - (discovered by :func:`_walk_terms` via :class:`ManagerTermBaseCfg` - subclassing, not by hard-coded attribute names), replace any stable - ``func`` *or* ``class_type`` with its same-named warp twin. A missing - twin raises :class:`FrontendIncompatibleError` — partial coverage is - unsafe under the warp managers' kernel-only signature. + The class is a stateless namespace apart from the MDP route registry + (:attr:`_mdp_routes`), which experimental task packages populate at import + time via :meth:`register_mdp_route`. The stable and experimental task trees + mirror each other (``isaaclab_tasks.core.X`` ↔ + ``isaaclab_tasks_experimental.core.X``), so most routes are the mechanical + mirror; the registry keeps routing explicit and accommodates exceptions + (the Ant task borrows the Humanoid twins). """ - label = type(cfg).__name__ - _require_newton_physics(cfg, label) - _promote_scene_entity_cfgs(cfg) - _swap_mdp(cfg, label) - -def _require_newton_physics(cfg: Any, label: str) -> None: - """Block unless ``cfg.sim.physics`` is :class:`NewtonCfg`. - - The warp managers' assets read state through :class:`NewtonManager`; - a :class:`PhysxCfg` (or unresolved :class:`PresetCfg`) is a hard - incompatibility. The fix is to pass ``presets=newton_mjwarp`` on the CLI so - Hydra resolves the task's :class:`PresetCfg` wrapper to the Newton field - before construction. - """ - from isaaclab_newton.physics import NewtonCfg - - physics = getattr(getattr(cfg, "sim", None), "physics", None) - if isinstance(physics, NewtonCfg): - return - raise FrontendIncompatibleError( - f"warp env {label!r}: expected cfg.sim.physics to be NewtonCfg," - f" got {type(physics).__name__!r}. Pass `presets=newton_mjwarp` on the CLI so" - f" Hydra resolves the task's PresetCfg wrapper to the Newton variant." + # Prefixes that identify warp-origin symbols and registrations. + WARP_ROOT_PREFIXES: ClassVar[tuple[str, ...]] = ("isaaclab_experimental", "isaaclab_tasks_experimental") + + # Top-level cfg groups whose managers run warp-first. Only terms under these + # are adapted (SceneEntityCfg promotion + MDP twin swap). The event manager + # is warp-first too — it invokes term funcs with a Warp env-mask, so a stable + # event func (which expects torch ``env_ids``) breaks at runtime; its funcs + # must be swapped to warp twins. The curriculum, recorder and command + # managers run on the stable (torch) implementation, so their terms are left + # untouched. A stable term left on a warp manager would break, so a missing + # twin in these groups is a hard error; a stable term on a stable manager is + # correct, so those groups are skipped. + WARP_MANAGED_GROUPS: ClassVar[frozenset[str]] = frozenset( + {"observations", "rewards", "terminations", "actions", "events"} ) + # Root package that provides task-specific warp MDP twins. Imported (lazily, + # at swap time) so its task packages run their register_mdp_route() calls + # even when the caller only imported the stable task package. + _TWIN_PROVIDER_PACKAGE: ClassVar[str] = "isaaclab_tasks_experimental" + + # Shared fallback module holding generic warp twins. + _FALLBACK_MDP_MODULE: ClassVar[str] = "isaaclab_experimental.envs.mdp" + + # Stable package prefixes mapped to the warp MDP module that twins their + # terms. Populated via register_mdp_route(); never edited here. + _mdp_routes: ClassVar[dict[str, str]] = {} + + # ------------------------------------------------------------------ + # Route registry + # ------------------------------------------------------------------ + + @classmethod + def register_mdp_route(cls, stable_package: str, warp_mdp_module: str) -> None: + """Register the warp MDP module that twins a stable package's terms. + + Experimental task packages call this at import time to declare where + the warp twins of a stable task's MDP terms live. Twin lookup is by + symbol name on ``warp_mdp_module`` itself, mirroring how a stable cfg + consumes its ``mdp`` namespace — so the module should re-export + everything the task needs (task-specific twins plus the generic + forwards). + + Two kinds of keys are useful: + + * A *task package* (e.g. ``"isaaclab_tasks.core.reach"``) routes every + cfg class defined under it. This also covers terms whose stable + functions are defined in core/shared packages but overridden by + task-specific twins. + * An *MDP package* (e.g. ``"isaaclab_tasks.core.locomotion.humanoid.mdp"``) + routes symbols *defined* under it, which serves other tasks that + borrow those terms. + + Args: + stable_package: Stable package prefix to route (longest match wins). + warp_mdp_module: Warp MDP module providing the twins, e.g. + ``"isaaclab_tasks_experimental.core.cartpole.mdp"``. + + Raises: + ValueError: If ``stable_package`` is already routed to a different module. + """ + existing = cls._mdp_routes.get(stable_package) + if existing is not None and existing != warp_mdp_module: + raise ValueError( + f"MDP route conflict for {stable_package!r}: already routed to {existing!r}, got {warp_mdp_module!r}." + ) + cls._mdp_routes[stable_package] = warp_mdp_module + + # ------------------------------------------------------------------ + # Env construction + # ------------------------------------------------------------------ + + @classmethod + def build_env(cls, env_cfg: Any, task_id: str, **construct_kwargs: Any) -> gym.Env: + """Construct the warp env for ``task_id`` from a stable env cfg.""" + if cls._detect_workflow(env_cfg) is Workflow.DIRECT: + return cls._build_direct_env(env_cfg, task_id, **construct_kwargs) + return cls._build_manager_env(env_cfg, **construct_kwargs) + + @classmethod + def _build_manager_env(cls, env_cfg: Any, **construct_kwargs: Any) -> gym.Env: + """Construct :class:`ManagerBasedRLEnvWarp`; the env adapts the cfg itself. + + Imported lazily so that ``--frontend=torch`` callers don't pay the + ``isaaclab_experimental.envs`` import cost. The env runs + :meth:`adapt_cfg` in ``__init__``, so stable cfgs and warp-native cfgs + take the same path. + """ + from isaaclab_experimental.envs import ManagerBasedRLEnvWarp + + return ManagerBasedRLEnvWarp(cfg=env_cfg, **construct_kwargs) + + @classmethod + def _build_direct_env(cls, env_cfg: Any, task_id: str, **construct_kwargs: Any) -> gym.Env: + """Construct a direct warp env. + + Direct workflows aren't cfg-adapted: a hand-written warp env class + implements the task. Preferred path: the stable registration declares + it via ``warp_entry_point`` and shares the stable cfg. Fallback: the + task itself is registered under the warp packages (warp-native cfg). + """ + env_class = cls._resolve_direct_warp_class(task_id) + if env_class is None: + cls._assert_direct_warp_registration(task_id) + return gym.make(task_id, cfg=env_cfg, **construct_kwargs) + cls._require_newton_physics(env_cfg, type(env_cfg).__name__) + return env_class(cfg=env_cfg, **construct_kwargs) -def _promote_scene_entity_cfgs(cfg: Any) -> None: - """Replace stable :class:`SceneEntityCfg` instances with the warp variant. - - Iterates every term cfg in the tree (via :func:`_walk_terms`) and rebuilds - any stable :class:`SceneEntityCfg` value under ``term.params`` through - :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable`. The - warp variant subclasses the stable one, so type checks elsewhere stay - valid; the new fields (``joint_mask`` / ``joint_ids_wp`` / ``body_ids_wp``) - are filled at :meth:`resolve` time by the warp scene. - """ - from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as _WarpSceneEntityCfg - - promoted: list[str] = [] - for path, term in _walk_terms(cfg): - if not path or path[0] not in _WARP_MANAGED_GROUPS: - continue - params = getattr(term, "params", None) - if not isinstance(params, dict): - continue - for key, value in list(params.items()): - if isinstance(value, _WarpSceneEntityCfg) or not isinstance(value, _StableSceneEntityCfg): - continue - params[key] = _WarpSceneEntityCfg.from_stable(value) - promoted.append(f"{'.'.join(path)}.params[{key!r}] ({value.name!r})") - if promoted: - logger.info( - "frontend.warp: promoted %d SceneEntityCfg instance(s) to warp variant:\n %s", - len(promoted), - "\n ".join(promoted), + # ------------------------------------------------------------------ + # Cfg adaptation (manager-based) + # ------------------------------------------------------------------ + + @classmethod + def adapt_cfg(cls, cfg: Any) -> None: + """Mutate a stable manager-based ``cfg`` in place for the warp managers. + + Idempotent: re-running on an already-adapted cfg is a no-op (the steps + below skip warp-origin symbols and already-promoted entities). + + Three steps, each independently testable: + + 1. :meth:`_require_newton_physics` — hard check that ``cfg.sim.physics`` + is :class:`~isaaclab_newton.physics.NewtonCfg`. The user is + responsible for selecting the Newton variant of the task's + :class:`PresetCfg` via ``presets=newton_mjwarp``; we don't auto-inject. + 2. :meth:`_promote_scene_entity_cfgs` — replace stable + :class:`~isaaclab.managers.SceneEntityCfg` instances under each + term's ``params`` with the warp variant (which adds warp-cached + ``joint_mask``, ``joint_ids_wp``, ``body_ids_wp`` fields). + 3. :meth:`_swap_mdp` — for every MDP term found anywhere in the cfg tree + (discovered by :meth:`_walk_terms` via :class:`ManagerTermBaseCfg` + subclassing, not by hard-coded attribute names), replace any stable + ``func`` *or* ``class_type`` with its same-named warp twin. A missing + twin raises :class:`FrontendIncompatibleError` — partial coverage is + unsafe under the warp managers' kernel-only signature. + """ + label = type(cfg).__name__ + cls._require_newton_physics(cfg, label) + cls._promote_scene_entity_cfgs(cfg) + cls._swap_mdp(cfg, label) + + @staticmethod + def _require_newton_physics(cfg: Any, label: str) -> None: + """Block unless ``cfg.sim.physics`` is :class:`NewtonCfg`. + + The warp managers' assets read state through :class:`NewtonManager`; + a :class:`PhysxCfg` (or unresolved :class:`PresetCfg`) is a hard + incompatibility. The fix is to pass ``presets=newton_mjwarp`` on the CLI + so Hydra resolves the task's :class:`PresetCfg` wrapper to the Newton + field before construction. + """ + from isaaclab_newton.physics import NewtonCfg + + physics = getattr(getattr(cfg, "sim", None), "physics", None) + if isinstance(physics, NewtonCfg): + return + raise FrontendIncompatibleError( + f"warp env {label!r}: expected cfg.sim.physics to be NewtonCfg," + f" got {type(physics).__name__!r}. Pass `presets=newton_mjwarp` on the CLI so" + f" Hydra resolves the task's PresetCfg wrapper to the Newton variant." ) - -def _swap_mdp(cfg: Any, label: str) -> None: - """Replace ``term.func`` and ``term.class_type`` with their warp twins. - - Iterates every term cfg in the tree (via :func:`_walk_terms`) and on each - term swaps whichever of ``func`` / ``class_type`` is a stable-origin - callable. Twin lookup is name-based, in the order given by - :func:`_twin_modules`: the warp mirror of the cfg's own task MDP namespace - first (task twins win, even for symbols defined in core packages), then the - mirror of the symbol's defining package (covers terms borrowed from another - task family, e.g. the Ant task reusing Humanoid MDP terms), then the shared - :mod:`isaaclab_experimental.envs.mdp` fallback. Any missing twin raises - :class:`FrontendIncompatibleError` listing every affected term — partial - swaps would leave torch-style callables in the cfg and the warp managers - would call them with the wrong signature. - - The warp-side declarations (``out_dim``, ``axes``, ``observation_type``) - that the warp managers need at init are *not* supplied by this swap; they - travel with the warp twin function itself via its own - ``@generic_io_descriptor_warp(out_dim=…)`` decorator. This function only - substitutes the callable; the manager reads the new func's annotations - when it parses the term cfg. - """ - _ensure_twin_providers_imported() - - cfg_route_modules = _cfg_route_modules(cfg) - module_cache: dict[str, list[ModuleType]] = {} - searched: set[str] = set() - - swapped = 0 - missing: list[tuple[str, str, str]] = [] # (location, attr, symbol) - for path, term in _walk_terms(cfg): - if not path or path[0] not in _WARP_MANAGED_GROUPS: - continue - location = ".".join(path) - for attr in ("func", "class_type"): - stable = getattr(term, attr, None) - if stable is None or not _is_swap_candidate(stable): + @classmethod + def _promote_scene_entity_cfgs(cls, cfg: Any) -> None: + """Replace stable :class:`SceneEntityCfg` instances with the warp variant. + + Iterates every term cfg in the tree (via :meth:`_walk_terms`) and + rebuilds any stable :class:`SceneEntityCfg` value under ``term.params`` + through :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable`. + The warp variant subclasses the stable one, so type checks elsewhere + stay valid; the new fields (``joint_mask`` / ``joint_ids_wp`` / + ``body_ids_wp``) are filled at :meth:`resolve` time by the warp scene. + """ + from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as _WarpSceneEntityCfg + + promoted: list[str] = [] + for path, term in cls._walk_terms(cfg): + if not path or path[0] not in cls.WARP_MANAGED_GROUPS: continue - origin = getattr(stable, "__module__", "") or "" - if origin not in module_cache: - module_cache[origin] = _twin_modules(origin, cfg_route_modules) - searched.update(m.__name__ for m in module_cache[origin]) - twin = _resolve_warp_twin(stable.__name__, module_cache[origin]) - if twin is None: - missing.append((location, attr, stable.__name__)) + params = getattr(term, "params", None) + if not isinstance(params, dict): continue - setattr(term, attr, twin) - swapped += 1 - - if missing: - lines = "\n ".join(f"{loc}.{attr}: no warp twin for {sym!r}" for loc, attr, sym in missing) + for key, value in list(params.items()): + if isinstance(value, _WarpSceneEntityCfg) or not isinstance(value, _StableSceneEntityCfg): + continue + params[key] = _WarpSceneEntityCfg.from_stable(value) + promoted.append(f"{'.'.join(path)}.params[{key!r}] ({value.name!r})") + if promoted: + logger.info( + "frontend.warp: promoted %d SceneEntityCfg instance(s) to warp variant:\n %s", + len(promoted), + "\n ".join(promoted), + ) + + @classmethod + def _swap_mdp(cls, cfg: Any, label: str) -> None: + """Replace ``term.func`` and ``term.class_type`` with their warp twins. + + Iterates every term cfg in the tree (via :meth:`_walk_terms`) and on + each term swaps whichever of ``func`` / ``class_type`` is a + stable-origin callable. Twin lookup is name-based, in the order given by + :meth:`_twin_modules`: the warp mirror of the cfg's own task MDP + namespace first (task twins win, even for symbols defined in core + packages), then the mirror of the symbol's defining package (covers + terms borrowed from another task family, e.g. the Ant task reusing + Humanoid MDP terms), then the shared fallback. Any missing twin raises + :class:`FrontendIncompatibleError` listing every affected term — + partial swaps would leave torch-style callables in the cfg and the warp + managers would call them with the wrong signature. + + The warp-side declarations (``out_dim``, ``axes``, ``observation_type``) + that the warp managers need at init are *not* supplied by this swap; + they travel with the warp twin function itself via its own + ``@generic_io_descriptor_warp(out_dim=…)`` decorator. This method only + substitutes the callable; the manager reads the new func's annotations + when it parses the term cfg. + """ + cls._ensure_twin_providers_imported() + + cfg_route_modules = cls._cfg_route_modules(cfg) + module_cache: dict[str, list[ModuleType]] = {} + searched: set[str] = set() + + swapped = 0 + missing: list[tuple[str, str, str]] = [] # (location, attr, symbol) + for path, term in cls._walk_terms(cfg): + if not path or path[0] not in cls.WARP_MANAGED_GROUPS: + continue + location = ".".join(path) + for attr in ("func", "class_type"): + stable = getattr(term, attr, None) + if stable is None or not cls._is_swap_candidate(stable): + continue + origin = getattr(stable, "__module__", "") or "" + if origin not in module_cache: + module_cache[origin] = cls._twin_modules(origin, cfg_route_modules) + searched.update(m.__name__ for m in module_cache[origin]) + twin = cls._resolve_warp_twin(stable.__name__, module_cache[origin]) + if twin is None: + missing.append((location, attr, stable.__name__)) + continue + setattr(term, attr, twin) + swapped += 1 + + if missing: + lines = "\n ".join(f"{loc}.{attr}: no warp twin for {sym!r}" for loc, attr, sym in missing) + raise FrontendIncompatibleError( + f"warp env {label!r}: missing warp MDP twins (searched {sorted(searched)}):\n {lines}" + ) + + logger.info("frontend.warp: swapped %d MDP symbol(s) to warp twins", swapped) + + # ------------------------------------------------------------------ + # Workflow detection and direct-task resolution + # ------------------------------------------------------------------ + + @staticmethod + def _detect_workflow(cfg: Any) -> Workflow: + """Classify the env cfg into manager-based or direct (used to pick build path). + + Note: + The env cfg roots (ManagerBasedRLEnvCfg, DirectRLEnvCfg, + DirectMARLEnvCfg) do not share a common base class. When a new cfg + root is added, extend the isinstance tuples below. + """ + if isinstance(cfg, ManagerBasedRLEnvCfg): + return Workflow.MANAGER_BASED + if isinstance(cfg, (DirectRLEnvCfg, DirectMARLEnvCfg)): + return Workflow.DIRECT raise FrontendIncompatibleError( - f"warp env {label!r}: missing warp MDP twins (searched {sorted(searched)}):\n {lines}" + f"Unrecognised env cfg type {type(cfg).__name__!r};" + f" expected ManagerBasedRLEnvCfg / DirectRLEnvCfg / DirectMARLEnvCfg subclass." ) - logger.info("frontend.warp: swapped %d MDP symbol(s) to warp twins", swapped) - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _detect_workflow(cfg: Any) -> Workflow: - """Classify the env cfg into manager-based or direct (used to pick build path). - - Note: - The four env cfg roots (ManagerBasedEnvCfg, ManagerBasedRLEnvCfg, - DirectRLEnvCfg, DirectMARLEnvCfg) do not share a common base class. - When a new cfg root is added, extend the isinstance tuples below. - """ - if isinstance(cfg, ManagerBasedRLEnvCfg): - return Workflow.MANAGER_BASED - if isinstance(cfg, (DirectRLEnvCfg, DirectMARLEnvCfg)): - return Workflow.DIRECT - raise FrontendIncompatibleError( - f"Unrecognised env cfg type {type(cfg).__name__!r};" - f" expected ManagerBasedRLEnvCfg / DirectRLEnvCfg / DirectMARLEnvCfg subclass." - ) - - -def _resolve_direct_warp_class(task_id: str) -> type | None: - """Return the warp env class a stable direct task declares, if any. - - Stable direct registrations may carry a ``warp_entry_point`` kwarg of the - form ``"module.path:ClassName"`` (mirroring ``env_cfg_entry_point``) naming - the warp implementation of the same task. The class is constructed with the - *stable* cfg, so the task needs no parallel warp registration or duplicated - configuration. - """ - try: + @staticmethod + def _resolve_direct_warp_class(task_id: str) -> type | None: + """Return the warp env class a stable direct task declares, if any. + + Stable direct registrations may carry a ``warp_entry_point`` kwarg of + the form ``"module.path:ClassName"`` (mirroring ``env_cfg_entry_point``) + naming the warp implementation of the same task. The class is + constructed with the *stable* cfg, so the task needs no parallel warp + registration or duplicated configuration. + """ + try: + spec = gym.spec(task_id) + except gym.error.NameNotFound as exc: + raise FrontendIncompatibleError( + f"--frontend=warp: task {task_id!r} is not registered with gymnasium." + ) from exc + target = (spec.kwargs or {}).get("warp_entry_point") + if not isinstance(target, str): + return None + module_name, _, class_name = target.partition(":") + try: + module = importlib.import_module(module_name) + except ImportError as exc: + raise FrontendIncompatibleError( + f"--frontend=warp: task {task_id!r} declares warp_entry_point {target!r}," + f" but it failed to import: {exc}" + ) from exc + env_class = getattr(module, class_name, None) + if env_class is None: + raise FrontendIncompatibleError( + f"--frontend=warp: task {task_id!r} declares warp_entry_point {target!r}," + f" but {class_name!r} is not defined in {module_name!r}." + ) + return env_class + + @classmethod + def _assert_direct_warp_registration(cls, task_id: str) -> None: + """For direct workflows without a ``warp_entry_point``, the task must be warp-registered.""" spec = gym.spec(task_id) - except gym.error.NameNotFound as exc: - raise FrontendIncompatibleError(f"--frontend=warp: task {task_id!r} is not registered with gymnasium.") from exc - target = (spec.kwargs or {}).get("warp_entry_point") - if not isinstance(target, str): + ep = spec.entry_point + module = ep if isinstance(ep, str) else (getattr(ep, "__module__", "") or "") + if not module.startswith(cls.WARP_ROOT_PREFIXES): + raise FrontendIncompatibleError( + f"--frontend=warp on direct task {task_id!r}: entry_point {ep!r}" + f" is not under {list(cls.WARP_ROOT_PREFIXES)}. Direct tasks must either" + f" declare a `warp_entry_point` in their stable registration or be" + f" registered as a warp env class." + ) + + # ------------------------------------------------------------------ + # Twin resolution + # ------------------------------------------------------------------ + + @classmethod + def _match_route(cls, module: str) -> str | None: + """Return the routed warp MDP module for ``module``, longest prefix wins.""" + if not isinstance(module, str) or not module: + return None + best_key: str | None = None + for stable_prefix in cls._mdp_routes: + if module == stable_prefix or module.startswith(f"{stable_prefix}."): + if best_key is None or len(stable_prefix) > len(best_key): + best_key = stable_prefix + return cls._mdp_routes[best_key] if best_key is not None else None + + @staticmethod + def _import_routed_module(target: str) -> ModuleType: + """Import a registered route target; a broken registration is a hard error.""" + try: + return importlib.import_module(target) + except ImportError as exc: + raise FrontendIncompatibleError( + f"registered warp MDP route target {target!r} failed to import: {exc}" + ) from exc + + @classmethod + def _cfg_route_modules(cls, cfg: Any) -> list[ModuleType]: + """Warp MDP modules routed from the cfg's class hierarchy, subclass first. + + A stable cfg consumes terms through its task's ``mdp`` namespace, which + re-exports symbols that may be *defined* in core or shared packages. The + warp mirror of that namespace is therefore the primary place to resolve + twins, and it is found by routing the modules of ``type(cfg).__mro__`` — + the subclass module first (robot-specific cfgs), then base-cfg modules + (task-family bases). + """ + modules: list[ModuleType] = [] + for klass in type(cfg).__mro__: + target = cls._match_route(getattr(klass, "__module__", "") or "") + if target is None: + continue + module = cls._import_routed_module(target) + if module not in modules: + modules.append(module) + return modules + + @classmethod + def _twin_modules(cls, symbol_module: str, cfg_route_modules: list[ModuleType]) -> list[ModuleType]: + """Warp modules to consult for a stable symbol's twin, in preference order. + + 1. The warp mirrors of the cfg's own task MDP namespace + (:meth:`_cfg_route_modules`) — task-specific twins win, including + twins for symbols defined in core/shared packages. + 2. The warp mirror routed from the symbol's defining package — covers + terms borrowed from another task family's MDP package. + 3. The shared fallback module (where generic warp twins live). + """ + modules = list(cfg_route_modules) + target = cls._match_route(symbol_module) + if target is not None: + module = cls._import_routed_module(target) + if module not in modules: + modules.append(module) + try: + fallback_module = importlib.import_module(cls._FALLBACK_MDP_MODULE) + except ModuleNotFoundError as exc: + if exc.name != cls._FALLBACK_MDP_MODULE: + raise + logger.warning("frontend.warp: fallback mdp module %r not importable", cls._FALLBACK_MDP_MODULE) + else: + if fallback_module not in modules: + modules.append(fallback_module) + return modules + + @classmethod + def _ensure_twin_providers_imported(cls) -> None: + """Import the twin-provider package so its MDP routes are registered. + + Route registration happens as an import side effect of the experimental + task packages. A caller running ``--frontend=warp`` on a stable task id + has no other reason to import the provider package, so the swap + triggers it here. A missing provider package is not an error by itself + — the swap then fails with the explicit missing-twin report. + """ + try: + importlib.import_module(cls._TWIN_PROVIDER_PACKAGE) + except ModuleNotFoundError as exc: + # Only swallow "the provider package is not installed"; a genuine + # import error inside an existing package must surface. + if exc.name != cls._TWIN_PROVIDER_PACKAGE: + raise + logger.warning("frontend.warp: twin provider package %r is not installed", cls._TWIN_PROVIDER_PACKAGE) + + @classmethod + def _resolve_warp_twin(cls, name: str, modules: Iterable[ModuleType]) -> Any | None: + """Return the same-named symbol from ``modules`` that originates under the warp packages.""" + for module in modules: + candidate = getattr(module, name, None) + if candidate is None: + continue + origin = getattr(candidate, "__module__", "") or "" + if origin.startswith(cls.WARP_ROOT_PREFIXES): + return candidate return None - module_name, _, class_name = target.partition(":") - try: - module = importlib.import_module(module_name) - except ImportError as exc: - raise FrontendIncompatibleError( - f"--frontend=warp: task {task_id!r} declares warp_entry_point {target!r}, but it failed to import: {exc}" - ) from exc - env_class = getattr(module, class_name, None) - if env_class is None: - raise FrontendIncompatibleError( - f"--frontend=warp: task {task_id!r} declares warp_entry_point {target!r}," - f" but {class_name!r} is not defined in {module_name!r}." - ) - return env_class - -def _assert_direct_warp_registration(task_id: str) -> None: - """For direct workflows without a ``warp_entry_point``, the task must be warp-registered.""" - spec = gym.spec(task_id) - ep = spec.entry_point - module = ep if isinstance(ep, str) else (getattr(ep, "__module__", "") or "") - if not module.startswith(_WARP_ROOT_PREFIXES): - raise FrontendIncompatibleError( - f"--frontend=warp on direct task {task_id!r}: entry_point {ep!r}" - f" is not under {list(_WARP_ROOT_PREFIXES)}. Direct tasks must either" - f" declare a `warp_entry_point` in their stable registration or be" - f" registered as a warp env class." - ) + @classmethod + def _is_swap_candidate(cls, value: Any) -> bool: + """Heuristic: callable or class whose origin is *not already* warp.""" + if not callable(value): + return False + origin = getattr(value, "__module__", "") or "" + if origin.startswith(cls.WARP_ROOT_PREFIXES): + return False # already a warp twin (idempotent) + return True + + @staticmethod + def _walk_terms(node: Any, path: tuple[str, ...] = ()) -> Iterator[tuple[tuple[str, ...], Any]]: + """Yield ``(path, term)`` for every MDP term cfg in the cfg tree. + + A "term" is a :class:`ManagerTermBaseCfg` (observation/reward/ + termination/event/curriculum) *or* an :class:`ActionTermCfg` — the + latter is a separate base that is **not** a ``ManagerTermBaseCfg`` + subclass, yet carries a swappable ``class_type``, so it must be matched + explicitly. + + Behavior at each node: + + * Match (a term cfg instance): yield ``(path, node)`` and stop — do not + descend into ``term.params`` / ``term.func`` / ``term.class_type``. + * Configclass: don't yield; recurse into every non-underscore *instance* + attribute (``vars(node)``), extending the path. ``observations``, + ``rewards``, ``events``, ``actions``, sub-groups like + ``observations.policy``, and anything nested deeper are reached + transparently. + * Anything else (plain Python data, callables, non-configclass objects): + stop. No yield, no recursion. + + Iterating the instance ``__dict__`` mirrors how the warp managers + consume group cfgs (they iterate ``cfg.__dict__.items()``), so the + walker sees exactly the terms the managers will see — including terms + assigned in ``__post_init__`` — while never descending into methods or + nested class objects, which live on the class rather than the instance. + + Driven entirely by type — no attribute names are hardcoded — so future + cfg layouts (extra observation groups, new nesting, etc.) are picked up + automatically as long as their terms subclass one of the term base cfgs. + """ + from isaaclab.managers.manager_term_cfg import ActionTermCfg, ManagerTermBaseCfg + + if isinstance(node, (ManagerTermBaseCfg, ActionTermCfg)): + yield path, node + return + if not hasattr(node, "__dataclass_fields__"): + return + for name, value in vars(node).items(): + if name.startswith("_") or value is None: + continue + yield from WarpFrontend._walk_terms(value, path + (name,)) -def _match_route(module: str) -> str | None: - """Return the routed warp MDP module for ``module``, longest prefix wins.""" - if not isinstance(module, str) or not module: - return None - best_key: str | None = None - for stable_prefix in _WARP_MDP_MODULE_ROUTES: - if module == stable_prefix or module.startswith(f"{stable_prefix}."): - if best_key is None or len(stable_prefix) > len(best_key): - best_key = stable_prefix - return _WARP_MDP_MODULE_ROUTES[best_key] if best_key is not None else None - - -def _import_routed_module(target: str) -> ModuleType: - """Import a registered route target; a broken registration is a hard error.""" - try: - return importlib.import_module(target) - except ImportError as exc: - raise FrontendIncompatibleError(f"registered warp MDP route target {target!r} failed to import: {exc}") from exc - - -def _cfg_route_modules(cfg: Any) -> list[ModuleType]: - """Warp MDP modules routed from the cfg's class hierarchy, subclass first. - - A stable cfg consumes terms through its task's ``mdp`` namespace, which - re-exports symbols that may be *defined* in core or shared packages. The - warp mirror of that namespace is therefore the primary place to resolve - twins, and it is found by routing the modules of ``type(cfg).__mro__`` — - the subclass module first (robot-specific cfgs), then base-cfg modules - (task-family bases). - """ - modules: list[ModuleType] = [] - for klass in type(cfg).__mro__: - target = _match_route(getattr(klass, "__module__", "") or "") - if target is None: - continue - module = _import_routed_module(target) - if module not in modules: - modules.append(module) - return modules - - -def _twin_modules(symbol_module: str, cfg_route_modules: list[ModuleType]) -> list[ModuleType]: - """Warp modules to consult for a stable symbol's twin, in preference order. - - 1. The warp mirrors of the cfg's own task MDP namespace - (:func:`_cfg_route_modules`) — task-specific twins win, including twins - for symbols defined in core/shared packages. - 2. The warp mirror routed from the symbol's defining package — covers terms - borrowed from another task family's MDP package. - 3. The shared :mod:`isaaclab_experimental.envs.mdp` fallback (where generic - warp twins live). - """ - modules = list(cfg_route_modules) - target = _match_route(symbol_module) - if target is not None: - module = _import_routed_module(target) - if module not in modules: - modules.append(module) - fallback = "isaaclab_experimental.envs.mdp" - try: - fallback_module = importlib.import_module(fallback) - except ModuleNotFoundError as exc: - if exc.name != fallback: - raise - logger.warning("frontend.warp: fallback mdp module %r not importable", fallback) - else: - if fallback_module not in modules: - modules.append(fallback_module) - return modules - - -def _ensure_twin_providers_imported() -> None: - """Import the twin-provider package so its MDP routes are registered. - - Route registration happens as an import side effect of the experimental - task packages. A caller running ``--frontend=warp`` on a stable task id - has no other reason to import :mod:`isaaclab_tasks_experimental`, so the - swap triggers it here. A missing provider package is not an error by - itself — the swap then fails with the explicit missing-twin report. - """ - try: - importlib.import_module(_TWIN_PROVIDER_PACKAGE) - except ModuleNotFoundError as exc: - # Only swallow "the provider package is not installed"; a genuine import - # error inside an existing package must surface. - if exc.name != _TWIN_PROVIDER_PACKAGE: - raise - logger.warning("frontend.warp: twin provider package %r is not installed", _TWIN_PROVIDER_PACKAGE) - - -def _resolve_warp_twin(name: str, modules: Iterable[ModuleType]) -> Any | None: - """Return the same-named symbol from ``modules`` that originates under the warp packages.""" - for module in modules: - candidate = getattr(module, name, None) - if candidate is None: - continue - origin = getattr(candidate, "__module__", "") or "" - if origin.startswith(_WARP_ROOT_PREFIXES): - return candidate - return None - - -def _is_swap_candidate(value: Any) -> bool: - """Heuristic: callable or class whose origin is *not already* warp.""" - if not callable(value): - return False - origin = getattr(value, "__module__", "") or "" - if origin.startswith(_WARP_ROOT_PREFIXES): - return False # already a warp twin (idempotent) - return True - - -def _walk_terms(node: Any, path: tuple[str, ...] = ()) -> Iterator[tuple[tuple[str, ...], Any]]: - """Yield ``(path, term)`` for every MDP term cfg in the cfg tree. - - A "term" is a :class:`ManagerTermBaseCfg` (observation/reward/termination/ - event/curriculum) *or* an :class:`ActionTermCfg` — the latter is a separate - base that is **not** a ``ManagerTermBaseCfg`` subclass, yet carries a - swappable ``class_type``, so it must be matched explicitly. - - Behavior at each node: - - * Match (a term cfg instance): yield ``(path, node)`` and stop — do not - descend into ``term.params`` / ``term.func`` / ``term.class_type``. - * Configclass: don't yield; recurse into every non-underscore *instance* - attribute (``vars(node)``), extending the path. ``observations``, - ``rewards``, ``events``, ``actions``, sub-groups like - ``observations.policy`` / ``observations.perception``, and anything - nested deeper are reached transparently. - * Anything else (plain Python data, callables, non-configclass objects): - stop. No yield, no recursion. - - Iterating the instance ``__dict__`` mirrors how the warp managers consume - group cfgs (they iterate ``cfg.__dict__.items()``), so the walker sees - exactly the terms the managers will see — including terms assigned in - ``__post_init__`` — while never descending into methods or nested class - objects, which live on the class rather than the instance. - - Driven entirely by type — no attribute names are hardcoded — so future - cfg layouts (extra observation groups, new nesting, etc.) are picked up - automatically as long as their terms subclass one of the term base cfgs. - """ - from isaaclab.managers.manager_term_cfg import ActionTermCfg, ManagerTermBaseCfg - - if isinstance(node, (ManagerTermBaseCfg, ActionTermCfg)): - yield path, node - return - if not hasattr(node, "__dataclass_fields__"): - return - for name, value in vars(node).items(): - if name.startswith("_") or value is None: - continue - yield from _walk_terms(value, path + (name,)) +# Module-level alias used by the task packages at import time. +register_mdp_route = WarpFrontend.register_mdp_route diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py index f112a6ddf562..e32d80fd2b99 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend.py +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -135,9 +135,9 @@ def _term(func=None, params: dict | None = None) -> RewardTermCfg: @pytest.fixture def fake_routes(monkeypatch: pytest.MonkeyPatch) -> dict[str, str]: - """Isolate the module-level route registry for a test.""" + """Isolate the WarpFrontend route registry for a test.""" routes: dict[str, str] = {} - monkeypatch.setattr(fe, "_WARP_MDP_MODULE_ROUTES", routes) + monkeypatch.setattr(fe.WarpFrontend, "_mdp_routes", routes) return routes @@ -245,19 +245,19 @@ def _cfg_with_physics(physics: Any) -> Any: def test_require_newton_passes_for_newton(): - fe._require_newton_physics(_cfg_with_physics(NewtonCfg()), "Isaac-Test-v0") # no raise + fe.WarpFrontend._require_newton_physics(_cfg_with_physics(NewtonCfg()), "Isaac-Test-v0") # no raise def test_require_newton_rejects_physx(): with pytest.raises(FrontendIncompatibleError) as exc: - fe._require_newton_physics(_cfg_with_physics(PhysxCfg()), "Isaac-Test-v0") + fe.WarpFrontend._require_newton_physics(_cfg_with_physics(PhysxCfg()), "Isaac-Test-v0") assert "presets=newton_mjwarp" in str(exc.value) assert "PhysxCfg" in str(exc.value) def test_require_newton_rejects_none(): with pytest.raises(FrontendIncompatibleError): - fe._require_newton_physics(_cfg_with_physics(None), "Isaac-Test-v0") + fe.WarpFrontend._require_newton_physics(_cfg_with_physics(None), "Isaac-Test-v0") # ====================================================================== @@ -271,7 +271,7 @@ def test_walk_terms_yields_each_term_with_its_path(): events=_EventsCfg(e1=EventTermCfg(func=_stable_func, mode="reset")), ) # Configclass instances aren't hashable, so collect paths only. - paths = {".".join(p) for p, _ in fe._walk_terms(cfg)} + paths = {".".join(p) for p, _ in fe.WarpFrontend._walk_terms(cfg)} assert paths == {"rewards.r1", "rewards.r2", "events.e1"} @@ -282,7 +282,7 @@ def test_walk_terms_descends_into_obs_subgroups(): perception=_ExtraObsGroup(o3=ObservationTermCfg(func=_stable_func)), ), ) - paths = {".".join(p) for p, _ in fe._walk_terms(cfg)} + paths = {".".join(p) for p, _ in fe.WarpFrontend._walk_terms(cfg)} # Discovery is purely type-driven; no obs group name is hardcoded. assert paths == {"observations.policy.o1", "observations.perception.o3"} @@ -291,7 +291,7 @@ def test_walk_terms_stops_at_terms(): # The walker must not descend into term.params / term.func — yields the term itself. nested_se_cfg = StableSceneEntityCfg(name="robot") cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(params={"asset_cfg": nested_se_cfg}))) - terms = list(fe._walk_terms(cfg)) + terms = list(fe.WarpFrontend._walk_terms(cfg)) assert len(terms) == 1 _, term = terms[0] assert isinstance(term, RewardTermCfg) @@ -300,12 +300,12 @@ def test_walk_terms_stops_at_terms(): def test_walk_terms_skips_non_configclass_attrs(): # A namespace without __dataclass_fields__ is not descended into. cfg = types.SimpleNamespace(some_plain_attr="hello") - assert list(fe._walk_terms(cfg)) == [] + assert list(fe.WarpFrontend._walk_terms(cfg)) == [] def test_walk_terms_skips_none_subtrees(): cfg = _CfgFixture(rewards=None, events=None) - assert list(fe._walk_terms(cfg)) == [] + assert list(fe.WarpFrontend._walk_terms(cfg)) == [] @configclass @@ -324,7 +324,7 @@ def test_walk_terms_ignores_nested_class_objects(): # ``cfg.__dict__`` consumption), so a nested class object — reachable as an # attribute on the instance — must never contribute paths. cfg = _GroupWithNestedClass(o1=ObservationTermCfg(func=_stable_func)) - paths = {".".join(p) for p, _ in fe._walk_terms(cfg)} + paths = {".".join(p) for p, _ in fe.WarpFrontend._walk_terms(cfg)} assert paths == {"o1"} @@ -338,13 +338,13 @@ def test_register_mdp_route_and_longest_prefix_match(fake_routes: dict[str, str] register_mdp_route("isaaclab_tasks.core.locomotion.humanoid", "warp.humanoid.mdp") # Longest registered prefix wins. - assert fe._match_route("isaaclab_tasks.core.locomotion.humanoid.mdp.rewards") == "warp.humanoid.mdp" - assert fe._match_route("isaaclab_tasks.core.locomotion.ant.mdp") == "warp.locomotion.mdp" + assert fe.WarpFrontend._match_route("isaaclab_tasks.core.locomotion.humanoid.mdp.rewards") == "warp.humanoid.mdp" + assert fe.WarpFrontend._match_route("isaaclab_tasks.core.locomotion.ant.mdp") == "warp.locomotion.mdp" # Exact package match works too. - assert fe._match_route("isaaclab_tasks.core.locomotion.humanoid") == "warp.humanoid.mdp" + assert fe.WarpFrontend._match_route("isaaclab_tasks.core.locomotion.humanoid") == "warp.humanoid.mdp" # Prefix matching is per package segment, not per character. - assert fe._match_route("isaaclab_tasks.core.locomotion_extras.mdp") is None - assert fe._match_route("isaaclab.envs.mdp.rewards") is None + assert fe.WarpFrontend._match_route("isaaclab_tasks.core.locomotion_extras.mdp") is None + assert fe.WarpFrontend._match_route("isaaclab.envs.mdp.rewards") is None def test_register_mdp_route_is_idempotent(fake_routes: dict[str, str]): @@ -362,13 +362,13 @@ def test_register_mdp_route_rejects_conflicts(fake_routes: dict[str, str]): def test_cfg_route_modules_resolves_cfg_class_module(fake_routes: dict[str, str]): # Route this test module's own name so the fixture cfg class matches. register_mdp_route(__name__, "math") - assert fe._cfg_route_modules(_CfgFixture()) == [importlib.import_module("math")] + assert fe.WarpFrontend._cfg_route_modules(_CfgFixture()) == [importlib.import_module("math")] def test_cfg_route_modules_raises_on_broken_target(fake_routes: dict[str, str]): register_mdp_route(__name__, "definitely_not_an_importable_module") with pytest.raises(FrontendIncompatibleError): - fe._cfg_route_modules(_CfgFixture()) + fe.WarpFrontend._cfg_route_modules(_CfgFixture()) # ====================================================================== @@ -391,9 +391,11 @@ def _fake_mdp_module(name_to_symbol: dict[str, Any]) -> _FakeMdpModule: def _patch_twin_modules(monkeypatch: pytest.MonkeyPatch, modules: list[Any]) -> None: """Pin twin resolution to ``modules`` and keep unit tests hermetic.""" - monkeypatch.setattr(fe, "_ensure_twin_providers_imported", lambda: None) - monkeypatch.setattr(fe, "_cfg_route_modules", lambda cfg: []) - monkeypatch.setattr(fe, "_twin_modules", lambda symbol_module, cfg_route_modules: modules) + monkeypatch.setattr(fe.WarpFrontend, "_ensure_twin_providers_imported", classmethod(lambda cls: None)) + monkeypatch.setattr(fe.WarpFrontend, "_cfg_route_modules", classmethod(lambda cls, cfg: [])) + monkeypatch.setattr( + fe.WarpFrontend, "_twin_modules", classmethod(lambda cls, symbol_module, cfg_route_modules: modules) + ) def test_swap_mdp_swaps_func_and_class_type(monkeypatch: pytest.MonkeyPatch): @@ -403,7 +405,7 @@ def test_swap_mdp_swaps_func_and_class_type(monkeypatch: pytest.MonkeyPatch): term_action = _term() term_action.class_type = _StableActionCls # set attr to exercise class_type swap cfg = _CfgFixture(rewards=_RewardsCfg(r1=term_reward, r2=term_action)) - fe._swap_mdp(cfg, "Isaac-Test-v0") + fe.WarpFrontend._swap_mdp(cfg, "Isaac-Test-v0") assert cfg.rewards.r1.func is _warp_twin_func assert cfg.rewards.r2.class_type is _WarpActionCls @@ -412,7 +414,7 @@ def test_swap_mdp_missing_twin_raises_with_path_list(monkeypatch: pytest.MonkeyP _patch_twin_modules(monkeypatch, [_fake_mdp_module({})]) cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(func=_stable_func))) with pytest.raises(FrontendIncompatibleError) as exc: - fe._swap_mdp(cfg, "Isaac-Test-v0") + fe.WarpFrontend._swap_mdp(cfg, "Isaac-Test-v0") msg = str(exc.value) assert "rewards.r1.func" in msg assert "_stable_func" in msg @@ -423,7 +425,7 @@ def test_swap_mdp_missing_twin_raises_with_path_list(monkeypatch: pytest.MonkeyP def test_swap_mdp_skips_already_warp(monkeypatch: pytest.MonkeyPatch): _patch_twin_modules(monkeypatch, [_fake_mdp_module({})]) cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(func=_warp_twin_func))) - fe._swap_mdp(cfg, "Isaac-Test-v0") # no raise + fe.WarpFrontend._swap_mdp(cfg, "Isaac-Test-v0") # no raise assert cfg.rewards.r1.func is _warp_twin_func @@ -438,7 +440,7 @@ def test_promote_scene_entity_cfgs_promotes_in_params(): r1=_term(params={"asset_cfg": StableSceneEntityCfg(name="robot", joint_names=["lf_hip"]), "scale": 1.0}) ) ) - fe._promote_scene_entity_cfgs(cfg) + fe.WarpFrontend._promote_scene_entity_cfgs(cfg) promoted = cfg.rewards.r1.params["asset_cfg"] assert isinstance(promoted, WarpSceneEntityCfg) assert promoted.name == "robot" @@ -455,7 +457,7 @@ def test_promote_scene_entity_cfgs_skips_already_warp(): warp = WarpSceneEntityCfg(name="robot", joint_names=["lf_hip"]) cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(params={"asset_cfg": warp}))) before = cfg.rewards.r1.params["asset_cfg"] - fe._promote_scene_entity_cfgs(cfg) + fe.WarpFrontend._promote_scene_entity_cfgs(cfg) after = cfg.rewards.r1.params["asset_cfg"] assert isinstance(after, WarpSceneEntityCfg) # The asset_cfg object was not replaced by another from_stable call. @@ -480,7 +482,7 @@ def test_promote_scene_entity_cfgs_walks_all_term_groups(): c1=EventTermCfg(func=_stable_func, mode="reset", params={"asset_cfg": StableSceneEntityCfg(name="c")}) ), ) - fe._promote_scene_entity_cfgs(cfg) + fe.WarpFrontend._promote_scene_entity_cfgs(cfg) # Warp-managed groups (rewards, observations incl. sub-groups, events) are promoted. assert isinstance(cfg.rewards.r1.params["asset_cfg"], WarpSceneEntityCfg) assert isinstance(cfg.observations.policy.o1.params["asset_cfg"], WarpSceneEntityCfg) @@ -502,29 +504,29 @@ def test_promote_scene_entity_cfgs_walks_all_term_groups(): def test_resolve_warp_twin_accepts_warp_origin(): module = types.SimpleNamespace(foo=_warp_twin_func) - assert fe._resolve_warp_twin("foo", [module]) is _warp_twin_func + assert fe.WarpFrontend._resolve_warp_twin("foo", [module]) is _warp_twin_func def test_resolve_warp_twin_rejects_stable_origin(): module = types.SimpleNamespace(foo=_stable_func) # same name, stable origin - assert fe._resolve_warp_twin("foo", [module]) is None + assert fe.WarpFrontend._resolve_warp_twin("foo", [module]) is None def test_resolve_warp_twin_returns_none_when_absent(): - assert fe._resolve_warp_twin("missing", [types.SimpleNamespace()]) is None + assert fe.WarpFrontend._resolve_warp_twin("missing", [types.SimpleNamespace()]) is None def test_is_swap_candidate_stable_callable(): - assert fe._is_swap_candidate(_stable_func) + assert fe.WarpFrontend._is_swap_candidate(_stable_func) def test_is_swap_candidate_rejects_warp_callable(): - assert not fe._is_swap_candidate(_warp_twin_func) + assert not fe.WarpFrontend._is_swap_candidate(_warp_twin_func) def test_is_swap_candidate_rejects_non_callables(): - assert not fe._is_swap_candidate(42) - assert not fe._is_swap_candidate("string") + assert not fe.WarpFrontend._is_swap_candidate(42) + assert not fe.WarpFrontend._is_swap_candidate("string") # ====================================================================== @@ -555,35 +557,36 @@ def _register_direct_stub_tasks(): def test_direct_guard_accepts_warp_rooted(): - fe._assert_direct_warp_registration("_Frontend-Test-Warp-Direct-v0") # no raise + fe.WarpFrontend._assert_direct_warp_registration("_Frontend-Test-Warp-Direct-v0") # no raise def test_direct_guard_rejects_stable_rooted(): with pytest.raises(FrontendIncompatibleError) as exc: - fe._assert_direct_warp_registration("_Frontend-Test-Stable-Direct-v0") + fe.WarpFrontend._assert_direct_warp_registration("_Frontend-Test-Stable-Direct-v0") assert "isaaclab_experimental" in str(exc.value) def test_resolve_direct_warp_class_returns_declared_class(): import types as types_module - assert fe._resolve_direct_warp_class("_Frontend-Test-Declared-Direct-v0") is types_module.SimpleNamespace + env_class = fe.WarpFrontend._resolve_direct_warp_class("_Frontend-Test-Declared-Direct-v0") + assert env_class is types_module.SimpleNamespace def test_resolve_direct_warp_class_returns_none_without_declaration(): - assert fe._resolve_direct_warp_class("_Frontend-Test-Stable-Direct-v0") is None + assert fe.WarpFrontend._resolve_direct_warp_class("_Frontend-Test-Stable-Direct-v0") is None def test_resolve_direct_warp_class_rejects_broken_module(): with pytest.raises(FrontendIncompatibleError): - fe._resolve_direct_warp_class("_Frontend-Test-Broken-Direct-v0") + fe.WarpFrontend._resolve_direct_warp_class("_Frontend-Test-Broken-Direct-v0") def test_resolve_direct_warp_class_rejects_missing_class(): with pytest.raises(FrontendIncompatibleError): - fe._resolve_direct_warp_class("_Frontend-Test-Missing-Class-Direct-v0") + fe.WarpFrontend._resolve_direct_warp_class("_Frontend-Test-Missing-Class-Direct-v0") def test_resolve_direct_warp_class_rejects_unknown_task(): with pytest.raises(FrontendIncompatibleError): - fe._resolve_direct_warp_class("Frontend-Test-NotRegistered-v0") + fe.WarpFrontend._resolve_direct_warp_class("Frontend-Test-NotRegistered-v0") From e1066c21a80486b995fe6fcd1010d6358edd12de Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 16 Jul 2026 01:23:05 -0700 Subject: [PATCH 31/58] Make the warp frontend module warp-only The shared RL CLI already peels off the torch path before importing the optional isaaclab_experimental package, so the frontend's own torch branch (module-level build() with a Frontend enum) was dead code that duplicated the gym.make dispatch. Remove it: create_isaaclab_env owns the torch/warp dispatch and calls WarpFrontend.build_env directly, and the module's public surface shrinks to WarpFrontend, Workflow, register_mdp_route, and the error type. Callers use WarpFrontend.adapt_cfg in place of the adapt_cfg_for_warp alias. --- scripts/reinforcement_learning/common.py | 9 +- .../test_reinforcement_learning_common.py | 8 +- .../changelog.d/warp-manager-bridge.rst | 2 +- .../isaaclab_experimental/envs/frontend.py | 91 ++++--------------- .../envs/manager_based_rl_env_warp.py | 4 +- .../test/envs/test_frontend.py | 19 +--- .../test/envs/test_frontend_cfg_conversion.py | 6 +- 7 files changed, 38 insertions(+), 101 deletions(-) diff --git a/scripts/reinforcement_learning/common.py b/scripts/reinforcement_learning/common.py index efebd09747fe..cd00334e393c 100644 --- a/scripts/reinforcement_learning/common.py +++ b/scripts/reinforcement_learning/common.py @@ -440,13 +440,14 @@ def create_isaaclab_env( The created Gymnasium environment. """ render_mode = "rgb_array" if args_cli.video else None - frontend = args_cli.frontend - if frontend == "torch": + if args_cli.frontend == "torch": env = gym.make(task, cfg=env_cfg, render_mode=render_mode) else: - from isaaclab_experimental.envs.frontend import build as build_frontend_env + # Imported lazily: the warp frontend lives in the optional + # isaaclab_experimental package, and the torch path must work without it. + from isaaclab_experimental.envs.frontend import WarpFrontend - env = build_frontend_env(frontend, env_cfg, task, render_mode=render_mode) + env = WarpFrontend.build_env(env_cfg, task, render_mode=render_mode) if convert_marl_to_single_agent and isinstance(env.unwrapped.cfg, DirectMARLEnvCfg): from isaaclab.envs import multi_agent_to_single_agent diff --git a/source/isaaclab/test/test_reinforcement_learning_common.py b/source/isaaclab/test/test_reinforcement_learning_common.py index 9de4173bd8f7..9e2b4c24e701 100644 --- a/source/isaaclab/test/test_reinforcement_learning_common.py +++ b/source/isaaclab/test/test_reinforcement_learning_common.py @@ -262,17 +262,17 @@ def test_create_isaaclab_env_uses_selected_warp_frontend(monkeypatch: pytest.Mon env_cfg = object() calls: list[tuple[Any, ...]] = [] - def fake_build(frontend: str, cfg: Any, task: str, **kwargs: Any) -> Any: - calls.append((frontend, cfg, task, kwargs)) + def fake_build_env(cfg: Any, task: str, **kwargs: Any) -> Any: + calls.append((cfg, task, kwargs)) return expected_env - monkeypatch.setattr(frontend_module, "build", fake_build) + monkeypatch.setattr(frontend_module.WarpFrontend, "build_env", fake_build_env) args_cli = argparse.Namespace(video=True, frontend="warp") env = create_isaaclab_env("Isaac-Test", env_cfg, args_cli, convert_marl_to_single_agent=False) assert env is expected_env - assert calls == [("warp", env_cfg, "Isaac-Test", {"render_mode": "rgb_array"})] + assert calls == [(env_cfg, "Isaac-Test", {"render_mode": "rgb_array"})] def test_dispatch_library_entrypoint_shows_help_without_library( diff --git a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst index 9cb34541bee0..1d56783ad0e9 100644 --- a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst +++ b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst @@ -19,7 +19,7 @@ Changed * Changed :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` to adapt its configuration in ``__init__`` via - :func:`~isaaclab_experimental.envs.frontend.adapt_cfg_for_warp`, so + :meth:`~isaaclab_experimental.envs.frontend.WarpFrontend.adapt_cfg`, so registered warp task variants can derive from stable configurations instead of duplicating them. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index ca5174c479f5..685fba4145bf 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -3,11 +3,10 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Runtime selector for IsaacLab tasks (``--frontend {torch,warp}``). +"""Warp runtime frontend for IsaacLab tasks (``--frontend warp``). -:func:`build` decides which runtime constructs the env: torch via ``gym.make``, -or warp via :class:`WarpFrontend`, which encapsulates everything the warp -runtime needs from a stable task definition: +:class:`WarpFrontend` encapsulates everything the warp runtime needs from a +stable task definition: * Manager-based tasks: adapt the stable cfg in place (Newton physics check, :class:`SceneEntityCfg` promotion, MDP twin swap) and construct @@ -18,7 +17,9 @@ via its ``warp_entry_point`` kwarg, sharing the stable cfg. Task packages declare where their warp MDP twins live with -:func:`register_mdp_route`; the frontend holds no per-task table. +:func:`register_mdp_route`; the frontend holds no per-task table. The torch +runtime never touches this module — the shared RL CLI dispatches to plain +``gym.make`` before importing it, so this (optional) package stays warp-only. """ from __future__ import annotations @@ -39,28 +40,13 @@ __all__ = [ - "Frontend", "FrontendIncompatibleError", "WarpFrontend", "Workflow", - "adapt_cfg_for_warp", - "build", "register_mdp_route", ] -# --------------------------------------------------------------------------- -# Enums + error -# --------------------------------------------------------------------------- - - -class Frontend(StrEnum): - """Runtime selector exposed by ``--frontend``.""" - - TORCH = "torch" - WARP = "warp" - - class Workflow(StrEnum): """Manager-based vs direct env workflow.""" @@ -69,55 +55,7 @@ class Workflow(StrEnum): class FrontendIncompatibleError(RuntimeError): - """Raised when the chosen frontend can't run the requested task.""" - - -# --------------------------------------------------------------------------- -# Public entry points -# --------------------------------------------------------------------------- - - -def build( - frontend: Frontend | str, - env_cfg: Any, - task_id: str, - **construct_kwargs: Any, -) -> gym.Env: - """Construct the env on the selected runtime. - - Args: - frontend: ``"torch"`` (default IsaacLab path) or ``"warp"`` - (:class:`WarpFrontend`). - env_cfg: Stable env cfg. Mutated in place when ``frontend == "warp"``. - task_id: Gym registration id, e.g. ``"Isaac-Cartpole"``. - **construct_kwargs: Forwarded to the env constructor (``render_mode``, …). - - Returns: - The constructed :class:`gym.Env`. - - Raises: - FrontendIncompatibleError: If the warp runtime can't run the task - (wrong physics, missing MDP twins, direct task without a warp - implementation). - """ - frontend = Frontend(frontend) - if frontend is Frontend.TORCH: - return gym.make(task_id, cfg=env_cfg, **construct_kwargs) - return WarpFrontend.build_env(env_cfg, task_id, **construct_kwargs) - - -def adapt_cfg_for_warp(cfg: Any) -> None: - """Mutate a stable manager-based ``cfg`` in place so warp managers can consume it. - - Functional alias for :meth:`WarpFrontend.adapt_cfg`, used by - :class:`ManagerBasedRLEnvWarp` in its ``__init__``. - """ - WarpFrontend.adapt_cfg(cfg) - - -# --------------------------------------------------------------------------- -# Warp frontend -# --------------------------------------------------------------------------- + """Raised when the warp frontend can't run the requested task.""" class WarpFrontend: @@ -206,7 +144,18 @@ def register_mdp_route(cls, stable_package: str, warp_mdp_module: str) -> None: @classmethod def build_env(cls, env_cfg: Any, task_id: str, **construct_kwargs: Any) -> gym.Env: - """Construct the warp env for ``task_id`` from a stable env cfg.""" + """Construct the warp env for ``task_id`` from a stable env cfg. + + Args: + env_cfg: Stable env cfg. Manager-based cfgs are mutated in place. + task_id: Gym registration id, e.g. ``"Isaac-Cartpole"``. + **construct_kwargs: Forwarded to the env constructor (``render_mode``, …). + + Raises: + FrontendIncompatibleError: If the warp runtime can't run the task + (wrong physics, missing MDP twins, direct task without a warp + implementation). + """ if cls._detect_workflow(env_cfg) is Workflow.DIRECT: return cls._build_direct_env(env_cfg, task_id, **construct_kwargs) return cls._build_manager_env(env_cfg, **construct_kwargs) @@ -390,7 +339,7 @@ def _swap_mdp(cls, cfg: Any, label: str) -> None: @staticmethod def _detect_workflow(cfg: Any) -> Workflow: - """Classify the env cfg into manager-based or direct (used to pick build path). + """Classify the env cfg into manager-based or direct (used to pick the build path). Note: The env cfg roots (ManagerBasedRLEnvCfg, DirectRLEnvCfg, diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index 502c9a132042..1c3535e887ec 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -97,9 +97,9 @@ def __init__(self, cfg: ManagerBasedRLEnvCfg, render_mode: str | None = None, ** # promotion, MDP twin swap). Idempotent: a warp-native cfg passes through # unchanged, and a stable-derived cfg (``--frontend=warp`` or a registered # warp task variant subclassing a stable cfg) is adapted in place. - from isaaclab_experimental.envs.frontend import adapt_cfg_for_warp + from isaaclab_experimental.envs.frontend import WarpFrontend - adapt_cfg_for_warp(cfg) + WarpFrontend.adapt_cfg(cfg) # -- counter for curriculum self.common_step_counter = 0 diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py index e32d80fd2b99..859f0b346769 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend.py +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -7,7 +7,7 @@ Pure-Python unit tests; no app launch. Covers: -* :class:`Frontend` / :class:`Workflow` enum surface. +* :class:`Workflow` enum surface. * :func:`SceneEntityCfg.from_stable` field copy. * :func:`_require_newton_physics` hard-check. * :func:`_walk_terms` recursive ManagerTermBaseCfg discovery over instance @@ -34,10 +34,9 @@ import isaaclab_experimental.envs.frontend as fe import pytest from isaaclab_experimental.envs.frontend import ( - Frontend, FrontendIncompatibleError, + WarpFrontend, Workflow, - build, register_mdp_route, ) from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as WarpSceneEntityCfg @@ -146,18 +145,6 @@ def fake_routes(monkeypatch: pytest.MonkeyPatch) -> dict[str, str]: # ====================================================================== -def test_frontend_values(): - assert Frontend.TORCH == "torch" - assert Frontend.WARP == "warp" - - -def test_frontend_coercion(): - assert Frontend("torch") is Frontend.TORCH - assert Frontend("warp") is Frontend.WARP - with pytest.raises(ValueError): - Frontend("kit") - - def test_workflow_values(): assert Workflow.MANAGER_BASED == "manager_based" assert Workflow.DIRECT == "direct" @@ -182,7 +169,7 @@ def fake_env(*, cfg: Any, **kwargs: Any) -> Any: return expected_env with patch.object(envs, "ManagerBasedRLEnvWarp", fake_env): - env = build("warp", cfg, "Isaac-Test", render_mode="rgb_array") + env = WarpFrontend.build_env(cfg, "Isaac-Test", render_mode="rgb_array") assert env is expected_env assert calls == [(cfg, {"render_mode": "rgb_array"})] diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index 4fc8d100675b..f94ba002fb6c 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -5,7 +5,7 @@ """Coverage tests for warp task configuration adaptation. -Two sweeps, both running :func:`adapt_cfg_for_warp` — the same adaptation +Two sweeps, both running :meth:`WarpFrontend.adapt_cfg` — the same adaptation :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` runs in its ``__init__``: @@ -27,7 +27,7 @@ import gymnasium as gym import isaaclab_tasks_experimental # noqa: F401 import pytest -from isaaclab_experimental.envs.frontend import adapt_cfg_for_warp +from isaaclab_experimental.envs.frontend import WarpFrontend from isaaclab_newton.physics import NewtonCfg # Registering the task packages is the whole point — import for side effects. @@ -74,7 +74,7 @@ def _load_adapted_cfg(cfg_entry_point: str): cfg = resolve_presets(cfg, selected=("newton_mjwarp",)) assert isinstance(cfg.sim.physics, NewtonCfg), "task does not provide a newton_mjwarp physics preset" # Raises FrontendIncompatibleError if any warp-managed term lacks a warp twin. - adapt_cfg_for_warp(cfg) + WarpFrontend.adapt_cfg(cfg) return cfg From 8694427778b94d829eb343d3afdaf21f0989fb6e Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 16 Jul 2026 15:36:54 -0700 Subject: [PATCH 32/58] Resolve warp twins by package mirror instead of a registry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The experimental tree now mirrors the stable tree exactly, so twin routing needs no declarations: mirror the module root (isaaclab -> isaaclab_experimental, isaaclab_tasks -> isaaclab_tasks_experimental) and look the symbol up on the nearest .mdp package of the mirrored path — the cfg's own task namespace first, then the symbol's defining package. Resolution is a pure function of the installed package tree; misses accumulate across the whole cfg and are reported in a single hard failure listing every missing twin and all searched modules. This deletes register_mdp_route, the route registry, the twin-provider side-effect import, and all five task-package registration calls. The humanoid MDP twins move to core/locomotion/mdp — the true mirror of the stable shared locomotion package — which also removes the Ant special case (both tasks now resolve mechanically). --- .../isaaclab_experimental/envs/frontend.py | 197 +++++++----------- .../test/envs/test_frontend.py | 85 ++++---- .../core/cartpole/__init__.py | 5 - .../core/locomotion/__init__.py | 8 +- .../core/locomotion/ant/__init__.py | 20 +- .../core/locomotion/humanoid/__init__.py | 16 +- .../locomotion/{humanoid => }/mdp/__init__.py | 0 .../{humanoid => }/mdp/__init__.pyi | 0 .../{humanoid => }/mdp/observations.py | 0 .../locomotion/{humanoid => }/mdp/rewards.py | 0 .../core/reach/__init__.py | 5 - .../core/velocity/__init__.py | 6 - 12 files changed, 132 insertions(+), 210 deletions(-) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/{humanoid => }/mdp/__init__.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/{humanoid => }/mdp/__init__.pyi (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/{humanoid => }/mdp/observations.py (100%) rename source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/{humanoid => }/mdp/rewards.py (100%) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index 685fba4145bf..29d14191a9bd 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -16,15 +16,20 @@ * Direct tasks: construct the warp env class the stable registration declares via its ``warp_entry_point`` kwarg, sharing the stable cfg. -Task packages declare where their warp MDP twins live with -:func:`register_mdp_route`; the frontend holds no per-task table. The torch -runtime never touches this module — the shared RL CLI dispatches to plain -``gym.make`` before importing it, so this (optional) package stays warp-only. +Twin resolution is purely mechanical: the experimental packages mirror the +stable tree (``isaaclab.* ↔ isaaclab_experimental.*``, +``isaaclab_tasks.* ↔ isaaclab_tasks_experimental.*``), and a stable symbol's +twin is looked up by name on the nearest ``.mdp`` package of the mirrored +path. There is no registration; anything unresolved is collected and reported +in one hard failure. The torch runtime never touches this module — the shared +RL CLI dispatches to plain ``gym.make`` before importing it, so this +(optional) package stays warp-only. """ from __future__ import annotations import importlib +import importlib.util import logging from collections.abc import Iterable, Iterator from enum import StrEnum @@ -43,7 +48,6 @@ "FrontendIncompatibleError", "WarpFrontend", "Workflow", - "register_mdp_route", ] @@ -61,17 +65,21 @@ class FrontendIncompatibleError(RuntimeError): class WarpFrontend: """Builds environments on the warp runtime from stable task definitions. - The class is a stateless namespace apart from the MDP route registry - (:attr:`_mdp_routes`), which experimental task packages populate at import - time via :meth:`register_mdp_route`. The stable and experimental task trees - mirror each other (``isaaclab_tasks.core.X`` ↔ - ``isaaclab_tasks_experimental.core.X``), so most routes are the mechanical - mirror; the registry keeps routing explicit and accommodates exceptions - (the Ant task borrows the Humanoid twins). + The class is a stateless namespace: twin resolution is a pure function of + the stable module paths and the mirrored package tree (:attr:`MIRROR_ROOTS`); + there is no registry and no import-order dependence. """ + # Stable package roots mirrored by their warp implementation packages. Twin + # resolution is purely mechanical: mirror the root, then look up the nearest + # enclosing ``.mdp`` package on the mirrored path. + MIRROR_ROOTS: ClassVar[dict[str, str]] = { + "isaaclab_tasks": "isaaclab_tasks_experimental", + "isaaclab": "isaaclab_experimental", + } + # Prefixes that identify warp-origin symbols and registrations. - WARP_ROOT_PREFIXES: ClassVar[tuple[str, ...]] = ("isaaclab_experimental", "isaaclab_tasks_experimental") + WARP_ROOT_PREFIXES: ClassVar[tuple[str, ...]] = tuple(MIRROR_ROOTS.values()) # Top-level cfg groups whose managers run warp-first. Only terms under these # are adapted (SceneEntityCfg promotion + MDP twin swap). The event manager @@ -86,58 +94,6 @@ class WarpFrontend: {"observations", "rewards", "terminations", "actions", "events"} ) - # Root package that provides task-specific warp MDP twins. Imported (lazily, - # at swap time) so its task packages run their register_mdp_route() calls - # even when the caller only imported the stable task package. - _TWIN_PROVIDER_PACKAGE: ClassVar[str] = "isaaclab_tasks_experimental" - - # Shared fallback module holding generic warp twins. - _FALLBACK_MDP_MODULE: ClassVar[str] = "isaaclab_experimental.envs.mdp" - - # Stable package prefixes mapped to the warp MDP module that twins their - # terms. Populated via register_mdp_route(); never edited here. - _mdp_routes: ClassVar[dict[str, str]] = {} - - # ------------------------------------------------------------------ - # Route registry - # ------------------------------------------------------------------ - - @classmethod - def register_mdp_route(cls, stable_package: str, warp_mdp_module: str) -> None: - """Register the warp MDP module that twins a stable package's terms. - - Experimental task packages call this at import time to declare where - the warp twins of a stable task's MDP terms live. Twin lookup is by - symbol name on ``warp_mdp_module`` itself, mirroring how a stable cfg - consumes its ``mdp`` namespace — so the module should re-export - everything the task needs (task-specific twins plus the generic - forwards). - - Two kinds of keys are useful: - - * A *task package* (e.g. ``"isaaclab_tasks.core.reach"``) routes every - cfg class defined under it. This also covers terms whose stable - functions are defined in core/shared packages but overridden by - task-specific twins. - * An *MDP package* (e.g. ``"isaaclab_tasks.core.locomotion.humanoid.mdp"``) - routes symbols *defined* under it, which serves other tasks that - borrow those terms. - - Args: - stable_package: Stable package prefix to route (longest match wins). - warp_mdp_module: Warp MDP module providing the twins, e.g. - ``"isaaclab_tasks_experimental.core.cartpole.mdp"``. - - Raises: - ValueError: If ``stable_package`` is already routed to a different module. - """ - existing = cls._mdp_routes.get(stable_package) - if existing is not None and existing != warp_mdp_module: - raise ValueError( - f"MDP route conflict for {stable_package!r}: already routed to {existing!r}, got {warp_mdp_module!r}." - ) - cls._mdp_routes[stable_package] = warp_mdp_module - # ------------------------------------------------------------------ # Env construction # ------------------------------------------------------------------ @@ -298,8 +254,6 @@ def _swap_mdp(cls, cfg: Any, label: str) -> None: substitutes the callable; the manager reads the new func's annotations when it parses the term cfg. """ - cls._ensure_twin_providers_imported() - cfg_route_modules = cls._cfg_route_modules(cfg) module_cache: dict[str, list[ModuleType]] = {} searched: set[str] = set() @@ -409,44 +363,68 @@ def _assert_direct_warp_registration(cls, task_id: str) -> None: # ------------------------------------------------------------------ @classmethod - def _match_route(cls, module: str) -> str | None: - """Return the routed warp MDP module for ``module``, longest prefix wins.""" + def _mirror_module(cls, module: str) -> str | None: + """Mechanical mirror of a stable module path, or ``None`` if not mirrored.""" if not isinstance(module, str) or not module: return None - best_key: str | None = None - for stable_prefix in cls._mdp_routes: - if module == stable_prefix or module.startswith(f"{stable_prefix}."): - if best_key is None or len(stable_prefix) > len(best_key): - best_key = stable_prefix - return cls._mdp_routes[best_key] if best_key is not None else None + root, _, rest = module.partition(".") + mirror_root = cls.MIRROR_ROOTS.get(root) + if mirror_root is None: + return None + return f"{mirror_root}.{rest}" if rest else mirror_root @staticmethod - def _import_routed_module(target: str) -> ModuleType: - """Import a registered route target; a broken registration is a hard error.""" + def _nearest_mdp_module(mirrored: str) -> str | None: + """Deepest existing ``.mdp`` package on a mirrored module path. + + If the path already contains an ``.mdp`` segment, the search starts at + that boundary (``...velocity.mdp.rewards`` resolves to + ``...velocity.mdp``); otherwise the path is walked upward until a + mirrored ``.mdp`` package exists. Existence is probed with + :func:`importlib.util.find_spec`; a prefix that does not exist simply + ends that probe, so resolution is a pure function of the installed + package tree. + """ + parts = mirrored.split(".") + if "mdp" in parts: + parts = parts[: parts.index("mdp")] + for depth in range(len(parts), 0, -1): + candidate = ".".join([*parts[:depth], "mdp"]) + try: + if importlib.util.find_spec(candidate) is not None: + return candidate + except ModuleNotFoundError: + continue + return None + + @staticmethod + def _import_twin_module(target: str) -> ModuleType: + """Import a resolved twin module; a module that exists but breaks is a hard error.""" try: return importlib.import_module(target) except ImportError as exc: - raise FrontendIncompatibleError( - f"registered warp MDP route target {target!r} failed to import: {exc}" - ) from exc + raise FrontendIncompatibleError(f"warp twin module {target!r} exists but failed to import: {exc}") from exc @classmethod def _cfg_route_modules(cls, cfg: Any) -> list[ModuleType]: - """Warp MDP modules routed from the cfg's class hierarchy, subclass first. + """Warp MDP mirrors of the cfg's class hierarchy, subclass first. A stable cfg consumes terms through its task's ``mdp`` namespace, which re-exports symbols that may be *defined* in core or shared packages. The warp mirror of that namespace is therefore the primary place to resolve - twins, and it is found by routing the modules of ``type(cfg).__mro__`` — + twins; it is found by mirroring the modules of ``type(cfg).__mro__`` — the subclass module first (robot-specific cfgs), then base-cfg modules (task-family bases). """ modules: list[ModuleType] = [] for klass in type(cfg).__mro__: - target = cls._match_route(getattr(klass, "__module__", "") or "") + mirrored = cls._mirror_module(getattr(klass, "__module__", "") or "") + if mirrored is None: + continue + target = cls._nearest_mdp_module(mirrored) if target is None: continue - module = cls._import_routed_module(target) + module = cls._import_twin_module(target) if module not in modules: modules.append(module) return modules @@ -458,46 +436,21 @@ def _twin_modules(cls, symbol_module: str, cfg_route_modules: list[ModuleType]) 1. The warp mirrors of the cfg's own task MDP namespace (:meth:`_cfg_route_modules`) — task-specific twins win, including twins for symbols defined in core/shared packages. - 2. The warp mirror routed from the symbol's defining package — covers - terms borrowed from another task family's MDP package. - 3. The shared fallback module (where generic warp twins live). + 2. The mirror of the symbol's defining package — covers terms borrowed + from another task family and, for core-defined symbols, the shared + :mod:`isaaclab_experimental.envs.mdp` twins (the mirror of + :mod:`isaaclab.envs.mdp`). """ modules = list(cfg_route_modules) - target = cls._match_route(symbol_module) - if target is not None: - module = cls._import_routed_module(target) - if module not in modules: - modules.append(module) - try: - fallback_module = importlib.import_module(cls._FALLBACK_MDP_MODULE) - except ModuleNotFoundError as exc: - if exc.name != cls._FALLBACK_MDP_MODULE: - raise - logger.warning("frontend.warp: fallback mdp module %r not importable", cls._FALLBACK_MDP_MODULE) - else: - if fallback_module not in modules: - modules.append(fallback_module) + mirrored = cls._mirror_module(symbol_module) + if mirrored is not None: + target = cls._nearest_mdp_module(mirrored) + if target is not None: + module = cls._import_twin_module(target) + if module not in modules: + modules.append(module) return modules - @classmethod - def _ensure_twin_providers_imported(cls) -> None: - """Import the twin-provider package so its MDP routes are registered. - - Route registration happens as an import side effect of the experimental - task packages. A caller running ``--frontend=warp`` on a stable task id - has no other reason to import the provider package, so the swap - triggers it here. A missing provider package is not an error by itself - — the swap then fails with the explicit missing-twin report. - """ - try: - importlib.import_module(cls._TWIN_PROVIDER_PACKAGE) - except ModuleNotFoundError as exc: - # Only swallow "the provider package is not installed"; a genuine - # import error inside an existing package must surface. - if exc.name != cls._TWIN_PROVIDER_PACKAGE: - raise - logger.warning("frontend.warp: twin provider package %r is not installed", cls._TWIN_PROVIDER_PACKAGE) - @classmethod def _resolve_warp_twin(cls, name: str, modules: Iterable[ModuleType]) -> Any | None: """Return the same-named symbol from ``modules`` that originates under the warp packages.""" @@ -563,7 +516,3 @@ def _walk_terms(node: Any, path: tuple[str, ...] = ()) -> Iterator[tuple[tuple[s if name.startswith("_") or value is None: continue yield from WarpFrontend._walk_terms(value, path + (name,)) - - -# Module-level alias used by the task packages at import time. -register_mdp_route = WarpFrontend.register_mdp_route diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py index 859f0b346769..b0ad962082f1 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend.py +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -12,8 +12,8 @@ * :func:`_require_newton_physics` hard-check. * :func:`_walk_terms` recursive ManagerTermBaseCfg discovery over instance attributes (nested class objects are not descended into). -* :func:`register_mdp_route` / :func:`_match_route` route registry. -* :func:`_cfg_route_modules` cfg-hierarchy route resolution. +* :meth:`_mirror_module` / :meth:`_nearest_mdp_module` deterministic mirror + resolution and :meth:`_cfg_route_modules` cfg-hierarchy resolution. * :func:`_promote_scene_entity_cfgs` walks ``term.params`` dicts. * :func:`_swap_mdp` swaps ``func`` *and* ``class_type``; raises with a path list when twins are missing. @@ -25,7 +25,6 @@ from __future__ import annotations import contextlib -import importlib import types from typing import Any from unittest.mock import patch @@ -37,7 +36,6 @@ FrontendIncompatibleError, WarpFrontend, Workflow, - register_mdp_route, ) from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as WarpSceneEntityCfg from isaaclab_newton.physics import NewtonCfg @@ -132,14 +130,6 @@ def _term(func=None, params: dict | None = None) -> RewardTermCfg: return RewardTermCfg(func=func or _stable_func, weight=1.0, params=params or {}) -@pytest.fixture -def fake_routes(monkeypatch: pytest.MonkeyPatch) -> dict[str, str]: - """Isolate the WarpFrontend route registry for a test.""" - routes: dict[str, str] = {} - monkeypatch.setattr(fe.WarpFrontend, "_mdp_routes", routes) - return routes - - # ====================================================================== # Enums # ====================================================================== @@ -316,46 +306,52 @@ def test_walk_terms_ignores_nested_class_objects(): # ====================================================================== -# Route registry +# Mirror resolution # ====================================================================== -def test_register_mdp_route_and_longest_prefix_match(fake_routes: dict[str, str]): - register_mdp_route("isaaclab_tasks.core.locomotion", "warp.locomotion.mdp") - register_mdp_route("isaaclab_tasks.core.locomotion.humanoid", "warp.humanoid.mdp") - - # Longest registered prefix wins. - assert fe.WarpFrontend._match_route("isaaclab_tasks.core.locomotion.humanoid.mdp.rewards") == "warp.humanoid.mdp" - assert fe.WarpFrontend._match_route("isaaclab_tasks.core.locomotion.ant.mdp") == "warp.locomotion.mdp" - # Exact package match works too. - assert fe.WarpFrontend._match_route("isaaclab_tasks.core.locomotion.humanoid") == "warp.humanoid.mdp" - # Prefix matching is per package segment, not per character. - assert fe.WarpFrontend._match_route("isaaclab_tasks.core.locomotion_extras.mdp") is None - assert fe.WarpFrontend._match_route("isaaclab.envs.mdp.rewards") is None - - -def test_register_mdp_route_is_idempotent(fake_routes: dict[str, str]): - register_mdp_route("isaaclab_tasks.core.cartpole", "warp.cartpole.mdp") - register_mdp_route("isaaclab_tasks.core.cartpole", "warp.cartpole.mdp") # no raise - assert fake_routes == {"isaaclab_tasks.core.cartpole": "warp.cartpole.mdp"} - +def test_mirror_module_maps_stable_roots(): + assert ( + fe.WarpFrontend._mirror_module("isaaclab_tasks.core.velocity.mdp.rewards") + == "isaaclab_tasks_experimental.core.velocity.mdp.rewards" + ) + assert fe.WarpFrontend._mirror_module("isaaclab.envs.mdp.events") == "isaaclab_experimental.envs.mdp.events" + assert fe.WarpFrontend._mirror_module("isaaclab_tasks") == "isaaclab_tasks_experimental" + # Unmirrored roots are not resolvable. + assert fe.WarpFrontend._mirror_module("torch.nn.functional") is None + assert fe.WarpFrontend._mirror_module("") is None + + +def test_nearest_mdp_module_truncates_at_mdp_boundary(): + # A symbol module resolves to its own mdp package on the mirrored path. + assert ( + fe.WarpFrontend._nearest_mdp_module("isaaclab_tasks_experimental.core.locomotion.mdp.rewards") + == "isaaclab_tasks_experimental.core.locomotion.mdp" + ) + assert fe.WarpFrontend._nearest_mdp_module("isaaclab_experimental.envs.mdp.events") == ( + "isaaclab_experimental.envs.mdp" + ) -def test_register_mdp_route_rejects_conflicts(fake_routes: dict[str, str]): - register_mdp_route("isaaclab_tasks.core.cartpole", "warp.cartpole.mdp") - with pytest.raises(ValueError): - register_mdp_route("isaaclab_tasks.core.cartpole", "warp.other.mdp") +def test_nearest_mdp_module_walks_up_cfg_paths(): + # A cfg module path walks up to the task family's mdp package. + assert ( + fe.WarpFrontend._nearest_mdp_module("isaaclab_tasks_experimental.core.velocity.config.anymal_d.flat_env_cfg") + == "isaaclab_tasks_experimental.core.velocity.mdp" + ) + # Nothing on the path provides an mdp package. + assert fe.WarpFrontend._nearest_mdp_module("isaaclab_tasks_experimental.nonexistent.task_pkg") is None -def test_cfg_route_modules_resolves_cfg_class_module(fake_routes: dict[str, str]): - # Route this test module's own name so the fixture cfg class matches. - register_mdp_route(__name__, "math") - assert fe.WarpFrontend._cfg_route_modules(_CfgFixture()) == [importlib.import_module("math")] +def test_cfg_route_modules_resolves_mirrored_hierarchy(): + # A class whose module mirrors into the experimental tree resolves to the + # task family's mdp package; unmirrored classes contribute nothing. + fixture_cls = type("FakeVelocityCfg", (), {}) + fixture_cls.__module__ = "isaaclab_tasks.core.velocity.config.anymal_d.flat_env_cfg" + modules = fe.WarpFrontend._cfg_route_modules(fixture_cls()) + assert [m.__name__ for m in modules] == ["isaaclab_tasks_experimental.core.velocity.mdp"] -def test_cfg_route_modules_raises_on_broken_target(fake_routes: dict[str, str]): - register_mdp_route(__name__, "definitely_not_an_importable_module") - with pytest.raises(FrontendIncompatibleError): - fe.WarpFrontend._cfg_route_modules(_CfgFixture()) + assert fe.WarpFrontend._cfg_route_modules(_CfgFixture()) == [] # ====================================================================== @@ -378,7 +374,6 @@ def _fake_mdp_module(name_to_symbol: dict[str, Any]) -> _FakeMdpModule: def _patch_twin_modules(monkeypatch: pytest.MonkeyPatch, modules: list[Any]) -> None: """Pin twin resolution to ``modules`` and keep unit tests hermetic.""" - monkeypatch.setattr(fe.WarpFrontend, "_ensure_twin_providers_imported", classmethod(lambda cls: None)) monkeypatch.setattr(fe.WarpFrontend, "_cfg_route_modules", classmethod(lambda cls, cfg: [])) monkeypatch.setattr( fe.WarpFrontend, "_twin_modules", classmethod(lambda cls, symbol_module, cfg_route_modules: modules) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py index 9a0eaa256a77..e8bc71c28dc3 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py @@ -10,8 +10,3 @@ declared by ``Isaac-Cartpole-Direct``. There is no separate task registration: run the stable ids with ``--frontend warp`` and ``presets=newton_mjwarp``. """ - -from isaaclab_experimental.envs.frontend import register_mdp_route - -# Warp twins for the stable cartpole MDP terms live in this package's ``mdp``. -register_mdp_route("isaaclab_tasks.core.cartpole", f"{__name__}.mdp") diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/__init__.py index e37ca86717b5..3fba47b723b6 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/__init__.py @@ -3,4 +3,10 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Warp-first locomotion task implementations.""" +"""Warp implementations for the stable locomotion tasks. + +The ``mdp`` package holds the twins for the shared +:mod:`isaaclab_tasks.core.locomotion.mdp` terms used by the Humanoid and Ant +tasks. There is no separate task registration: run the stable ids with +``--frontend warp`` and ``presets=newton_mjwarp``. +""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/__init__.py index 277fea940e58..1be981faf7f0 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/__init__.py @@ -3,20 +3,12 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Warp implementations for the stable Ant tasks. +"""Warp implementation of the stable direct Ant task. -Holds the direct +The manager-based MDP twins live in the shared +:mod:`isaaclab_tasks_experimental.core.locomotion.mdp` package (mirroring the +stable layout). The stable ``Isaac-Ant-Direct`` registration declares :class:`~isaaclab_tasks_experimental.core.locomotion.ant.ant_env_warp.AntWarpEnv` -declared by ``Isaac-Ant-Direct``; the manager-based MDP twins are borrowed from -the Humanoid package. There is no separate task registration: run the stable -ids with ``--frontend warp`` and ``presets=newton_mjwarp``. +as its ``warp_entry_point``; run the stable ids with ``--frontend warp`` and +``presets=newton_mjwarp``. """ - -from isaaclab_experimental.envs.frontend import register_mdp_route - -# The stable Ant task borrows Humanoid MDP terms, so its warp twins live in the -# experimental humanoid package. -register_mdp_route( - "isaaclab_tasks.core.locomotion.ant", - "isaaclab_tasks_experimental.core.locomotion.humanoid.mdp", -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/__init__.py index 7f05dee5608e..b87d773f7b6d 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/__init__.py @@ -3,16 +3,12 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Warp implementations for the stable Humanoid tasks. +"""Warp implementation of the stable direct Humanoid task. -Holds the MDP twins for the manager-based ``Isaac-Humanoid`` task (also -serving the Ant task, which borrows them) and the direct +The manager-based MDP twins live in the shared +:mod:`isaaclab_tasks_experimental.core.locomotion.mdp` package (mirroring the +stable layout). The stable ``Isaac-Humanoid-Direct`` registration declares :class:`~isaaclab_tasks_experimental.core.locomotion.humanoid.humanoid_warp_env.HumanoidWarpEnv` -declared by ``Isaac-Humanoid-Direct``. There is no separate task registration: -run the stable ids with ``--frontend warp`` and ``presets=newton_mjwarp``. +as its ``warp_entry_point``; run the stable ids with ``--frontend warp`` and +``presets=newton_mjwarp``. """ - -from isaaclab_experimental.envs.frontend import register_mdp_route - -# Warp twins for the stable humanoid MDP terms live in this package's ``mdp``. -register_mdp_route("isaaclab_tasks.core.locomotion.humanoid", f"{__name__}.mdp") diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/__init__.pyi similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/__init__.pyi rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/__init__.pyi diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/observations.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/observations.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/observations.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/observations.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/rewards.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/mdp/rewards.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/rewards.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/__init__.py index 09abb3d8aafc..1a9be132962a 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/__init__.py @@ -10,8 +10,3 @@ and ``presets=newton_mjwarp``. The IK and OSC variants require controller action twins that do not exist yet. """ - -from isaaclab_experimental.envs.frontend import register_mdp_route - -# Warp twins for the stable reach MDP terms live in this package's ``mdp``. -register_mdp_route("isaaclab_tasks.core.reach", f"{__name__}.mdp") diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py index 436c1da089b6..b4d5fd3d214f 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py @@ -16,15 +16,9 @@ from typing import TYPE_CHECKING -from isaaclab_experimental.envs.frontend import register_mdp_route - if TYPE_CHECKING: from isaaclab.envs import ManagerBasedRLEnvCfg -# Warp twins for the stable velocity MDP terms live in this package's ``mdp``. -register_mdp_route("isaaclab_tasks.core.velocity", f"{__name__}.mdp") - - def disable_unsupported_randomization_events(cfg: ManagerBasedRLEnvCfg) -> None: """Disable stable randomization events that have no warp twins yet. From b510d31ea22eac7dc71dcd51994e7537387e21ff Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 16 Jul 2026 15:50:28 -0700 Subject: [PATCH 33/58] Run velocity tasks on warp via the frontend; drop the last variants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add warp adapters for the two stable startup randomization events (randomize_rigid_body_material, randomize_rigid_body_mass): the stable terms already dispatch to the active physics backend, so the adapters only convert the warp event manager's env-mask calling convention to the stable env-ids one and inherit the warp ManagerTermBase so the managers accept them as class terms. With those twins in place the stable flat velocity tasks adapt cleanly, so the 18 velocity *-Warp-v0 registrations and their delta configs are deleted — no manager-based warp registration remains, and the conversion test pins that end state. The Allegro reorient task id drops its stale -v0 suffix (Isaac-Reorient-Cube-Allegro-Direct-Warp). --- .../newton/warp-environments.rst | 30 ++---- docs/source/overview/environments.rst | 47 +--------- .../changelog.d/warp-manager-bridge.rst | 13 ++- .../isaaclab_experimental/envs/mdp/events.py | 93 +++++++++++++++++++ .../test/envs/test_frontend_cfg_conversion.py | 36 +++---- .../changelog.d/warp-manager-bridge.minor.rst | 20 ++-- .../reorient/config/allegro_hand/__init__.py | 2 +- .../core/velocity/__init__.py | 31 ++----- .../core/velocity/config/__init__.py | 9 -- .../core/velocity/config/a1/__init__.py | 37 -------- .../core/velocity/config/a1/flat_env_cfg.py | 35 ------- .../core/velocity/config/anymal_b/__init__.py | 37 -------- .../velocity/config/anymal_b/flat_env_cfg.py | 35 ------- .../core/velocity/config/anymal_c/__init__.py | 39 -------- .../velocity/config/anymal_c/flat_env_cfg.py | 35 ------- .../core/velocity/config/anymal_d/__init__.py | 35 ------- .../velocity/config/anymal_d/flat_env_cfg.py | 35 ------- .../core/velocity/config/cassie/__init__.py | 35 ------- .../velocity/config/cassie/flat_env_cfg.py | 35 ------- .../core/velocity/config/g1/__init__.py | 36 ------- .../core/velocity/config/g1/flat_env_cfg.py | 33 ------- .../core/velocity/config/go1/__init__.py | 35 ------- .../core/velocity/config/go1/flat_env_cfg.py | 35 ------- .../core/velocity/config/go2/__init__.py | 35 ------- .../core/velocity/config/go2/flat_env_cfg.py | 35 ------- .../core/velocity/config/h1/__init__.py | 35 ------- .../core/velocity/config/h1/flat_env_cfg.py | 33 ------- 27 files changed, 143 insertions(+), 773 deletions(-) delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/flat_env_cfg.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/flat_env_cfg.py diff --git a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst index f1c1aee930ad..f572c68c6e4f 100644 --- a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst +++ b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst @@ -43,7 +43,7 @@ environment class. Stable tasks with a declared warp implementation: Warp-native direct task with its own registration and configuration (the cube is modeled differently than in the stable task): -- ``Isaac-Reorient-Cube-Allegro-Direct-Warp-v0`` — Allegro hand cube reorient +- ``Isaac-Reorient-Cube-Allegro-Direct-Warp`` — Allegro hand cube reorient Manager-Based Warp Execution (``--frontend warp``) @@ -58,28 +58,14 @@ MDP term for its warp twin). Select the Newton solver explicitly with - ``Isaac-Cartpole`` - ``Isaac-Ant`` - ``Isaac-Humanoid`` -- ``Isaac-Reach-Franka`` -- ``Isaac-Reach-UR10`` +- ``Isaac-Reach-Franka``, ``Isaac-Reach-UR10`` +- ``Isaac-Velocity-Flat-AnymalD``, ``-Cassie``, ``-G1``, ``-H1``, ``-UnitreeGo2`` + (and their ``-Play`` variants) A missing twin is a hard error listing the affected terms, so a partially covered task fails at build time rather than silently changing behavior. - -**Locomotion (Flat) — registered warp variants** - -The velocity tasks still register ``*-Warp-v0`` variants because their stable -configurations contain randomization events without warp twins yet; the -variants reuse the stable flat configurations and only disable those events. -They also require ``presets=newton_mjwarp``. - -- ``Isaac-Velocity-Flat-AnymalB-Warp-v0`` -- ``Isaac-Velocity-Flat-AnymalC-Warp-v0`` -- ``Isaac-Velocity-Flat-AnymalD-Warp-v0`` -- ``Isaac-Velocity-Flat-Cassie-Warp-v0`` -- ``Isaac-Velocity-Flat-G1-Warp-v0`` -- ``Isaac-Velocity-Flat-H1-Warp-v0`` -- ``Isaac-Velocity-Flat-UnitreeA1-Warp-v0`` -- ``Isaac-Velocity-Flat-UnitreeGo1-Warp-v0`` -- ``Isaac-Velocity-Flat-UnitreeGo2-Warp-v0`` +Rough-terrain velocity tasks remain unsupported until +:class:`~isaaclab.terrains.TerrainImporter` gains Warp APIs. Quick Start @@ -95,9 +81,9 @@ Quick Start ./isaaclab.sh train --rl_library rsl_rl \ --task Isaac-Cartpole --frontend warp presets=newton_mjwarp --num_envs 4096 - # Manager-based workflow: registered warp variant + # Manager-based workflow: velocity task on the warp runtime ./isaaclab.sh train --rl_library rsl_rl \ - --task Isaac-Velocity-Flat-AnymalD-Warp-v0 presets=newton_mjwarp --num_envs 4096 + --task Isaac-Velocity-Flat-AnymalD --frontend warp presets=newton_mjwarp --num_envs 4096 All RL libraries with warp-compatible wrappers are supported: RSL-RL, RL Games, SKRL, and Stable-Baselines3. diff --git a/docs/source/overview/environments.rst b/docs/source/overview/environments.rst index 5568f77e2820..8819ddd50dbb 100644 --- a/docs/source/overview/environments.rst +++ b/docs/source/overview/environments.rst @@ -1231,7 +1231,7 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_mjwarp``, ``ovphysx``, ``physx`` - * - Isaac-Reorient-Cube-Allegro-Direct-Warp-v0 + * - Isaac-Reorient-Cube-Allegro-Direct-Warp - - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) @@ -1397,11 +1397,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Manager Based - **rsl_rl** (PPO) - **physics=** ``newton_mjwarp``, ``physx`` - * - Isaac-Velocity-Flat-AnymalB-Warp-v0 - - Isaac-Velocity-Flat-AnymalB-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO) - - * - IsaacContrib-Velocity-Flat-AnymalB - IsaacContrib-Velocity-Flat-AnymalB-Play - Manager Based @@ -1412,11 +1407,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Velocity-Flat-AnymalC-Warp-v0 - - Isaac-Velocity-Flat-AnymalC-Warp-Play-v0 - - Manager Based - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - * - IsaacContrib-Velocity-Flat-AnymalC - IsaacContrib-Velocity-Flat-AnymalC-Play - Manager Based @@ -1432,11 +1422,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Manager Based - **rsl_rl** (PPO, DISTILLATION, DISTILLATION_RECURRENT, RECURRENT), **skrl** (PPO) - **physics=** ``newton_mjwarp``, ``ovphysx``, ``physx`` - * - Isaac-Velocity-Flat-AnymalD-Warp-v0 - - Isaac-Velocity-Flat-AnymalD-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Velocity-Flat-Cassie - - Manager Based @@ -1447,11 +1432,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_mjwarp``, ``physx`` - * - Isaac-Velocity-Flat-Cassie-Warp-v0 - - Isaac-Velocity-Flat-Cassie-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Velocity-Flat-Digit - - Manager Based @@ -1472,11 +1452,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_mjwarp``, ``physx`` - * - Isaac-Velocity-Flat-G1-Warp-v0 - - Isaac-Velocity-Flat-G1-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Velocity-Flat-H1 - - Manager Based @@ -1487,11 +1462,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_mjwarp``, ``physx`` - * - Isaac-Velocity-Flat-H1-Warp-v0 - - Isaac-Velocity-Flat-H1-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Velocity-Flat-Spot - - Manager Based @@ -1502,21 +1472,11 @@ inferencing, including reading from an already trained checkpoint and disabling - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_mjwarp``, ``physx`` - * - Isaac-Velocity-Flat-UnitreeA1-Warp-v0 - - Isaac-Velocity-Flat-UnitreeA1-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) - - * - IsaacContrib-Velocity-Flat-UnitreeA1 - IsaacContrib-Velocity-Flat-UnitreeA1-Play - Manager Based - **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) - **physics=** ``newton_mjwarp``, ``physx`` - * - Isaac-Velocity-Flat-UnitreeGo1-Warp-v0 - - Isaac-Velocity-Flat-UnitreeGo1-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO) - - * - IsaacContrib-Velocity-Flat-UnitreeGo1 - IsaacContrib-Velocity-Flat-UnitreeGo1-Play - Manager Based @@ -1532,11 +1492,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_mjwarp``, ``physx`` - * - Isaac-Velocity-Flat-UnitreeGo2-Warp-v0 - - Isaac-Velocity-Flat-UnitreeGo2-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO) - - * - IsaacContrib-Velocity-Rough-AnymalB - IsaacContrib-Velocity-Rough-AnymalB-Play - Manager Based diff --git a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst index 1d56783ad0e9..2a28c88089be 100644 --- a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst +++ b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst @@ -7,12 +7,19 @@ Added * Added :mod:`isaaclab_experimental.envs.frontend` runtime selector and :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable` used by ``--frontend=warp`` to adapt stable cfgs onto the warp runtime. -* Added :func:`isaaclab_experimental.envs.frontend.register_mdp_route` so task - packages declare where their warp MDP twins live; the frontend holds no - per-task table. * Added ``warp_entry_point`` registration support: a stable direct task may declare its warp environment class, which ``--frontend=warp`` constructs with the stable configuration. +* Added warp adapters for the stable + :class:`~isaaclab.envs.mdp.events.randomize_rigid_body_material` and + :class:`~isaaclab.envs.mdp.events.randomize_rigid_body_mass` startup event + terms, so tasks using them run under ``--frontend=warp`` with randomization + active. +* Added deterministic warp twin resolution: a stable symbol's twin is looked + up on the mirrored experimental package tree (``isaaclab.* ↔ + isaaclab_experimental.*``, ``isaaclab_tasks.* ↔ + isaaclab_tasks_experimental.*``); missing twins are collected and reported + in one failure. No registration is needed. Changed ^^^^^^^ diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py index d9e61041a7ad..90e5696a6b10 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py @@ -26,13 +26,19 @@ from typing import TYPE_CHECKING +import torch import warp as wp +from isaaclab.envs.mdp.events import randomize_rigid_body_mass as _StableRandomizeRigidBodyMass +from isaaclab.envs.mdp.events import randomize_rigid_body_material as _StableRandomizeRigidBodyMaterial + +from isaaclab_experimental.managers import ManagerTermBase as _WarpManagerTermBase from isaaclab_experimental.managers import SceneEntityCfg from isaaclab_experimental.utils.warp import WarpCapturable if TYPE_CHECKING: from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedEnv # --------------------------------------------------------------------------- # Randomize rigid body center of mass @@ -585,3 +591,90 @@ def reset_joints_by_scale( # Sync derived buffers (_previous_joint_vel, joint_acc) for reset envs. asset.write_joint_position_to_sim_mask(position=asset.data.joint_pos.warp, env_mask=env_mask) asset.write_joint_velocity_to_sim_mask(velocity=asset.data.joint_vel.warp, env_mask=env_mask) + + +# --------------------------------------------------------------------------- +# Startup-mode randomization (mask-signature adapters) +# --------------------------------------------------------------------------- +# +# The stable material/mass randomization terms already dispatch to the active +# physics backend (Newton included); only their calling convention differs — +# the warp EventManager invokes terms with a Warp env-mask while the stable +# class terms expect torch env indices. Both terms run in ``startup`` mode, +# once at construction and outside any captured path, so a host-side +# mask-to-ids conversion is acceptable. + + +def _mask_to_env_ids(env_mask: wp.array) -> torch.Tensor: + """Convert a Warp boolean env-mask to the torch index tensor the stable terms expect.""" + return torch.nonzero(wp.to_torch(env_mask), as_tuple=False).squeeze(-1) + + +class randomize_rigid_body_material(_StableRandomizeRigidBodyMaterial, _WarpManagerTermBase): + """Warp adapter for the stable, backend-dispatched material randomization term. + + Converts the warp event manager's env-mask calling convention to the stable + term's env-ids convention and delegates. Startup mode only. Inherits the + warp :class:`ManagerTermBase` so the warp managers accept it as a class term. + """ + + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array, + static_friction_range: tuple[float, float], + dynamic_friction_range: tuple[float, float], + restitution_range: tuple[float, float], + num_buckets: int, + asset_cfg: SceneEntityCfg, + make_consistent: bool = False, + ) -> None: + super().__call__( + env, + _mask_to_env_ids(env_mask), + static_friction_range, + dynamic_friction_range, + restitution_range, + num_buckets, + asset_cfg, + make_consistent, + ) + + def reset(self, env_mask: wp.array | None = None) -> None: + """Nothing to reset; materials are randomized once at startup.""" + return + + +class randomize_rigid_body_mass(_StableRandomizeRigidBodyMass, _WarpManagerTermBase): + """Warp adapter for the stable, backend-dispatched mass randomization term. + + Converts the warp event manager's env-mask calling convention to the stable + term's env-ids convention and delegates. Startup mode only. Inherits the + warp :class:`ManagerTermBase` so the warp managers accept it as a class term. + """ + + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array, + asset_cfg: SceneEntityCfg, + mass_distribution_params: tuple[float, float], + operation: str, + distribution: str = "uniform", + recompute_inertia: bool = True, + min_mass: float = 1e-6, + ) -> None: + super().__call__( + env, + _mask_to_env_ids(env_mask), + asset_cfg, + mass_distribution_params, + operation, + distribution, + recompute_inertia, + min_mass, + ) + + def reset(self, env_mask: wp.array | None = None) -> None: + """Nothing to reset; masses are randomized once at startup.""" + return diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index f94ba002fb6c..fc308dd8ca63 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -45,6 +45,16 @@ "Isaac-Reach-Franka-Play", "Isaac-Reach-UR10", "Isaac-Reach-UR10-Play", + "Isaac-Velocity-Flat-AnymalD", + "Isaac-Velocity-Flat-AnymalD-Play", + "Isaac-Velocity-Flat-Cassie", + "Isaac-Velocity-Flat-Cassie-Play", + "Isaac-Velocity-Flat-G1", + "Isaac-Velocity-Flat-G1-Play", + "Isaac-Velocity-Flat-H1", + "Isaac-Velocity-Flat-H1-Play", + "Isaac-Velocity-Flat-UnitreeGo2", + "Isaac-Velocity-Flat-UnitreeGo2-Play", ] # Manager-based warp tasks are exactly those whose entry point is the shared warp @@ -84,33 +94,15 @@ def _cfg_entry_point(task_id: str) -> str: return cfg_entry -_MANAGER_WARP_TASKS = _manager_warp_tasks() - - -def test_manager_warp_tasks_are_registered(): - """Sanity: every expected velocity warp task variant is registered.""" - expected = { - f"Isaac-Velocity-Flat-{robot}-Warp{suffix}-v0" - for robot in ("AnymalB", "AnymalC", "AnymalD", "Cassie", "G1", "H1", "UnitreeA1", "UnitreeGo1", "UnitreeGo2") - for suffix in ("", "-Play") - } - registered = {task_id for task_id, _ in _MANAGER_WARP_TASKS} - assert registered == expected, f"missing: {sorted(expected - registered)}; extra: {sorted(registered - expected)}" +def test_no_manager_warp_variants_remain(): + """End-state pin: manager-based warp execution needs no parallel registrations.""" + assert _manager_warp_tasks() == [], "unexpected ManagerBasedRLEnvWarp registrations reappeared" @pytest.mark.parametrize("task_id", _STABLE_TASKS_WITH_WARP_COVERAGE, ids=_STABLE_TASKS_WITH_WARP_COVERAGE) def test_stable_task_cfg_adapts_to_warp(task_id: str): """Each covered stable task adapts without a missing twin (the --frontend warp path).""" - _load_adapted_cfg(_cfg_entry_point(task_id)) - - -@pytest.mark.parametrize( - "task_id, cfg_entry_point", - [pytest.param(task_id, cfg_entry, id=task_id) for task_id, cfg_entry in _MANAGER_WARP_TASKS], -) -def test_registered_warp_variant_cfg_adapts(task_id: str, cfg_entry_point: str): - """Each registered ``*-Warp-v0`` variant cfg adapts without a missing twin.""" - cfg = _load_adapted_cfg(cfg_entry_point) + cfg = _load_adapted_cfg(_cfg_entry_point(task_id)) # Action terms carry a ``class_type`` (not a ``func``) and live on a base that # is not a ManagerTermBaseCfg; guard that the adapter still swaps them to the diff --git a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst index 425d89c1651b..dd6acda10e2e 100644 --- a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst +++ b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst @@ -1,11 +1,9 @@ Changed ^^^^^^^ -* **Breaking:** Changed the manager-based velocity ``*-Warp-v0`` task variants - to reuse the stable flat configurations, disabling only the randomization - events that have no warp twins yet. The variants now require selecting the - Newton solver on the CLI via ``presets=newton_mjwarp`` instead of - hard-coding it in the configuration. +* **Breaking:** Renamed ``Isaac-Reorient-Cube-Allegro-Direct-Warp-v0`` to + ``Isaac-Reorient-Cube-Allegro-Direct-Warp`` to follow the stable id plus + ``-Warp`` suffix convention. * **Breaking:** Changed the package layout to mirror :mod:`isaaclab_tasks.core` (``isaaclab_tasks_experimental.core.``), replacing the previous ``manager_based``/``direct`` split. Update imports of the old paths to the @@ -14,12 +12,14 @@ Changed Removed ^^^^^^^ -* **Breaking:** Removed the duplicated manager-based warp task registrations +* **Breaking:** Removed all manager-based ``*-Warp-v0`` task registrations — ``Isaac-Cartpole-Warp-v0``, ``Isaac-Humanoid-Warp-v0``, ``Isaac-Ant-Warp-v0``, - ``Isaac-Reach-Franka-Warp-v0``, and ``Isaac-Reach-Franka-Warp-Play-v0`` - together with their duplicated environment configurations. Run the stable - task ids with ``--frontend warp`` and ``presets=newton_mjwarp`` instead, - e.g. ``--task Isaac-Cartpole --frontend warp presets=newton_mjwarp``. + ``Isaac-Reach-Franka-Warp-v0``, ``Isaac-Reach-Franka-Warp-Play-v0``, and the + velocity variants (``Isaac-Velocity-Flat--Warp-v0`` and their + ``-Warp-Play-v0`` forms) — together with their environment configurations. + Run the stable task ids with ``--frontend warp`` and + ``presets=newton_mjwarp`` instead, e.g. + ``--task Isaac-Velocity-Flat-AnymalD --frontend warp presets=newton_mjwarp``. * **Breaking:** Removed the duplicated direct warp task registrations ``Isaac-Cartpole-Direct-Warp-v0``, ``Isaac-Ant-Direct-Warp-v0``, and ``Isaac-Humanoid-Direct-Warp-v0`` together with their duplicated diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/__init__.py index 13ab818b132d..d3ccc4edc51b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/__init__.py @@ -17,7 +17,7 @@ stable_agents = "isaaclab_tasks.core.reorient.config.allegro_hand.agents" gym.register( - id="Isaac-Reorient-Cube-Allegro-Direct-Warp-v0", + id="Isaac-Reorient-Cube-Allegro-Direct-Warp", entry_point=f"{inhand_task_entry}.inhand_manipulation_warp_env:InHandManipulationWarpEnv", disable_env_checker=True, kwargs={ diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py index b4d5fd3d214f..f6510ad92c91 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py @@ -3,30 +3,11 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Velocity locomotion experimental task registrations (manager-based). +"""Warp MDP twins for the stable velocity locomotion tasks. -The per-robot ``*-Warp-v0`` variants reuse the stable flat cfgs and only -disable the randomization events that have no warp twins yet (see -:func:`disable_unsupported_randomization_events`). Once those twins exist, -the variants can be dropped in favor of ``--frontend warp`` on the stable -task ids. +There is no separate task registration: run the stable flat velocity ids +(e.g. ``Isaac-Velocity-Flat-AnymalD``) with ``--frontend warp`` and +``presets=newton_mjwarp``. Rough-terrain variants remain unsupported until +:class:`~isaaclab.terrains.TerrainImporter` gains Warp APIs (no +``height_scan`` twin). """ - -from __future__ import annotations - -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from isaaclab.envs import ManagerBasedRLEnvCfg - -def disable_unsupported_randomization_events(cfg: ManagerBasedRLEnvCfg) -> None: - """Disable stable randomization events that have no warp twins yet. - - The warp event manager invokes term functions with a Warp env mask, and no - warp twins exist yet for the rigid-body material/mass randomization events. - A stable term on a warp manager is a hard error at adaptation time, so the - warp task variants disable these terms until twins are available. - """ - for name in ("physics_material", "add_base_mass"): - if getattr(cfg.events, name, None) is not None: - setattr(cfg.events, name, None) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/__init__.py deleted file mode 100644 index 26f3257daef4..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# 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 - -"""Configurations for velocity-based locomotion environments.""" - -# We leave this file empty since we don't want to expose any configs in this package directly. -# We still need this file to import the "config" module in the parent package. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/__init__.py deleted file mode 100644 index 5b79532a456e..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/__init__.py +++ /dev/null @@ -1,37 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.contrib.velocity.config.a1 import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Velocity-Flat-UnitreeA1-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeA1FlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeA1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-UnitreeA1-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeA1FlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeA1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/flat_env_cfg.py deleted file mode 100644 index c9948d38a278..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/a1/flat_env_cfg.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -"""Warp variants of the stable Unitree A1 flat velocity task configuration.""" - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks.contrib.velocity.config.a1.flat_env_cfg import UnitreeA1FlatEnvCfg as _StableUnitreeA1FlatEnvCfg -from isaaclab_tasks.contrib.velocity.config.a1.flat_env_cfg import ( - UnitreeA1FlatEnvCfg_PLAY as _StableUnitreeA1FlatEnvCfg_PLAY, -) - -from isaaclab_tasks_experimental.core.velocity import ( - disable_unsupported_randomization_events, -) - - -@configclass -class UnitreeA1FlatEnvCfg(_StableUnitreeA1FlatEnvCfg): - """Stable flat cfg with the randomization events lacking warp twins disabled.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) - - -@configclass -class UnitreeA1FlatEnvCfg_PLAY(_StableUnitreeA1FlatEnvCfg_PLAY): - """Play variant of :class:`UnitreeA1FlatEnvCfg` for the warp runtime.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/__init__.py deleted file mode 100644 index a6753cfb95ef..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/__init__.py +++ /dev/null @@ -1,37 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.contrib.velocity.config.anymal_b import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Velocity-Flat-AnymalB-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalBFlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBFlatPPORunnerCfg", - "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBFlatPPORunnerWithSymmetryCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-AnymalB-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalBFlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBFlatPPORunnerCfg", - "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBFlatPPORunnerWithSymmetryCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/flat_env_cfg.py deleted file mode 100644 index 0c01243aa7d0..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_b/flat_env_cfg.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -"""Warp variants of the stable AnymalB flat velocity task configuration.""" - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks.contrib.velocity.config.anymal_b.flat_env_cfg import AnymalBFlatEnvCfg as _StableAnymalBFlatEnvCfg -from isaaclab_tasks.contrib.velocity.config.anymal_b.flat_env_cfg import ( - AnymalBFlatEnvCfg_PLAY as _StableAnymalBFlatEnvCfg_PLAY, -) - -from isaaclab_tasks_experimental.core.velocity import ( - disable_unsupported_randomization_events, -) - - -@configclass -class AnymalBFlatEnvCfg(_StableAnymalBFlatEnvCfg): - """Stable flat cfg with the randomization events lacking warp twins disabled.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) - - -@configclass -class AnymalBFlatEnvCfg_PLAY(_StableAnymalBFlatEnvCfg_PLAY): - """Play variant of :class:`AnymalBFlatEnvCfg` for the warp runtime.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/__init__.py deleted file mode 100644 index 9c9c2530e118..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/__init__.py +++ /dev/null @@ -1,39 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.contrib.velocity.config.anymal_c import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Velocity-Flat-AnymalC-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalCFlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerCfg", - "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerWithSymmetryCfg", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_flat_ppo_cfg.yaml", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-AnymalC-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalCFlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerCfg", - "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerWithSymmetryCfg", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_flat_ppo_cfg.yaml", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/flat_env_cfg.py deleted file mode 100644 index 7777d8274e6a..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_c/flat_env_cfg.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -"""Warp variants of the stable AnymalC flat velocity task configuration.""" - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks.contrib.velocity.config.anymal_c.flat_env_cfg import AnymalCFlatEnvCfg as _StableAnymalCFlatEnvCfg -from isaaclab_tasks.contrib.velocity.config.anymal_c.flat_env_cfg import ( - AnymalCFlatEnvCfg_PLAY as _StableAnymalCFlatEnvCfg_PLAY, -) - -from isaaclab_tasks_experimental.core.velocity import ( - disable_unsupported_randomization_events, -) - - -@configclass -class AnymalCFlatEnvCfg(_StableAnymalCFlatEnvCfg): - """Stable flat cfg with the randomization events lacking warp twins disabled.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) - - -@configclass -class AnymalCFlatEnvCfg_PLAY(_StableAnymalCFlatEnvCfg_PLAY): - """Play variant of :class:`AnymalCFlatEnvCfg` for the warp runtime.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/__init__.py deleted file mode 100644 index e54ed0b4d6ca..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/__init__.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.velocity.config.anymal_d import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Velocity-Flat-AnymalD-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalDFlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-AnymalD-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalDFlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/flat_env_cfg.py deleted file mode 100644 index ccd874878e2f..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/anymal_d/flat_env_cfg.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -"""Warp variants of the stable AnymalD flat velocity task configuration.""" - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks.core.velocity.config.anymal_d.flat_env_cfg import AnymalDFlatEnvCfg as _StableAnymalDFlatEnvCfg -from isaaclab_tasks.core.velocity.config.anymal_d.flat_env_cfg import ( - AnymalDFlatEnvCfg_PLAY as _StableAnymalDFlatEnvCfg_PLAY, -) - -from isaaclab_tasks_experimental.core.velocity import ( - disable_unsupported_randomization_events, -) - - -@configclass -class AnymalDFlatEnvCfg(_StableAnymalDFlatEnvCfg): - """Stable flat cfg with the randomization events lacking warp twins disabled.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) - - -@configclass -class AnymalDFlatEnvCfg_PLAY(_StableAnymalDFlatEnvCfg_PLAY): - """Play variant of :class:`AnymalDFlatEnvCfg` for the warp runtime.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/__init__.py deleted file mode 100644 index e3af41334541..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/__init__.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.velocity.config.cassie import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Velocity-Flat-Cassie-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:CassieFlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CassieFlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-Cassie-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:CassieFlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CassieFlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/flat_env_cfg.py deleted file mode 100644 index 37696377d675..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/cassie/flat_env_cfg.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -"""Warp variants of the stable Cassie flat velocity task configuration.""" - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks.core.velocity.config.cassie.flat_env_cfg import CassieFlatEnvCfg as _StableCassieFlatEnvCfg -from isaaclab_tasks.core.velocity.config.cassie.flat_env_cfg import ( - CassieFlatEnvCfg_PLAY as _StableCassieFlatEnvCfg_PLAY, -) - -from isaaclab_tasks_experimental.core.velocity import ( - disable_unsupported_randomization_events, -) - - -@configclass -class CassieFlatEnvCfg(_StableCassieFlatEnvCfg): - """Stable flat cfg with the randomization events lacking warp twins disabled.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) - - -@configclass -class CassieFlatEnvCfg_PLAY(_StableCassieFlatEnvCfg_PLAY): - """Play variant of :class:`CassieFlatEnvCfg` for the warp runtime.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/__init__.py deleted file mode 100644 index 47d6909f9356..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/__init__.py +++ /dev/null @@ -1,36 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.velocity.config.g1 import agents - -## -# Register Gym environments. -## - - -gym.register( - id="Isaac-Velocity-Flat-G1-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:G1FlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-G1-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:G1FlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/flat_env_cfg.py deleted file mode 100644 index e11424804874..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/g1/flat_env_cfg.py +++ /dev/null @@ -1,33 +0,0 @@ -# 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 - -"""Warp variants of the stable G1 flat velocity task configuration.""" - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks.core.velocity.config.g1.flat_env_cfg import G1FlatEnvCfg as _StableG1FlatEnvCfg -from isaaclab_tasks.core.velocity.config.g1.flat_env_cfg import G1FlatEnvCfg_PLAY as _StableG1FlatEnvCfg_PLAY - -from isaaclab_tasks_experimental.core.velocity import ( - disable_unsupported_randomization_events, -) - - -@configclass -class G1FlatEnvCfg(_StableG1FlatEnvCfg): - """Stable flat cfg with the randomization events lacking warp twins disabled.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) - - -@configclass -class G1FlatEnvCfg_PLAY(_StableG1FlatEnvCfg_PLAY): - """Play variant of :class:`G1FlatEnvCfg` for the warp runtime.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/__init__.py deleted file mode 100644 index 417865511f5f..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/__init__.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.contrib.velocity.config.go1 import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Velocity-Flat-UnitreeGo1-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo1FlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-UnitreeGo1-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo1FlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/flat_env_cfg.py deleted file mode 100644 index bbe78df0bd1b..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go1/flat_env_cfg.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -"""Warp variants of the stable Unitree Go1 flat velocity task configuration.""" - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks.contrib.velocity.config.go1.flat_env_cfg import UnitreeGo1FlatEnvCfg as _StableUnitreeGo1FlatEnvCfg -from isaaclab_tasks.contrib.velocity.config.go1.flat_env_cfg import ( - UnitreeGo1FlatEnvCfg_PLAY as _StableUnitreeGo1FlatEnvCfg_PLAY, -) - -from isaaclab_tasks_experimental.core.velocity import ( - disable_unsupported_randomization_events, -) - - -@configclass -class UnitreeGo1FlatEnvCfg(_StableUnitreeGo1FlatEnvCfg): - """Stable flat cfg with the randomization events lacking warp twins disabled.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) - - -@configclass -class UnitreeGo1FlatEnvCfg_PLAY(_StableUnitreeGo1FlatEnvCfg_PLAY): - """Play variant of :class:`UnitreeGo1FlatEnvCfg` for the warp runtime.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/__init__.py deleted file mode 100644 index 2d7f3b3212b8..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/__init__.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.velocity.config.go2 import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Velocity-Flat-UnitreeGo2-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo2FlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo2FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-UnitreeGo2-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo2FlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo2FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/flat_env_cfg.py deleted file mode 100644 index 87eb63e6e2fa..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/go2/flat_env_cfg.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -"""Warp variants of the stable Unitree Go2 flat velocity task configuration.""" - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks.core.velocity.config.go2.flat_env_cfg import UnitreeGo2FlatEnvCfg as _StableUnitreeGo2FlatEnvCfg -from isaaclab_tasks.core.velocity.config.go2.flat_env_cfg import ( - UnitreeGo2FlatEnvCfg_PLAY as _StableUnitreeGo2FlatEnvCfg_PLAY, -) - -from isaaclab_tasks_experimental.core.velocity import ( - disable_unsupported_randomization_events, -) - - -@configclass -class UnitreeGo2FlatEnvCfg(_StableUnitreeGo2FlatEnvCfg): - """Stable flat cfg with the randomization events lacking warp twins disabled.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) - - -@configclass -class UnitreeGo2FlatEnvCfg_PLAY(_StableUnitreeGo2FlatEnvCfg_PLAY): - """Play variant of :class:`UnitreeGo2FlatEnvCfg` for the warp runtime.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/__init__.py deleted file mode 100644 index 07deee5ee458..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/__init__.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.velocity.config.h1 import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Velocity-Flat-H1-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:H1FlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-H1-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:H1FlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/flat_env_cfg.py deleted file mode 100644 index 6c7ca5701261..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/config/h1/flat_env_cfg.py +++ /dev/null @@ -1,33 +0,0 @@ -# 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 - -"""Warp variants of the stable H1 flat velocity task configuration.""" - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks.core.velocity.config.h1.flat_env_cfg import H1FlatEnvCfg as _StableH1FlatEnvCfg -from isaaclab_tasks.core.velocity.config.h1.flat_env_cfg import H1FlatEnvCfg_PLAY as _StableH1FlatEnvCfg_PLAY - -from isaaclab_tasks_experimental.core.velocity import ( - disable_unsupported_randomization_events, -) - - -@configclass -class H1FlatEnvCfg(_StableH1FlatEnvCfg): - """Stable flat cfg with the randomization events lacking warp twins disabled.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) - - -@configclass -class H1FlatEnvCfg_PLAY(_StableH1FlatEnvCfg_PLAY): - """Play variant of :class:`H1FlatEnvCfg` for the warp runtime.""" - - def __post_init__(self): - super().__post_init__() - disable_unsupported_randomization_events(self) From ad86948fecc65b0aa14801557102adfc3c812481 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 16 Jul 2026 17:00:13 -0700 Subject: [PATCH 34/58] Annotate frontend and fix stale benchmark docs Add one-line comments at the dense decision points of WarpFrontend (twin lookup order, mirror walk, collect-all-misses reporting). Replace the removed -Warp-v0 ids in the benchmarking guide with a same-task A/B that differs only in --frontend, with identical Newton physics on both runs. Add the missing changelog entry for the RewardManager merging class-term reset() extras into its episode logs. --- .../newton/warp-environments.rst | 21 ++++++++---- .../changelog.d/warp-manager-bridge.rst | 4 +++ .../isaaclab_experimental/envs/frontend.py | 33 +++++++++++-------- 3 files changed, 37 insertions(+), 21 deletions(-) diff --git a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst index f572c68c6e4f..369b8c24b66e 100644 --- a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst +++ b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst @@ -273,25 +273,32 @@ to estimate the gain for your own task before committing to a migration. **Single-task A/B** +Run the same stable task id twice, differing only in ``--frontend``. Pass +``presets=newton_mjwarp`` to *both* runs so the physics backend is identical and the +measured difference isolates the frontend (manager pipeline + CUDA graph capture): + .. code-block:: bash - # Stable variant + # Torch frontend (stable managers) uv run isaaclab benchmark training \ --rl_library rsl_rl \ - --task -v0 \ + --task \ --num_envs 4096 \ --max_iterations 500 \ --benchmark_formatter summary \ - --output_path benchmarks/stable + --output_path benchmarks/torch \ + presets=newton_mjwarp - # Warp variant — same task with -Warp- suffix + # Warp frontend (warp managers, CUDA-graph captured) — same task id uv run isaaclab benchmark training \ --rl_library rsl_rl \ - --task -Warp-v0 \ + --task \ + --frontend warp \ --num_envs 4096 \ --max_iterations 500 \ --benchmark_formatter summary \ - --output_path benchmarks/warp + --output_path benchmarks/warp \ + presets=newton_mjwarp The ``summary`` formatter prints step time (min / mean / max) and total throughput. Compare "step time" between the two runs to estimate the gain per env step. @@ -299,7 +306,7 @@ The ``summary`` formatter prints step time (min / mean / max) and total throughp **Sweep across all available tasks** Run ``isaaclab benchmark training`` for each task in the stable set (cartpole, ant, humanoid, -locomotion, manipulation) and again with the ``-Warp-`` suffixed task ids, then diff the two output +locomotion, manipulation) and again with ``--frontend warp``, then diff the two output directories. **What to look at in the output** diff --git a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst index 2a28c88089be..de5082942259 100644 --- a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst +++ b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst @@ -29,6 +29,10 @@ Changed :meth:`~isaaclab_experimental.envs.frontend.WarpFrontend.adapt_cfg`, so registered warp task variants can derive from stable configurations instead of duplicating them. +* Changed the experimental warp ``RewardManager`` to merge dictionaries + returned by class-based reward term ``reset()`` into its episode-log extras + (persistent Warp buffer views, CUDA-graph safe); used to report + ``Metrics/success_rate`` under ``--frontend warp``. Fixed ^^^^^ diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index 29d14191a9bd..04d9d1886969 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -140,8 +140,10 @@ def _build_direct_env(cls, env_cfg: Any, task_id: str, **construct_kwargs: Any) """ env_class = cls._resolve_direct_warp_class(task_id) if env_class is None: + # No warp_entry_point: the task itself must be a warp-native registration. cls._assert_direct_warp_registration(task_id) return gym.make(task_id, cfg=env_cfg, **construct_kwargs) + # Declared warp class + stable cfg: swap only the env class. cls._require_newton_physics(env_cfg, type(env_cfg).__name__) return env_class(cfg=env_cfg, **construct_kwargs) @@ -215,11 +217,12 @@ def _promote_scene_entity_cfgs(cls, cfg: Any) -> None: promoted: list[str] = [] for path, term in cls._walk_terms(cfg): if not path or path[0] not in cls.WARP_MANAGED_GROUPS: - continue + continue # term runs on a stable manager; keep the stable entity params = getattr(term, "params", None) if not isinstance(params, dict): continue for key, value in list(params.items()): + # Skip non-entities and already-promoted values (idempotence). if isinstance(value, _WarpSceneEntityCfg) or not isinstance(value, _StableSceneEntityCfg): continue params[key] = _WarpSceneEntityCfg.from_stable(value) @@ -254,17 +257,18 @@ def _swap_mdp(cls, cfg: Any, label: str) -> None: substitutes the callable; the manager reads the new func's annotations when it parses the term cfg. """ + # Highest-priority twin source: warp mirrors of the cfg's own task namespace. cfg_route_modules = cls._cfg_route_modules(cfg) - module_cache: dict[str, list[ModuleType]] = {} - searched: set[str] = set() + module_cache: dict[str, list[ModuleType]] = {} # defining module -> modules to search + searched: set[str] = set() # module names, for the error message swapped = 0 missing: list[tuple[str, str, str]] = [] # (location, attr, symbol) for path, term in cls._walk_terms(cfg): if not path or path[0] not in cls.WARP_MANAGED_GROUPS: - continue + continue # term runs on a stable manager; leave it stable location = ".".join(path) - for attr in ("func", "class_type"): + for attr in ("func", "class_type"): # a term implements via either attr stable = getattr(term, attr, None) if stable is None or not cls._is_swap_candidate(stable): continue @@ -274,7 +278,7 @@ def _swap_mdp(cls, cfg: Any, label: str) -> None: searched.update(m.__name__ for m in module_cache[origin]) twin = cls._resolve_warp_twin(stable.__name__, module_cache[origin]) if twin is None: - missing.append((location, attr, stable.__name__)) + missing.append((location, attr, stable.__name__)) # collect every miss; report once below continue setattr(term, attr, twin) swapped += 1 @@ -367,7 +371,7 @@ def _mirror_module(cls, module: str) -> str | None: """Mechanical mirror of a stable module path, or ``None`` if not mirrored.""" if not isinstance(module, str) or not module: return None - root, _, rest = module.partition(".") + root, _, rest = module.partition(".") # only the top-level package root is swapped mirror_root = cls.MIRROR_ROOTS.get(root) if mirror_root is None: return None @@ -387,8 +391,8 @@ def _nearest_mdp_module(mirrored: str) -> str | None: """ parts = mirrored.split(".") if "mdp" in parts: - parts = parts[: parts.index("mdp")] - for depth in range(len(parts), 0, -1): + parts = parts[: parts.index("mdp")] # start the walk at the .mdp boundary + for depth in range(len(parts), 0, -1): # deepest prefix first candidate = ".".join([*parts[:depth], "mdp"]) try: if importlib.util.find_spec(candidate) is not None: @@ -417,7 +421,7 @@ def _cfg_route_modules(cls, cfg: Any) -> list[ModuleType]: (task-family bases). """ modules: list[ModuleType] = [] - for klass in type(cfg).__mro__: + for klass in type(cfg).__mro__: # subclass first: robot cfg wins over task-family base mirrored = cls._mirror_module(getattr(klass, "__module__", "") or "") if mirrored is None: continue @@ -442,6 +446,7 @@ def _twin_modules(cls, symbol_module: str, cfg_route_modules: list[ModuleType]) :mod:`isaaclab.envs.mdp`). """ modules = list(cfg_route_modules) + # Fallback: the mirror of the package that defines the symbol. mirrored = cls._mirror_module(symbol_module) if mirrored is not None: target = cls._nearest_mdp_module(mirrored) @@ -459,7 +464,7 @@ def _resolve_warp_twin(cls, name: str, modules: Iterable[ModuleType]) -> Any | N if candidate is None: continue origin = getattr(candidate, "__module__", "") or "" - if origin.startswith(cls.WARP_ROOT_PREFIXES): + if origin.startswith(cls.WARP_ROOT_PREFIXES): # re-exported stable symbols don't count return candidate return None @@ -508,11 +513,11 @@ def _walk_terms(node: Any, path: tuple[str, ...] = ()) -> Iterator[tuple[tuple[s from isaaclab.managers.manager_term_cfg import ActionTermCfg, ManagerTermBaseCfg if isinstance(node, (ManagerTermBaseCfg, ActionTermCfg)): - yield path, node + yield path, node # a term: yield and stop; never descend into params/func return if not hasattr(node, "__dataclass_fields__"): - return - for name, value in vars(node).items(): + return # plain data / callables: not a cfg subtree + for name, value in vars(node).items(): # instance attrs — the same view the managers iterate if name.startswith("_") or value is None: continue yield from WarpFrontend._walk_terms(value, path + (name,)) From fc37fa65ad95d1ac6e3f32e721329ef715899d37 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 16 Jul 2026 17:34:39 -0700 Subject: [PATCH 35/58] Fold Allegro warp reorient into the stable registration The stable cfg already resolves its cube to an articulation under the newton_mjwarp preset, so the duplicated warp cfg diverged only by pinning ls_iterations=15; the stable preset values now apply. Isaac-Reorient-Cube-Allegro-Direct declares the warp env class via warp_entry_point; the separate -Warp registration and its cfg package are deleted. A new end-state pin asserts no -Warp task ids remain, and the frontend stub-task fixture now unregisters its ids so later test modules see the real registry. --- .../newton/warp-environments.rst | 6 +- docs/source/overview/environments.rst | 5 - .../test/envs/test_frontend.py | 5 + .../test/envs/test_frontend_cfg_conversion.py | 13 +- .../changelog.d/warp-manager-bridge.rst | 7 +- .../reorient/config/allegro_hand/__init__.py | 3 + .../changelog.d/warp-manager-bridge.minor.rst | 14 +- .../core/reorient/config/__init__.py | 6 - .../reorient/config/allegro_hand/__init__.py | 29 ---- .../allegro_hand/allegro_hand_warp_env_cfg.py | 135 ------------------ .../reorient/inhand_manipulation_warp_env.py | 10 +- 11 files changed, 35 insertions(+), 198 deletions(-) delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/__init__.py delete mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/allegro_hand_warp_env_cfg.py diff --git a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst index 369b8c24b66e..988fe7ac561a 100644 --- a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst +++ b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst @@ -39,11 +39,7 @@ environment class. Stable tasks with a declared warp implementation: - ``Isaac-Cartpole-Direct`` — Cartpole balance - ``Isaac-Ant-Direct`` — Ant locomotion - ``Isaac-Humanoid-Direct`` — Humanoid locomotion - -Warp-native direct task with its own registration and configuration (the cube -is modeled differently than in the stable task): - -- ``Isaac-Reorient-Cube-Allegro-Direct-Warp`` — Allegro hand cube reorient +- ``Isaac-Reorient-Cube-Allegro-Direct`` — Allegro hand cube reorient Manager-Based Warp Execution (``--frontend warp``) diff --git a/docs/source/overview/environments.rst b/docs/source/overview/environments.rst index 8819ddd50dbb..dc17f390c2cf 100644 --- a/docs/source/overview/environments.rst +++ b/docs/source/overview/environments.rst @@ -1231,11 +1231,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_mjwarp``, ``ovphysx``, ``physx`` - * - Isaac-Reorient-Cube-Allegro-Direct-Warp - - - - Direct - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Reorient-Cube-Allegro-Play - - Manager Based diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py index b0ad962082f1..a5bbedd02214 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend.py +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -531,11 +531,16 @@ def _register_direct_stub_tasks(): # are never invoked (the guard only inspects them), so a fake import # path is fine. ``warp_entry_point`` targets are imported, so the valid # stub points at a real stdlib symbol. + registered = [] for task_id, (entry_point, warp_entry_point) in _DIRECT_TEST_TASKS.items(): kwargs = {"warp_entry_point": warp_entry_point} if warp_entry_point else {} with contextlib.suppress(gym.error.Error): gym.register(id=task_id, entry_point=entry_point, disable_env_checker=True, kwargs=kwargs) + registered.append(task_id) yield + # Unregister so later test modules see the real, stub-free registry. + for task_id in registered: + gym.registry.pop(task_id, None) def test_direct_guard_accepts_warp_rooted(): diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index fc308dd8ca63..ea534f3a4b33 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -99,6 +99,12 @@ def test_no_manager_warp_variants_remain(): assert _manager_warp_tasks() == [], "unexpected ManagerBasedRLEnvWarp registrations reappeared" +def test_no_warp_task_registrations_remain(): + """End-state pin: warp execution needs no parallel ``-Warp`` task ids at all.""" + warp_ids = sorted(task_id for task_id in gym.registry if "-Warp" in task_id) + assert warp_ids == [], f"unexpected warp task registrations: {warp_ids}" + + @pytest.mark.parametrize("task_id", _STABLE_TASKS_WITH_WARP_COVERAGE, ids=_STABLE_TASKS_WITH_WARP_COVERAGE) def test_stable_task_cfg_adapts_to_warp(task_id: str): """Each covered stable task adapts without a missing twin (the --frontend warp path).""" @@ -134,7 +140,12 @@ def _direct_tasks_with_warp_entry_point() -> list[tuple[str, str]]: def test_direct_tasks_declare_expected_warp_entry_points(): """Sanity: the stable direct tasks with warp implementations declare them.""" declared = {task_id for task_id, _ in _DIRECT_WARP_TASKS} - assert {"Isaac-Cartpole-Direct", "Isaac-Ant-Direct", "Isaac-Humanoid-Direct"} <= declared + assert { + "Isaac-Cartpole-Direct", + "Isaac-Ant-Direct", + "Isaac-Humanoid-Direct", + "Isaac-Reorient-Cube-Allegro-Direct", + } <= declared @pytest.mark.parametrize( diff --git a/source/isaaclab_tasks/changelog.d/warp-manager-bridge.rst b/source/isaaclab_tasks/changelog.d/warp-manager-bridge.rst index f34db06907a5..7a9bdd89d0b5 100644 --- a/source/isaaclab_tasks/changelog.d/warp-manager-bridge.rst +++ b/source/isaaclab_tasks/changelog.d/warp-manager-bridge.rst @@ -2,6 +2,7 @@ Added ^^^^^ * Added ``warp_entry_point`` declarations to the ``Isaac-Cartpole-Direct``, - ``Isaac-Ant-Direct``, and ``Isaac-Humanoid-Direct`` registrations so - ``--frontend warp`` constructs the Warp implementation of these tasks with - the stable task configuration. + ``Isaac-Ant-Direct``, ``Isaac-Humanoid-Direct``, and + ``Isaac-Reorient-Cube-Allegro-Direct`` registrations so ``--frontend warp`` + constructs the Warp implementation of these tasks with the stable task + configuration. diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/__init__.py index c5ec71b59cfe..35393451e9e1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/__init__.py @@ -47,6 +47,9 @@ disable_env_checker=True, kwargs={ "env_cfg_entry_point": f"{__name__}.allegro_hand_direct_env_cfg:AllegroHandEnvCfg", + "warp_entry_point": ( + "isaaclab_tasks_experimental.core.reorient.inhand_manipulation_warp_env:InHandManipulationWarpEnv" + ), "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_direct_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AllegroHandPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_direct_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst index dd6acda10e2e..f29f10f3e6ab 100644 --- a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst +++ b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst @@ -1,9 +1,6 @@ Changed ^^^^^^^ -* **Breaking:** Renamed ``Isaac-Reorient-Cube-Allegro-Direct-Warp-v0`` to - ``Isaac-Reorient-Cube-Allegro-Direct-Warp`` to follow the stable id plus - ``-Warp`` suffix convention. * **Breaking:** Changed the package layout to mirror :mod:`isaaclab_tasks.core` (``isaaclab_tasks_experimental.core.``), replacing the previous ``manager_based``/``direct`` split. Update imports of the old paths to the @@ -21,11 +18,12 @@ Removed ``presets=newton_mjwarp`` instead, e.g. ``--task Isaac-Velocity-Flat-AnymalD --frontend warp presets=newton_mjwarp``. * **Breaking:** Removed the duplicated direct warp task registrations - ``Isaac-Cartpole-Direct-Warp-v0``, ``Isaac-Ant-Direct-Warp-v0``, and - ``Isaac-Humanoid-Direct-Warp-v0`` together with their duplicated - environment configurations; the stable registrations declare the warp - environment classes via ``warp_entry_point``. Run the stable task ids with - ``--frontend warp`` and ``presets=newton_mjwarp`` instead, e.g. + ``Isaac-Cartpole-Direct-Warp-v0``, ``Isaac-Ant-Direct-Warp-v0``, + ``Isaac-Humanoid-Direct-Warp-v0``, and + ``Isaac-Reorient-Cube-Allegro-Direct-Warp-v0`` together with their + duplicated environment configurations; the stable registrations declare the + warp environment classes via ``warp_entry_point``. Run the stable task ids + with ``--frontend warp`` and ``presets=newton_mjwarp`` instead, e.g. ``--task Isaac-Cartpole-Direct --frontend warp presets=newton_mjwarp``. * Removed the unregistered rough velocity warp configurations; rough-terrain warp tasks remain unsupported until :class:`~isaaclab.terrains.TerrainImporter` diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/__init__.py deleted file mode 100644 index 771b55fc3605..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# 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 - -"""Warp task registrations for reorientation robots.""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/__init__.py deleted file mode 100644 index d3ccc4edc51b..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -# 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 - -""" -Allegro Inhand Manipulation environment. -""" - -import gymnasium as gym - -## -# Register Gym environments. -## - -inhand_task_entry = "isaaclab_tasks_experimental.core.reorient" -stable_agents = "isaaclab_tasks.core.reorient.config.allegro_hand.agents" - -gym.register( - id="Isaac-Reorient-Cube-Allegro-Direct-Warp", - entry_point=f"{inhand_task_entry}.inhand_manipulation_warp_env:InHandManipulationWarpEnv", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.allegro_hand_warp_env_cfg:AllegroHandWarpEnvCfg", - "rl_games_cfg_entry_point": f"{stable_agents}:rl_games_direct_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{stable_agents}.rsl_rl_ppo_cfg:AllegroHandPPORunnerCfg", - "skrl_cfg_entry_point": f"{stable_agents}:skrl_direct_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/allegro_hand_warp_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/allegro_hand_warp_env_cfg.py deleted file mode 100644 index 573b4f1d302d..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/allegro_hand/allegro_hand_warp_env_cfg.py +++ /dev/null @@ -1,135 +0,0 @@ -# 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 - - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg # , RigidObjectCfg -from isaaclab.envs import DirectRLEnvCfg -from isaaclab.markers import VisualizationMarkersCfg -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.sim.spawners.materials.physics_materials_cfg import RigidBodyMaterialCfg -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR -from isaaclab.utils.configclass import configclass - -from isaaclab_assets.robots.allegro import ALLEGRO_HAND_CFG - - -@configclass -class AllegroHandWarpEnvCfg(DirectRLEnvCfg): - # env - decimation = 4 - episode_length_s = 10.0 - action_space = 16 - observation_space = 124 # (full) - state_space = 0 - asymmetric_obs = False - obs_type = "full" - - solver_cfg = MJWarpSolverCfg( - solver="newton", - integrator="implicitfast", - njmax=80, - nconmax=70, - impratio=10.0, - cone="elliptic", - update_data_interval=2, - iterations=100, - ls_iterations=15, - # save_to_mjcf="AllegroHand.xml", - ) - - newton_cfg = NewtonCfg( - solver_cfg=solver_cfg, - num_substeps=2, - debug_mode=False, - ) - # simulation - sim: SimulationCfg = SimulationCfg( - dt=1 / 120, - render_interval=decimation, - physics_material=RigidBodyMaterialCfg( - static_friction=1.0, - dynamic_friction=1.0, - ), - physics=newton_cfg, - ) - # robot - robot_cfg: ArticulationCfg = ALLEGRO_HAND_CFG.replace(prim_path="/World/envs/env_.*/Robot") - - actuated_joint_names = [ - "index_joint_0", - "middle_joint_0", - "ring_joint_0", - "thumb_joint_0", - "index_joint_1", - "index_joint_2", - "index_joint_3", - "middle_joint_1", - "middle_joint_2", - "middle_joint_3", - "ring_joint_1", - "ring_joint_2", - "ring_joint_3", - "thumb_joint_1", - "thumb_joint_2", - "thumb_joint_3", - ] - fingertip_body_names = [ - "index_link_3", - "middle_link_3", - "ring_link_3", - "thumb_link_3", - ] - - # in-hand object - object_cfg: ArticulationCfg = ArticulationCfg( - prim_path="/World/envs/env_.*/object", - spawn=sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - mass_props=sim_utils.MassPropertiesCfg(density=400.0), - scale=(1.2, 1.2, 1.2), - ), - # FIXME: it does seem to be a bug for ArticulationCfg for handling empty joint list - init_state=ArticulationCfg.InitialStateCfg( - pos=(0.0, -0.17, 0.565), rot=(0.0, 0.0, 0.0, 1.0), joint_pos={}, joint_vel={} - ), - actuators={}, - articulation_root_prim_path="", - ) - # goal object - goal_object_cfg: VisualizationMarkersCfg = VisualizationMarkersCfg( - prim_path="/Visuals/goal_marker", - markers={ - "goal": sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - scale=(1.2, 1.2, 1.2), - ) - }, - ) - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg( - num_envs=8192, env_spacing=0.75, replicate_physics=True, clone_in_fabric=True - ) - # reset - reset_position_noise = 0.01 # range of position at reset - reset_dof_pos_noise = 0.2 # range of dof pos at reset - reset_dof_vel_noise = 0.0 # range of dof vel at reset - # reward scales - dist_reward_scale = -10.0 - rot_reward_scale = 1.0 - rot_eps = 0.1 - action_penalty_scale = -0.0002 - reach_goal_bonus = 250 - fall_penalty = 0 - fall_dist = 0.24 - vel_obs_scale = 0.2 - success_tolerance = 0.2 - max_consecutive_success = 0 - av_factor = 0.1 - act_moving_average = 1.0 - force_torque_obs_scale = 10.0 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/inhand_manipulation_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/inhand_manipulation_warp_env.py index ae8a8b30222c..e0deee854208 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/inhand_manipulation_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/inhand_manipulation_warp_env.py @@ -20,9 +20,7 @@ from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane if TYPE_CHECKING: - from isaaclab_tasks_experimental.core.reorient.config.allegro_hand.allegro_hand_warp_env_cfg import ( - AllegroHandWarpEnvCfg, - ) + from isaaclab_tasks.core.reorient.config.allegro_hand.allegro_hand_direct_env_cfg import AllegroHandEnvCfg @wp.kernel @@ -540,10 +538,10 @@ def rotation_distance(object_rot: wp.quatf, target_rot: wp.quatf) -> wp.float32: class InHandManipulationWarpEnv(DirectRLEnvWarp): - cfg: AllegroHandWarpEnvCfg # | ShadowHandWarpEnvCfg + cfg: AllegroHandEnvCfg # | ShadowHandEnvCfg - # def __init__(self, cfg: AllegroHandWarpEnvCfg | ShadowHandWarpEnvCfg, render_mode: str | None = None, **kwargs): - def __init__(self, cfg: AllegroHandWarpEnvCfg, render_mode: str | None = None, **kwargs): + # def __init__(self, cfg: AllegroHandEnvCfg | ShadowHandEnvCfg, render_mode: str | None = None, **kwargs): + def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwargs): super().__init__(cfg, render_mode, **kwargs) # --------------------------------------------------------------------- From 1216df95d6f1ee3e1372ae358b89d4190e3f2dd4 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Sun, 19 Jul 2026 03:20:44 -0700 Subject: [PATCH 36/58] Make Warp frontend mask-first Keep reset selection and registered manager hot paths on Warp-owned buffers, with explicit host boundaries only for legacy consumers. Default stateful frontend stages to eager execution so record/replay and capture can be layered on after correctness is established. --- .../jichuanh-warp-frontend-cleanup.minor.rst | 11 + .../isaaclab/envs/mdp/commands/_debug_vis.py | 73 ++++ .../envs/mdp/commands/pose_command.py | 35 +- .../envs/mdp/commands/velocity_command.py | 62 +-- .../isaaclab/terrains/terrain_importer.py | 118 ++++++ .../jichuanh-warp-frontend-cleanup.minor.rst | 28 ++ .../envs/direct_rl_env_warp.py | 47 +-- .../envs/interactive_scene_warp.py | 51 ++- .../envs/manager_based_env_warp.py | 164 +++++--- .../envs/manager_based_rl_env_warp.py | 369 +++++++++--------- .../envs/mdp/__init__.pyi | 6 + .../envs/mdp/commands/__init__.py | 10 + .../envs/mdp/commands/__init__.pyi | 15 + .../envs/mdp/commands/commands_cfg.py | 32 ++ .../envs/mdp/commands/pose_command.py | 206 ++++++++++ .../envs/mdp/commands/velocity_command.py | 299 ++++++++++++++ .../isaaclab_experimental/envs/mdp/events.py | 297 +++++++++----- .../envs/mdp/observations.py | 15 +- .../isaaclab_experimental/envs/mdp/rewards.py | 25 +- .../managers/__init__.pyi | 2 + .../managers/action_manager.py | 37 +- .../managers/command_manager.py | 79 ++-- .../managers/curriculum_manager.py | 245 ++++++++++++ .../managers/event_manager.py | 10 +- .../managers/manager_base.py | 45 ++- .../managers/manager_term_cfg.py | 14 + .../managers/observation_manager.py | 10 +- .../managers/reward_manager.py | 54 ++- .../managers/termination_manager.py | 42 +- .../utils/manager_call_switch.py | 50 ++- .../utils/warp/__init__.py | 12 +- .../utils/warp/kernels.py | 30 ++ .../isaaclab_experimental/utils/warp/utils.py | 88 ----- .../utils/warp_graph_cache.py | 14 +- .../test/envs/mdp/commands/test_commands.py | 302 ++++++++++++++ .../test/envs/mdp/parity_helpers.py | 4 + .../test/envs/mdp/test_events_warp_parity.py | 270 +++++++++++-- .../envs/mdp/test_observations_warp_parity.py | 9 - .../test/envs/mdp/test_rewards_warp_parity.py | 2 +- .../test/envs/test_interactive_scene_warp.py | 64 +++ .../envs/test_manager_based_rl_env_warp.py | 192 +++++++++ .../test/managers/test_curriculum_manager.py | 322 +++++++++++++++ .../test/managers/test_manager_base.py | 23 ++ .../test/managers/test_reward_manager.py | 41 ++ .../test/managers/test_termination_manager.py | 37 ++ .../test/utils/test_manager_call_switch.py | 40 +- .../test/utils/test_warp_graph_cache.py | 14 + .../jichuanh-warp-frontend-cleanup.minor.rst | 5 + .../isaaclab_newton/actuators/adapter.py | 52 ++- .../assets/articulation/articulation.py | 10 +- .../assets/rigid_object/rigid_object.py | 4 +- .../rigid_object_collection.py | 4 +- .../test/actuators/test_adapter_mask_reset.py | 65 +++ .../test_newton_mask_reset_forwarding.py | 63 +++ .../jichuanh-warp-frontend-cleanup.rst | 5 + .../assets/articulation/articulation.py | 12 +- .../assets/rigid_object/rigid_object.py | 4 +- .../rigid_object_collection.py | 6 +- .../test_physx_mask_reset_forwarding.py | 70 ++++ .../jichuanh-warp-frontend-cleanup.rst | 4 + .../dexsuite/mdp/commands/pose_commands.py | 2 +- .../jichuanh-warp-frontend-cleanup.minor.rst | 5 + .../inhand_manipulation_warp_env.py | 3 +- .../locomotion/velocity/mdp/__init__.pyi | 3 +- .../locomotion/velocity/mdp/curriculums.py | 121 +++++- .../locomotion/velocity/mdp/rewards.py | 60 ++- .../locomotion/velocity/velocity_env_cfg.py | 4 +- .../manipulation/reach/mdp/__init__.pyi | 2 + .../manipulation/reach/mdp/curriculums.py | 98 +++++ .../manipulation/reach/mdp/rewards.py | 28 +- .../manipulation/reach/reach_env_cfg.py | 8 +- 71 files changed, 3710 insertions(+), 843 deletions(-) create mode 100644 source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst create mode 100644 source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.py create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/commands_cfg.py create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py create mode 100644 source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py create mode 100644 source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py create mode 100644 source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py create mode 100644 source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py create mode 100644 source/isaaclab_experimental/test/managers/test_curriculum_manager.py create mode 100644 source/isaaclab_experimental/test/managers/test_manager_base.py create mode 100644 source/isaaclab_experimental/test/managers/test_reward_manager.py create mode 100644 source/isaaclab_experimental/test/managers/test_termination_manager.py create mode 100644 source/isaaclab_newton/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst create mode 100644 source/isaaclab_newton/test/actuators/test_adapter_mask_reset.py create mode 100644 source/isaaclab_newton/test/assets/test_newton_mask_reset_forwarding.py create mode 100644 source/isaaclab_physx/changelog.d/jichuanh-warp-frontend-cleanup.rst create mode 100644 source/isaaclab_physx/test/assets/test_physx_mask_reset_forwarding.py create mode 100644 source/isaaclab_tasks/changelog.d/jichuanh-warp-frontend-cleanup.rst create mode 100644 source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst create mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/curriculums.py diff --git a/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst new file mode 100644 index 000000000000..e4bb7eec7da8 --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -0,0 +1,11 @@ +Added +^^^^^ + +* Added :attr:`~isaaclab.terrains.TerrainImporter.terrain_levels_wp` and + :meth:`~isaaclab.terrains.TerrainImporter.update_env_origins_mask` for + mask-based Warp terrain curriculum updates. + +Fixed +^^^^^ + +* Fixed identity quaternion initialization for uniform pose commands. diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py b/source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py new file mode 100644 index 000000000000..72fb7503ef72 --- /dev/null +++ b/source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py @@ -0,0 +1,73 @@ +# 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 + +"""Shared debug-visualization helpers for command terms.""" + +from __future__ import annotations + +import torch + +import isaaclab.utils.math as math_utils +from isaaclab.markers import VisualizationMarkers + + +class _VelocityCommandDebugVis: + """Debug visualization shared by stable and Warp velocity commands.""" + + def _set_debug_vis_impl(self, debug_vis: bool): + if debug_vis: + if not hasattr(self, "goal_vel_visualizer"): + self.goal_vel_visualizer = VisualizationMarkers(self.cfg.goal_vel_visualizer_cfg) + self.current_vel_visualizer = VisualizationMarkers(self.cfg.current_vel_visualizer_cfg) + self.goal_vel_visualizer.set_visibility(True) + self.current_vel_visualizer.set_visibility(True) + elif hasattr(self, "goal_vel_visualizer"): + self.goal_vel_visualizer.set_visibility(False) + self.current_vel_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event): + if not self.robot.is_initialized: + return + base_pos_w = self.robot.data.root_pos_w.torch.clone() + base_pos_w[:, 2] += 0.5 + vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2]) + vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow( + self.robot.data.root_lin_vel_b.torch[:, :2] + ) + self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) + self.current_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale) + + def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Convert an XY velocity [m/s] to arrow scale and orientation.""" + default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale + arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1) + arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0 + heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0]) + zeros = torch.zeros_like(heading_angle) + arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) + arrow_quat = math_utils.quat_mul(self.robot.data.root_quat_w.torch, arrow_quat) + return arrow_scale, arrow_quat + + +class _PoseCommandDebugVis: + """Debug visualization shared by stable and Warp pose commands.""" + + def _set_debug_vis_impl(self, debug_vis: bool): + if debug_vis: + if not hasattr(self, "goal_pose_visualizer"): + self.goal_pose_visualizer = VisualizationMarkers(self.cfg.goal_pose_visualizer_cfg) + self.current_pose_visualizer = VisualizationMarkers(self.cfg.current_pose_visualizer_cfg) + self.goal_pose_visualizer.set_visibility(True) + self.current_pose_visualizer.set_visibility(True) + elif hasattr(self, "goal_pose_visualizer"): + self.goal_pose_visualizer.set_visibility(False) + self.current_pose_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event): + if not self.robot.is_initialized: + return + self.goal_pose_visualizer.visualize(self.pose_command_w[:, :3], self.pose_command_w[:, 3:]) + body_link_pose_w = self.robot.data.body_link_pose_w.torch[:, self.body_idx] + self.current_pose_visualizer.visualize(body_link_pose_w[:, :3], body_link_pose_w[:, 3:7]) diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py index d98a6b5e7f81..876b5fb9ee30 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py @@ -14,17 +14,18 @@ from isaaclab.assets import Articulation from isaaclab.managers import CommandTerm -from isaaclab.markers import VisualizationMarkers from isaaclab.utils.leapp import POSE7_ELEMENT_NAMES from isaaclab.utils.math import combine_frame_transforms, compute_pose_error, quat_from_euler_xyz, quat_unique +from ._debug_vis import _PoseCommandDebugVis + if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv from .commands_cfg import UniformPoseCommandCfg -class UniformPoseCommand(CommandTerm): +class UniformPoseCommand(_PoseCommandDebugVis, CommandTerm): """Command generator for generating pose commands uniformly. The command generator generates poses by sampling positions uniformly within specified @@ -63,7 +64,7 @@ def __init__(self, cfg: UniformPoseCommandCfg, env: ManagerBasedEnv): # create buffers # -- commands: (x, y, z, qx, qy, qz, qw) in root frame self.pose_command_b = torch.zeros(self.num_envs, 7, device=self.device) - self.pose_command_b[:, 3] = 1.0 + self.pose_command_b[:, 6] = 1.0 self.pose_command_w = torch.zeros_like(self.pose_command_b) # -- metrics self.metrics["position_error"] = torch.zeros(self.num_envs, device=self.device) @@ -151,31 +152,3 @@ def _resample_command(self, env_ids: Sequence[int]): def _update_command(self): pass - - def _set_debug_vis_impl(self, debug_vis: bool): - # create markers if necessary for the first time - if debug_vis: - if not hasattr(self, "goal_pose_visualizer"): - # -- goal pose - self.goal_pose_visualizer = VisualizationMarkers(self.cfg.goal_pose_visualizer_cfg) - # -- current body pose - self.current_pose_visualizer = VisualizationMarkers(self.cfg.current_pose_visualizer_cfg) - # set their visibility to true - self.goal_pose_visualizer.set_visibility(True) - self.current_pose_visualizer.set_visibility(True) - else: - if hasattr(self, "goal_pose_visualizer"): - self.goal_pose_visualizer.set_visibility(False) - self.current_pose_visualizer.set_visibility(False) - - def _debug_vis_callback(self, event): - # check if robot is initialized - # note: this is needed in-case the robot is de-initialized. we can't access the data - if not self.robot.is_initialized: - return - # update the markers - # -- goal pose - self.goal_pose_visualizer.visualize(self.pose_command_w[:, :3], self.pose_command_w[:, 3:]) - # -- current body pose - body_link_pose_w = self.robot.data.body_link_pose_w.torch[:, self.body_idx] - self.current_pose_visualizer.visualize(body_link_pose_w[:, :3], body_link_pose_w[:, 3:7]) diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py index 930f663f65d3..31fad8296bc4 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py @@ -16,7 +16,8 @@ import isaaclab.utils.math as math_utils from isaaclab.assets import Articulation from isaaclab.managers import CommandTerm -from isaaclab.markers import VisualizationMarkers + +from ._debug_vis import _VelocityCommandDebugVis if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -27,7 +28,7 @@ logger = logging.getLogger(__name__) -class UniformVelocityCommand(CommandTerm): +class UniformVelocityCommand(_VelocityCommandDebugVis, CommandTerm): r"""Command generator that generates a velocity command in SE(2) from uniform distribution. The command comprises of a linear velocity in x and y direction and an angular velocity around @@ -197,63 +198,6 @@ def _update_command(self): standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten() self.vel_command_b[standing_env_ids, :] = 0.0 - def _set_debug_vis_impl(self, debug_vis: bool): - # set visibility of markers - # note: parent only deals with callbacks. not their visibility - if debug_vis: - # create markers if necessary for the first time - if not hasattr(self, "goal_vel_visualizer"): - # -- goal - self.goal_vel_visualizer = VisualizationMarkers(self.cfg.goal_vel_visualizer_cfg) - # -- current - self.current_vel_visualizer = VisualizationMarkers(self.cfg.current_vel_visualizer_cfg) - # set their visibility to true - self.goal_vel_visualizer.set_visibility(True) - self.current_vel_visualizer.set_visibility(True) - else: - if hasattr(self, "goal_vel_visualizer"): - self.goal_vel_visualizer.set_visibility(False) - self.current_vel_visualizer.set_visibility(False) - - def _debug_vis_callback(self, event): - # check if robot is initialized - # note: this is needed in-case the robot is de-initialized. we can't access the data - if not self.robot.is_initialized: - return - # get marker location - # -- base state - base_pos_w = self.robot.data.root_pos_w.torch.clone() - base_pos_w[:, 2] += 0.5 - # -- resolve the scales and quaternions - vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2]) - vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow( - self.robot.data.root_lin_vel_b.torch[:, :2] - ) - # display markers - self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) - self.current_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale) - - """ - Internal helpers. - """ - - def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Converts the XY base velocity command to arrow direction rotation.""" - # obtain default scale of the marker - default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale - # arrow-scale - arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1) - arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0 - # arrow-direction - heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0]) - zeros = torch.zeros_like(heading_angle) - arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) - # convert everything back from base to world frame - base_quat_w = self.robot.data.root_quat_w.torch - arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) - - return arrow_scale, arrow_quat - class NormalVelocityCommand(UniformVelocityCommand): """Command generator that generates a velocity command in SE(2) from a normal distribution. diff --git a/source/isaaclab/isaaclab/terrains/terrain_importer.py b/source/isaaclab/isaaclab/terrains/terrain_importer.py index c78f9ae3344e..2320996bd774 100644 --- a/source/isaaclab/isaaclab/terrains/terrain_importer.py +++ b/source/isaaclab/isaaclab/terrains/terrain_importer.py @@ -11,6 +11,7 @@ import numpy as np import torch import trimesh +import warp as wp import isaaclab.sim as sim_utils from isaaclab.markers import VisualizationMarkers @@ -25,6 +26,32 @@ logger = logging.getLogger(__name__) +@wp.kernel +def _update_env_origins_mask( + env_mask: wp.array(dtype=wp.bool), + move_up: wp.array(dtype=wp.bool), + move_down: wp.array(dtype=wp.bool), + rng_state: wp.array(dtype=wp.uint32), + terrain_levels: wp.array(dtype=wp.int64), + terrain_types: wp.array(dtype=wp.int64), + terrain_origins: wp.array(dtype=wp.vec3f, ndim=2), + env_origins: wp.array(dtype=wp.vec3f), + max_terrain_level: int, +): + env_id = wp.tid() + if env_mask[env_id]: + state = rng_state[env_id] + random_level = wp.randi(state, 0, max_terrain_level) + level = terrain_levels[env_id] + wp.int64(move_up[env_id]) - wp.int64(move_down[env_id]) + if level >= wp.int64(max_terrain_level): + level = wp.int64(random_level) + elif level < wp.int64(0): + level = wp.int64(0) + terrain_levels[env_id] = level + env_origins[env_id] = terrain_origins[level, terrain_types[env_id]] + rng_state[env_id] = state + + class TerrainImporter: r"""A class to handle terrain meshes and import them into the simulator. @@ -77,6 +104,10 @@ def __init__(self, cfg: TerrainImporterCfg): self.terrain_prim_paths = list() self.terrain_origins = None self.env_origins = None # assigned later when `configure_env_origins` is called + self._terrain_levels_wp = None + self._terrain_types_wp = None + self._terrain_origins_wp = None + self._env_origins_wp = None # private variables self._terrain_flat_patches = dict() @@ -144,6 +175,12 @@ def terrain_names(self) -> list[str]: """A list of names of the imported terrains.""" return [f"'{path.split('/')[-1]}'" for path in self.terrain_prim_paths] + @property + def terrain_levels_wp(self) -> wp.array(dtype=wp.int64): + """Pointer-stable Warp view of the per-environment terrain levels.""" + self._initialize_warp_origin_views() + return self._terrain_levels_wp + """ Operations - Visibility. """ @@ -310,6 +347,7 @@ def configure_env_origins(self, origins: np.ndarray | torch.Tensor | None = None raise ValueError("Environment spacing must be specified for configuring grid-like origins.") # compute environment origins self.env_origins = self._compute_env_origins_grid(self.cfg.num_envs, self.cfg.env_spacing) + self._invalidate_warp_origin_views() def update_env_origins(self, env_ids: torch.Tensor, move_up: torch.Tensor, move_down: torch.Tensor): """Update the environment origins based on the terrain levels.""" @@ -328,10 +366,90 @@ def update_env_origins(self, env_ids: torch.Tensor, move_up: torch.Tensor, move_ # update the env origins self.env_origins[env_ids] = self.terrain_origins[self.terrain_levels[env_ids], self.terrain_types[env_ids]] + def update_env_origins_mask( + self, + env_mask: wp.array(dtype=wp.bool), + move_up: wp.array(dtype=wp.bool), + move_down: wp.array(dtype=wp.bool), + rng_state: wp.array(dtype=wp.uint32), + ) -> None: + """Update terrain levels and origins for a boolean environment mask. + + This method preserves the level update semantics of :meth:`update_env_origins`: moving past the + maximum level wraps to a random level, moving below zero clamps to zero, and unselected environments + remain unchanged. + + Args: + env_mask: Boolean Warp mask selecting environments to update. Shape is + ``(num_envs,)`` on :attr:`device`. + move_up: Boolean Warp flags that increment selected terrain levels. + Shape is ``(num_envs,)`` on :attr:`device`. + move_down: Boolean Warp flags that decrement selected terrain levels. + Shape is ``(num_envs,)`` on :attr:`device`. + rng_state: Per-environment Warp random-number-generator state. Shape is + ``(num_envs,)`` on :attr:`device`. + + Raises: + ValueError: If an input array does not match the configured environment count. + TypeError: If an input array has the wrong Warp data type. + """ + if self.terrain_origins is None: + return + num_envs = self.terrain_levels.shape[0] + arrays = { + "env_mask": (env_mask, wp.bool), + "move_up": (move_up, wp.bool), + "move_down": (move_down, wp.bool), + "rng_state": (rng_state, wp.uint32), + } + for name, (array, dtype) in arrays.items(): + if not isinstance(array, wp.array) or array.dtype != dtype: + raise TypeError(f"{name} must be a Warp array with dtype {dtype}; received {array}.") + if array.ndim != 1 or array.shape[0] != num_envs: + raise ValueError(f"{name} must have shape ({num_envs},); received {array.shape}.") + if array.device != wp.get_device(self.device): + raise ValueError(f"{name} must be on device {self.device}; received {array.device}.") + + self._initialize_warp_origin_views() + wp.launch( + kernel=_update_env_origins_mask, + dim=num_envs, + inputs=[ + env_mask, + move_up, + move_down, + rng_state, + self._terrain_levels_wp, + self._terrain_types_wp, + self._terrain_origins_wp, + self._env_origins_wp, + self.max_terrain_level, + ], + device=self.device, + ) + """ Internal helpers. """ + def _invalidate_warp_origin_views(self) -> None: + """Invalidate cached Warp views after origin storage changes.""" + self._terrain_levels_wp = None + self._terrain_types_wp = None + self._terrain_origins_wp = None + self._env_origins_wp = None + + def _initialize_warp_origin_views(self) -> None: + """Create persistent Warp views of terrain-origin Torch storage.""" + if self._terrain_levels_wp is not None: + return + if self.terrain_origins is None or self.env_origins is None: + raise RuntimeError("Terrain origins are not configured for curriculum updates.") + self._terrain_levels_wp = wp.from_torch(self.terrain_levels, dtype=wp.int64) + self._terrain_types_wp = wp.from_torch(self.terrain_types, dtype=wp.int64) + self._terrain_origins_wp = wp.from_torch(self.terrain_origins, dtype=wp.vec3f) + self._env_origins_wp = wp.from_torch(self.env_origins, dtype=wp.vec3f) + def _compute_env_origins_curriculum(self, num_envs: int, origins: torch.Tensor) -> torch.Tensor: """Compute the origins of the environments defined by the sub-terrains origins.""" # extract number of rows and cols diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst new file mode 100644 index 000000000000..234876b28416 --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -0,0 +1,28 @@ +Added +^^^^^ + +* Added Warp-native uniform pose and velocity command terms with pointer-stable + command arrays. +* Added :class:`~isaaclab_experimental.managers.CurriculumManager` and a + boolean-mask reset path for Warp manager-based environments, retaining compact + environment IDs only for legacy host consumers. + +Changed +^^^^^^^ + +* Changed Warp manager and direct-environment stages to execute eagerly by + default while stateful capture semantics are validated. Set + ``MANAGER_CALL_CONFIG='{"default": 2}'`` for manager capture or + ``ISAACLAB_WARP_DIRECT_CAPTURE=1`` for the legacy direct capture path. + +Deprecated +^^^^^^^^^^ + +* Deprecated selecting stable managers inside ``ManagerBasedEnvWarp``. Use + ``ManagerBasedEnv`` for Torch managers or mode ``1`` for the Warp frontend. + +Fixed +^^^^^ + +* Fixed Warp event state leaking across environments and configurations, and + corrected center-of-mass sampling and termination metrics for partial resets. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py index d69b13a8feec..dcd6fdf29a38 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py @@ -38,6 +38,7 @@ from isaaclab.utils.timer import Timer from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp +from isaaclab_experimental.utils.warp import increment_all_int32, zero_masked_int32 from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache # from isaacsim.core.simulation_manager import SimulationManager @@ -53,24 +54,8 @@ DEBUG_TIMERS = os.environ.get("DEBUG_TIMERS", "0") == "1" """Enable all fine-grained inner timers (adds wp.synchronize per sub-phase). Set DEBUG_TIMERS=1 env var to enable.""" - -@wp.kernel -def zero_mask_int32( - mask: wp.array(dtype=wp.bool), - data: wp.array(dtype=wp.int32), -): - env_index = wp.tid() - if mask[env_index]: - data[env_index] = 0 - - -@wp.kernel -def add_to_env( - data: wp.array(dtype=wp.int32), - value: wp.int32, -): - env_index = wp.tid() - data[env_index] += value +WARP_DIRECT_CAPTURE = os.environ.get("ISAACLAB_WARP_DIRECT_CAPTURE", "0") == "1" +"""Enable direct-environment stage capture. Eager execution is the correctness-first default.""" class DirectRLEnvWarp(DirectRLEnv): @@ -117,6 +102,10 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs RuntimeError: If a simulation context already exists. The environment must always create one since it configures the simulation context and controls the simulation. """ + # Keep the Warp frontend on the Warp-native actuator path. This covers + # implicit, PD, delayed, DC, and neural actuator configurations without + # crossing through the Torch-based Isaac Lab actuator loop each step. + cfg.sim.use_newton_actuators = True # check that the config is valid cfg.validate() # store inputs to class @@ -243,8 +232,10 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.torch_reset_time_outs: torch.Tensor = None self.torch_episode_length_buf: torch.Tensor = None - # Warp CUDA graph cache for capture-or-replay - self._graph_cache = WarpGraphCache() + # Direct stages are Warp eager by default. The owner-held cache keeps + # execution dispatch out of call sites and can enable capture later once + # stateful warm-up semantics are defined. + self._graph_cache = WarpGraphCache(enabled=WARP_DIRECT_CAPTURE) # setup the action and observation spaces for Gym self._configure_gym_env_spaces() @@ -429,8 +420,8 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: with Timer(name="apply_action", msg="Action processing step took:", enable=DEBUG_TIMERS): self._graph_cache.capture_or_replay("action", self.step_warp_action) - # write_data_to_sim runs outside the CUDA graph because _apply_actuator_model - # uses torch ops (wp.to_torch + torch arithmetic) that cross CUDA streams. + # Keep scene writes outside the task graph until scene, sensor, and + # actuator capturability have been validated as one backend boundary. with Timer(name="write_data_to_sim_loop", msg="Write data to sim (loop) took:", enable=DEBUG_TIMERS): self.scene.write_data_to_sim() @@ -448,7 +439,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self.common_step_counter += 1 # total step (common for all envs) with Timer(name="end_pre_graph", msg="End pre-graph took:", enable=DEBUG_TIMERS): self._graph_cache.capture_or_replay("end_pre", self._step_warp_end_pre) - # write_data_to_sim runs uncaptured — it uses torch ops that cross CUDA streams. + # Keep the post-reset scene write at the explicit backend boundary. with Timer(name="write_data_to_sim_post", msg="Write data to sim (post-reset) took:", enable=DEBUG_TIMERS): self.scene.write_data_to_sim() with Timer(name="end_post_graph", msg="End post-graph took:", enable=DEBUG_TIMERS): @@ -479,15 +470,13 @@ def _post_step_visualize(self) -> None: def step_warp_action(self) -> None: self._apply_action() - # Note: scene.write_data_to_sim() is called separately outside the CUDA graph - # capture scope because it invokes _apply_actuator_model() which uses torch - # arithmetic (wp.to_torch + torch ops). This would cause a CUDA stream crossing - # error during graph capture. Moving it outside is safe since it runs every step. + # Scene writes remain a separate backend stage. This keeps the Warp-first + # task path correct without making capture support a prerequisite. def _step_warp_end_pre(self) -> None: """Capturable portion before write_data_to_sim (pure warp kernels).""" wp.launch( - add_to_env, + increment_all_int32, dim=self.num_envs, inputs=[ self._episode_length_buf_wp, @@ -741,7 +730,7 @@ def _reset_idx(self, mask: wp.array | None = None): # reset the episode length buffer wp.launch( - zero_mask_int32, + zero_masked_int32, dim=self.num_envs, inputs=[ mask, diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py index f792b994f67b..6722599f9032 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py @@ -21,24 +21,53 @@ class InteractiveSceneWarp(InteractiveScene): avoiding the need to convert between env_ids and masks. """ - def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None): - """Reset scene entities using either env_ids or a boolean env_mask. + def reset( + self, + env_ids: Sequence[int] | None = None, + env_mask: wp.array(dtype=wp.bool) | None = None, + ) -> None: + """Reset scene entities using environment IDs or a boolean mask. + + Mask-based calls stay on the Warp data path for every entity that supports + masks. Surface grippers remain an explicit ID-based host boundary and are + reset only when :paramref:`env_ids` is also provided. Calls that only pass + IDs preserve the behavior of :meth:`InteractiveScene.reset`. Args: env_ids: The indices of the environments to reset. Defaults to None (all instances). env_mask: Boolean warp mask of shape (num_envs,). Defaults to None. """ - # -- assets (support env_mask) + if env_mask is None: + super().reset(env_ids) + return + if env_mask.dtype != wp.bool or env_mask.ndim != 1: + raise TypeError(f"env_mask must be a one-dimensional Warp boolean array, got {env_mask}.") + if env_mask.shape[0] != self.cfg.num_envs: + raise ValueError(f"env_mask must have shape ({self.cfg.num_envs},), received {env_mask.shape}.") + if env_mask.device != wp.get_device(self.sim.device): + raise ValueError(f"env_mask must be on device {self.sim.device}, received {env_mask.device}.") + + # -- Warp-capable assets for articulation in self._articulations.values(): - articulation.reset(env_ids, env_mask=env_mask) + articulation.reset(env_mask=env_mask) for deformable_object in self._deformable_objects.values(): - deformable_object.reset(env_ids) + deformable_object.reset(env_mask=env_mask) for rigid_object in self._rigid_objects.values(): - rigid_object.reset(env_ids, env_mask=env_mask) - for surface_gripper in self._surface_grippers.values(): - surface_gripper.reset(env_ids) + rigid_object.reset(env_mask=env_mask) for rigid_object_collection in self._rigid_object_collections.values(): - rigid_object_collection.reset(env_ids, env_mask=env_mask) - # -- sensors (no env_mask support) + rigid_object_collection.reset(env_mask=env_mask) + # -- Warp-capable sensors for sensor in self._sensors.values(): - sensor.reset(env_ids) + sensor.reset(env_mask=env_mask) + + if env_ids is not None: + self.reset_host(env_ids) + + def reset_host(self, env_ids: Sequence[int] | None = None) -> None: + """Reset ID-based scene entities at an explicit host boundary. + + Args: + env_ids: The indices of the environments to reset. Defaults to all instances. + """ + for surface_gripper in self._surface_grippers.values(): + surface_gripper.reset(env_ids) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py index 2b4791d22a16..f6979f7052de 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py @@ -73,6 +73,10 @@ def __init__(self, cfg: ManagerBasedEnvCfg): RuntimeError: If a simulation context already exists. The environment must always create one since it configures the simulation context and controls the simulation. """ + # Keep the Warp frontend on the Warp-native actuator path. This covers + # implicit, PD, delayed, DC, and neural actuator configurations without + # crossing through the Torch-based Isaac Lab actuator loop each step. + cfg.sim.use_newton_actuators = True # check that the config is valid cfg.validate() # store inputs to class @@ -367,6 +371,7 @@ def load_managers(self): self.recorder_manager = self._manager_call_switch.resolve_manager_class("RecorderManager")( self.cfg.recorders, self ) + self._has_recorders = bool(self.recorder_manager.active_terms) print("[INFO] Recorder Manager: ", self.recorder_manager) # -- action manager self.action_manager = self._manager_call_switch.resolve_manager_class("ActionManager")(self.cfg.actions, self) @@ -396,7 +401,11 @@ def setup_manager_visualizers(self): """ def reset( - self, seed: int | None = None, env_ids: Sequence[int] | None = None, options: dict[str, Any] | None = None + self, + seed: int | None = None, + env_ids: Sequence[int] | None = None, + options: dict[str, Any] | None = None, + env_mask: wp.array | torch.Tensor | None = None, ) -> tuple[VecEnvObs, dict]: """Resets the specified environments and returns observations. @@ -412,14 +421,19 @@ def reset( Note: This argument is used for compatibility with Gymnasium environment definition. + env_mask: Boolean environment mask. This is the canonical Warp + frontend selection; :paramref:`env_ids` remains supported for + API compatibility. + Returns: A tuple containing the observations and extras. """ - if env_ids is None: - env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device) - - # trigger recorder terms for pre-reset calls - self.recorder_manager.record_pre_reset(env_ids) + reset_mask = self.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + host_env_ids = self._resolve_reset_host_ids(env_ids=env_ids, env_mask=reset_mask) + if host_env_ids is not None: + if self._has_recorders: + self.recorder_manager.record_pre_reset(host_env_ids) + self._reset_host_pre(host_env_ids) # set the seed if seed is not None: @@ -434,8 +448,9 @@ def reset( device=self.device, ) - # reset state of scene - self._reset_idx(env_ids) + self._reset_idx(env_mask=reset_mask, env_ids=host_env_ids) + if host_env_ids is not None: + self.extras["log"].update(self._reset_host_post(host_env_ids)) # update articulation kinematics self.scene.write_data_to_sim() @@ -445,8 +460,8 @@ def reset( for _ in range(self.cfg.num_rerenders_on_reset): self.sim.render() - # trigger recorder terms for post-reset calls - self.recorder_manager.record_post_reset(env_ids) + if self._has_recorders and host_env_ids is not None: + self.recorder_manager.record_post_reset(host_env_ids) # compute observations self.obs_buf = self.observation_manager.compute(update_history=True) @@ -482,14 +497,17 @@ def reset_to( if env_ids is None: env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device) - # trigger recorder terms for pre-reset calls - self.recorder_manager.record_pre_reset(env_ids) + if self._has_recorders: + self.recorder_manager.record_pre_reset(env_ids) # set the seed if seed is not None: self.seed(seed) - self._reset_idx(env_ids) + env_mask = self.resolve_env_mask(env_ids=env_ids) + self._reset_host_pre(env_ids) + self._reset_idx(env_mask=env_mask, env_ids=env_ids) + self.extras["log"].update(self._reset_host_post(env_ids)) # set the state self.scene.reset_to(state, env_ids, is_relative=is_relative) @@ -502,8 +520,8 @@ def reset_to( for _ in range(self.cfg.num_rerenders_on_reset): self.sim.render() - # trigger recorder terms for post-reset calls - self.recorder_manager.record_post_reset(env_ids) + if self._has_recorders: + self.recorder_manager.record_post_reset(env_ids) # compute observations self.obs_buf = self.observation_manager.compute(update_history=True) @@ -535,7 +553,8 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: action_wp = wp.from_torch(action_device, dtype=wp.float32) self.action_manager.process_action(action_wp) - self.recorder_manager.record_pre_step() + if self._has_recorders: + self.recorder_manager.record_pre_step() # check if we need to do rendering within the physics loop # note: hoisted out of the decimation loop; is_rendering does live settings lookups @@ -564,7 +583,8 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: # -- compute observations self.obs_buf = self.observation_manager.compute(update_history=True) - self.recorder_manager.record_post_step() + if self._has_recorders: + self.recorder_manager.record_post_step() # return observations and extras return self.obs_buf, self.extras @@ -614,11 +634,7 @@ def close(self): """ def _resolve_stable_cfg_counterpart(self) -> ManagerBasedEnvCfg | None: - """Resolve a stable task config counterpart for the current experimental task config. - - The lookup follows a module-name mirror convention: - ``isaaclab_tasks_experimental...`` -> ``isaaclab_tasks...`` with the same config class name. - """ + """Resolve the legacy stable task config counterpart.""" cfg_cls = self.cfg.__class__ cfg_module_name = cfg_cls.__module__ if "isaaclab_tasks_experimental" not in cfg_module_name: @@ -637,11 +653,7 @@ def _resolve_stable_cfg_counterpart(self) -> ManagerBasedEnvCfg | None: stable_cfg_cls = getattr(stable_module, cfg_cls.__name__, None) if stable_cfg_cls is None: - logger.warning( - "Stable task cfg class '%s' not found in module '%s'.", - cfg_cls.__name__, - stable_module_name, - ) + logger.warning("Stable task cfg class '%s' not found in module '%s'.", cfg_cls.__name__, stable_module_name) return None try: @@ -656,11 +668,7 @@ def _resolve_stable_cfg_counterpart(self) -> ManagerBasedEnvCfg | None: return None def _apply_manager_term_cfg_profile(self) -> None: - """Align term configs with manager modes for stable manager selections. - - When a manager is configured as STABLE (0), swap its corresponding config subtree - from the stable task counterpart to keep manager-term type/signature compatibility. - """ + """Preserve the deprecated mixed stable-manager configuration path.""" manager_to_cfg_attr = { "ActionManager": "actions", "ObservationManager": "observations", @@ -671,60 +679,63 @@ def _apply_manager_term_cfg_profile(self) -> None: "RewardManager": "rewards", "CurriculumManager": "curriculum", } - - stable_manager_names = [ - manager_name - for manager_name in manager_to_cfg_attr - if self._manager_call_switch.get_mode_for_manager(manager_name) == ManagerCallMode.STABLE + stable_managers = [ + name + for name in manager_to_cfg_attr + if self._manager_call_switch.get_mode_for_manager(name) == ManagerCallMode.STABLE ] - if not stable_manager_names: + if not stable_managers: return + warnings.warn( + "Selecting STABLE managers inside ManagerBasedEnvWarp is deprecated. Use ManagerBasedEnv for Torch " + "managers or WARP_NOT_CAPTURED for the Warp frontend.", + DeprecationWarning, + stacklevel=2, + ) stable_cfg = self._resolve_stable_cfg_counterpart() if stable_cfg is None: logger.warning( - "Stable managers requested (%s), but no stable cfg counterpart could be resolved." - " Keeping experimental term configs.", - ", ".join(stable_manager_names), + "Stable managers requested (%s), but no stable cfg counterpart could be resolved. Keeping " + "experimental term configs.", + ", ".join(stable_managers), ) return - replaced_items: list[str] = [] for manager_name, cfg_attr in manager_to_cfg_attr.items(): if self._manager_call_switch.get_mode_for_manager(manager_name) != ManagerCallMode.STABLE: continue - if not hasattr(self.cfg, cfg_attr) or not hasattr(stable_cfg, cfg_attr): - continue - setattr(self.cfg, cfg_attr, deepcopy(getattr(stable_cfg, cfg_attr))) - replaced_items.append(f"{manager_name} -> cfg.{cfg_attr}") + if hasattr(self.cfg, cfg_attr) and hasattr(stable_cfg, cfg_attr): + setattr(self.cfg, cfg_attr, deepcopy(getattr(stable_cfg, cfg_attr))) - if replaced_items: - print("[INFO] Applied stable term config profile for managers:") - for item in replaced_items: - print(f" - {item}") - - def _reset_idx(self, env_ids: Sequence[int]): - """Reset environments based on specified indices. + def _reset_idx( + self, + *, + env_mask: wp.array(dtype=wp.bool), + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Reset Warp-owned state for selected environments. Args: - env_ids: List of environment ids which must be reset + env_mask: Boolean Warp mask selecting environments to reset. + env_ids: Optional compact IDs for explicit host-only terms. """ + del env_ids # reset the internal buffers of the scene elements - self.scene.reset(env_ids) + self.scene.reset(env_mask=env_mask) # apply events such as randomization for environments that need a reset if "reset" in self.event_manager.available_modes: env_step_count = self._sim_step_counter // self.cfg.decimation self._global_env_step_count_wp.fill_(env_step_count) self.event_manager.apply( - mode="reset", env_ids=env_ids, global_env_step_count=self._global_env_step_count_wp + mode="reset", env_mask_wp=env_mask, global_env_step_count=self._global_env_step_count_wp ) # iterate over all managers and reset them # this returns a dictionary of information which is stored in the extras # note: This is order-sensitive! Certain things need be reset before others. self.extras["log"] = dict() - env_mask = self.resolve_env_mask(env_ids=env_ids) # -- observation manager info = self.observation_manager.reset(env_mask=env_mask) self.extras["log"].update(info) @@ -734,6 +745,39 @@ def _reset_idx(self, env_ids: Sequence[int]): # -- event manager info = self.event_manager.reset(env_mask=env_mask) self.extras["log"].update(info) - # -- recorder manager - info = self.recorder_manager.reset(env_ids) - self.extras["log"].update(info) + + def _reset_host_pre(self, env_ids: Sequence[int] | torch.Tensor) -> None: + """Reset ID-only scene state before the Warp reset stage.""" + self.scene.reset_host(env_ids) + + def _reset_host_post(self, env_ids: Sequence[int] | torch.Tensor) -> dict[str, Any]: + """Reset host-only managers after the Warp reset stage.""" + if self._has_recorders: + return self.recorder_manager.reset(env_ids) + return {} + + def _reset_requires_host_selection(self) -> bool: + """Return whether configured reset features require a host-visible selection.""" + return bool(self._has_recorders or self.scene.surface_grippers) + + def _resolve_reset_host_ids( + self, + *, + env_ids: Sequence[int] | wp.array | torch.Tensor | None, + env_mask: wp.array(dtype=wp.bool), + ) -> Sequence[int] | torch.Tensor | None: + """Materialize reset IDs only when a configured host feature consumes them. + + Args: + env_ids: Original environment ID selection, if provided. + env_mask: Canonical Warp reset mask. + + Returns: + Host-readable environment IDs, or ``None`` when the reset remains + entirely on the Warp data path. + """ + if not self._reset_requires_host_selection(): + return None + if env_ids is not None and not isinstance(env_ids, wp.array): + return env_ids + return wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index 2e1894e54003..1c3632ef88db 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -25,12 +25,12 @@ from isaaclab.envs.common import VecEnvStepReturn from isaaclab.envs.manager_based_rl_env_cfg import ManagerBasedRLEnvCfg -from isaaclab.managers import CommandManager from isaaclab.ui.widgets import ManagerLiveVisualizer from isaaclab.utils.timer import Timer -from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode +from isaaclab_experimental.managers import CommandManager, CurriculumManager from isaaclab_experimental.utils.torch_utils import clone_obs_buffer +from isaaclab_experimental.utils.warp import increment_all_int64, zero_masked_int64 from .manager_based_env_warp import ManagerBasedEnvWarp @@ -106,9 +106,9 @@ def __init__(self, cfg: ManagerBasedRLEnvCfg, render_mode: str | None = None, ** # store the render mode self.render_mode = render_mode - # The persistent reset mask needed for warp capture - # The intended use is to copy into this mask whenever capture is needed - # TODO: termination manager provides the same mask, investigate whether this can be replaced. + # Reset stages may be captured independently, so every stage must see + # one owner-held mask pointer. Copy each new selection into this buffer + # before dispatch instead of recording caller-owned mask addresses. self.reset_mask_wp = wp.zeros(cfg.scene.num_envs, dtype=wp.bool, device=cfg.sim.device) # Persistent action input buffer to keep pointer stable for captured graphs. @@ -155,7 +155,7 @@ def max_episode_length(self) -> int: def load_managers(self): # note: this order is important since observation manager needs to know the command and action managers # and the reward manager needs to know the termination manager - # -- command manager (stable impl — not routed through ManagerCallSwitch) + # -- Warp-first command manager self.command_manager = CommandManager(self.cfg.commands, self) print("[INFO] Command Manager: ", self.command_manager) @@ -171,10 +171,8 @@ def load_managers(self): # -- reward manager (experimental fork; Warp-compatible rewards) self.reward_manager = self._manager_call_switch.resolve_manager_class("RewardManager")(self.cfg.rewards, self) print("[INFO] Reward Manager: ", self.reward_manager) - # -- curriculum manager - self.curriculum_manager = self._manager_call_switch.resolve_manager_class("CurriculumManager")( - self.cfg.curriculum, self - ) + # -- Warp-first curriculum manager + self.curriculum_manager = CurriculumManager(self.cfg.curriculum, self) print("[INFO] Curriculum Manager: ", self.curriculum_manager) # setup the action and observation spaces for Gym @@ -243,13 +241,15 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: action_device = action.to(self.device) wp.copy(self._action_in_wp, wp.from_torch(action_device, dtype=wp.float32)) - self._manager_call_switch.call_stage( - stage="ActionManager_process_action", - warp_call={"fn": self.action_manager.process_action, "kwargs": {"action": self._action_in_wp}}, - timer=DEBUG_TIMER_STEP, + self._manager_call_switch.call( + "ActionManager_process_action", + self.action_manager.process_action, + action=self._action_in_wp, + _timer=DEBUG_TIMER_STEP, ) - self.recorder_manager.record_pre_step() + if self._has_recorders: + self.recorder_manager.record_pre_step() # check if we need to do rendering within the physics loop # note: checked here once to avoid multiple checks within the loop @@ -259,21 +259,22 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: for _ in range(self.cfg.decimation): self._sim_step_counter += 1 # set actions into buffers - self._manager_call_switch.call_stage( - stage="ActionManager_apply_action", - warp_call={"fn": self.action_manager.apply_action}, - timer=DEBUG_TIMER_STEP, + self._manager_call_switch.call( + "ActionManager_apply_action", + self.action_manager.apply_action, + _timer=DEBUG_TIMER_STEP, ) - self._manager_call_switch.call_stage( - stage="Scene_write_data_to_sim", - warp_call={"fn": self.scene.write_data_to_sim}, - timer=DEBUG_TIMER_STEP, + self._manager_call_switch.call( + "Scene_write_data_to_sim", + self.scene.write_data_to_sim, + _timer=DEBUG_TIMER_STEP, ) # simulate with Timer(name="simulate", msg="Newton simulation step took:", enable=DEBUG_TIMER_STEP, time_unit="us"): self.sim.step(render=False) - self.recorder_manager.record_post_physics_decimation_step() + if self._has_recorders: + self.recorder_manager.record_post_physics_decimation_step() # render between steps only if the GUI or an RTX sensor needs it # note: we assume the render interval to be the shortest accepted rendering interval. # If a camera needs rendering at a faster frequency, this will lead to unexpected behavior. @@ -290,90 +291,66 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: # post-step: # -- update env counters (used for curriculum generation) - self.episode_length_buf += 1 # step in current episode (per env) + wp.launch( + increment_all_int64, + dim=self.num_envs, + inputs=[self._episode_length_buf_wp, 1], + device=self.device, + ) self.common_step_counter += 1 # total step (common for all envs) # -- post-processing (termination + reward) as independently configurable stages - self._manager_call_switch.call_stage( - stage="TerminationManager_compute", - warp_call={"fn": self.step_warp_termination_compute}, - timer=DEBUG_TIMER_STEP, + self._manager_call_switch.call( + "TerminationManager_compute", + self.step_warp_termination_compute, + _timer=DEBUG_TIMER_STEP, ) - self.reward_buf = self._manager_call_switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": self.reward_manager.compute, "kwargs": {"dt": float(self.step_dt)}}, - timer=DEBUG_TIMER_STEP, + self.reward_buf = self._manager_call_switch.call( + "RewardManager_compute", + self.reward_manager.compute, + dt=float(self.step_dt), + _timer=DEBUG_TIMER_STEP, ) - if len(self.recorder_manager.active_terms) > 0: + if self._has_recorders: # update observations for recording if needed - self._manager_call_switch.call_stage( - stage="ObservationManager_compute_no_history", - warp_call={"fn": self.observation_manager.compute, "kwargs": {"return_cloned_output": False}}, - timer=DEBUG_TIMER_STEP, + self._manager_call_switch.call( + "ObservationManager_compute_no_history", + self.observation_manager.compute, + return_cloned_output=False, + _timer=DEBUG_TIMER_STEP, ) self.recorder_manager.record_post_step() - # -- reset envs that terminated/timed-out and log the episode information - # NOTE: Interim path (intentional). - # We still compact `reset_buf` into `env_ids` here because several reset-time managers/recorders - # are still `env_ids`-based. Do NOT remove/replace this until mask-based reset is end-to-end. - with Timer( - name="reset_selection", - msg="Reset selection took:", - enable=DEBUG_TIMER_STEP, - time_unit="us", - ): - # Keep the reset-mask handoff fully in Warp when experimental termination buffers exist. - # Stable termination manager path exposes torch-only dones/reset buffers. - termination_manager_mode = self._manager_call_switch.get_mode_for_manager("TerminationManager") - if termination_manager_mode == ManagerCallMode.STABLE: - # copy still needed as mask will be used if manager is set to mode > 0 - wp.copy(self.reset_mask_wp, wp.from_torch(self.reset_buf, dtype=wp.bool)) - else: - wp.copy(self.reset_mask_wp, self.termination_manager.dones_wp) - reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) - if len(reset_env_ids) > 0: - # trigger recorder terms for pre-reset calls - self.recorder_manager.record_pre_reset(reset_env_ids) - - with Timer( - name="reset_idx", - msg="Reset idx took:", - enable=DEBUG_TIMER_STEP, - time_unit="us", - ): - self._reset_idx(env_ids=reset_env_ids, env_mask=self.reset_mask_wp) - - # if sensors are added to the scene, make sure we render to reflect changes in reset - if self.has_rtx_sensors and self.cfg.num_rerenders_on_reset > 0: - for _ in range(self.cfg.num_rerenders_on_reset): - self.sim.render() - - # trigger recorder terms for post-reset calls - self.recorder_manager.record_post_reset(reset_env_ids) + self._reset_terminated_envs() # -- update command - self.command_manager.compute(dt=float(self.step_dt)) + self._manager_call_switch.call( + "CommandManager_compute", + self.command_manager.compute, + dt=float(self.step_dt), + _timer=DEBUG_TIMER_STEP, + ) # -- step interval events if "interval" in self.event_manager.available_modes: - self._manager_call_switch.call_stage( - stage="EventManager_apply_interval", - warp_call={"fn": self.event_manager.apply, "kwargs": {"mode": "interval", "dt": float(self.step_dt)}}, - timer=DEBUG_TIMER_STEP, + self._manager_call_switch.call( + "EventManager_apply_interval", + self.event_manager.apply, + mode="interval", + dt=float(self.step_dt), + _timer=DEBUG_TIMER_STEP, ) # -- compute observations # note: done after reset to get the correct observations for reset envs - self.obs_buf = self._manager_call_switch.call_stage( - stage="ObservationManager_compute_update_history", - warp_call={ - "fn": self.observation_manager.compute, - "kwargs": {"update_history": True, "return_cloned_output": False}, - "output": lambda r: clone_obs_buffer(r), - }, - timer=DEBUG_TIMER_STEP, + self.obs_buf = self._manager_call_switch.call( + "ObservationManager_compute_update_history", + self.observation_manager.compute, + update_history=True, + return_cloned_output=False, + _output=clone_obs_buffer, + _timer=DEBUG_TIMER_STEP, ) # return observations, rewards, resets and extras return self.obs_buf, self.reward_buf, self.reset_terminated, self.reset_time_outs, self.extras @@ -487,118 +464,104 @@ def _configure_gym_env_spaces(self): def _reset_idx( self, - env_ids: Sequence[int] | slice | torch.Tensor, *, - env_mask: wp.array | None = None, - ): - """Reset environments based on specified indices. - - IMPORTANT: - This function always uses the **TerminationManager-produced Warp env mask** (`self.reset_buf`) to select - which envs to reset. The ids/mask conversion is performed in `step()` before calling this function. - - In other words: - - If `env_mask` is provided, it **must** be `self.reset_buf` (Warp bool mask) - - If `env_mask` is not provided, this function will populate `self.reset_buf` from `env_ids` - - When `env_mask` is provided, `env_ids` **must** correspond to the same mask + env_mask: wp.array(dtype=wp.bool), + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Reset Warp-owned RL state for selected environments. Args: - env_ids: Environment indices to reset. - env_mask: Warp boolean env mask selecting envs to reset. Must be `self.reset_buf`. - If None, uses and populates `self.reset_buf` from `env_ids`. + env_mask: Boolean Warp mask selecting environments to reset. + env_ids: Optional compact IDs for legacy curriculum terms that + explicitly require a host selection. """ - if env_mask is None: - # Base `reset()` / `reset_to()` call-path provides only `env_ids`. - # Populate the stable TerminationManager-owned mask (`self.reset_buf`) from ids. - env_mask = self.reset_mask_wp - # Use the centralized env-id/mask resolution from the base Warp env, then copy into the - # stable TerminationManager-owned buffer (`self.reset_buf`) used by captured graphs. - resolved_mask = self.resolve_env_mask(env_ids=env_ids) - wp.copy(env_mask, resolved_mask) - - if not isinstance(env_mask, wp.array): - raise TypeError(f"env_mask must be a wp.array (got {type(env_mask)}).") - - # update the curriculum for environments that need a reset - with Timer( - name="curriculum_manager.compute_reset", - msg="CurriculumManager.compute (reset) took:", - enable=DEBUG_TIMER_RESET, - time_unit="us", - ): - self.curriculum_manager.compute(env_ids=env_ids) + if env_mask is not self.reset_mask_wp: + wp.copy(self.reset_mask_wp, env_mask) + env_mask = self.reset_mask_wp + + self._manager_call_switch.call( + "CurriculumManager_compute", + self.curriculum_manager.compute, + env_mask=env_mask, + env_ids=env_ids, + _timer=DEBUG_TIMER_RESET, + ) - # reset the internal buffers of the scene elements - self._manager_call_switch.call_stage( - stage="Scene_reset", - warp_call={"fn": self.scene.reset, "kwargs": {"env_ids": env_ids, "env_mask": env_mask}}, - timer=DEBUG_TIMER_RESET, + self._manager_call_switch.call( + "Scene_reset", + self.scene.reset, + env_mask=env_mask, + _timer=DEBUG_TIMER_RESET, ) if "reset" in self.event_manager.available_modes: self._global_env_step_count_wp.fill_(self._sim_step_counter // self.cfg.decimation) - self._manager_call_switch.call_stage( - stage="EventManager_apply_reset", - warp_call={ - "fn": self.event_manager.apply, - "kwargs": { - "mode": "reset", - "env_mask_wp": env_mask, - "global_env_step_count": self._global_env_step_count_wp, - }, - }, - timer=DEBUG_TIMER_RESET, + self._manager_call_switch.call( + "EventManager_apply_reset", + self.event_manager.apply, + mode="reset", + env_mask_wp=env_mask, + global_env_step_count=self._global_env_step_count_wp, + _timer=DEBUG_TIMER_RESET, ) # iterate over all managers and reset them # this returns a dictionary of information which is stored in the extras # note: This is order-sensitive! Certain things need be reset before others. # -- observation manager + action + reward managers - obs_info = self._manager_call_switch.call_stage( - stage="ObservationManager_reset", - warp_call={"fn": self.observation_manager.reset, "kwargs": {"env_mask": env_mask}}, - timer=DEBUG_TIMER_RESET, + obs_info = self._manager_call_switch.call( + "ObservationManager_reset", + self.observation_manager.reset, + env_mask=env_mask, + _timer=DEBUG_TIMER_RESET, ) - action_info = self._manager_call_switch.call_stage( - stage="ActionManager_reset", - warp_call={"fn": self.action_manager.reset, "kwargs": {"env_mask": env_mask}}, - timer=DEBUG_TIMER_RESET, + action_info = self._manager_call_switch.call( + "ActionManager_reset", + self.action_manager.reset, + env_mask=env_mask, + _timer=DEBUG_TIMER_RESET, ) - reward_info = self._manager_call_switch.call_stage( - stage="RewardManager_reset", - warp_call={"fn": self.reward_manager.reset, "kwargs": {"env_mask": env_mask}}, - timer=DEBUG_TIMER_RESET, + reward_info = self._manager_call_switch.call( + "RewardManager_reset", + self.reward_manager.reset, + env_mask=env_mask, + _timer=DEBUG_TIMER_RESET, + ) + curriculum_info = self._manager_call_switch.call( + "CurriculumManager_reset", + self.curriculum_manager.reset, + env_mask=env_mask, + env_ids=env_ids, + _timer=DEBUG_TIMER_RESET, ) - - # -- curriculum manager - with Timer( - name="curriculum_manager.reset", - msg="CurriculumManager.reset took:", - enable=DEBUG_TIMER_RESET, - time_unit="us", - ): - curriculum_info = self.curriculum_manager.reset(env_ids=env_ids) # -- command + event + termination managers - command_info = self.command_manager.reset(env_ids=env_ids) - event_info = self._manager_call_switch.call_stage( - stage="EventManager_reset", - warp_call={"fn": self.event_manager.reset, "kwargs": {"env_mask": env_mask}}, - stable_call={"fn": self.event_manager.reset, "kwargs": {"env_ids": env_ids}}, - timer=DEBUG_TIMER_RESET, + command_info = self._manager_call_switch.call( + "CommandManager_reset", + self.command_manager.reset, + env_mask=env_mask, + _timer=DEBUG_TIMER_RESET, ) - termination_info = self._manager_call_switch.call_stage( - stage="TerminationManager_reset", - warp_call={"fn": self.termination_manager.reset, "kwargs": {"env_mask": env_mask}}, - stable_call={"fn": self.termination_manager.reset, "kwargs": {"env_ids": env_ids}}, - timer=DEBUG_TIMER_RESET, + event_info = self._manager_call_switch.call( + "EventManager_reset", + self.event_manager.reset, + env_mask=env_mask, + _timer=DEBUG_TIMER_RESET, + ) + termination_info = self._manager_call_switch.call( + "TerminationManager_reset", + self.termination_manager.reset, + env_mask=env_mask, + _timer=DEBUG_TIMER_RESET, ) - - # -- recorder manager - recorder_info = self.recorder_manager.reset(env_ids=env_ids) # reset the episode length buffer - self.episode_length_buf[env_ids] = 0 + wp.launch( + zero_masked_int64, + dim=self.num_envs, + inputs=[env_mask, self._episode_length_buf_wp], + device=self.device, + ) # aggregate logging info log: dict[str, Any] = {} @@ -610,7 +573,57 @@ def _reset_idx( command_info, event_info, termination_info, - recorder_info, ): log.update(info) self.extras["log"] = log + + def _reset_terminated_envs(self) -> None: + """Reset terminated environments, compacting IDs only for host consumers.""" + reset_mask = self.termination_manager.dones_wp + reset_env_ids = None + if self._reset_requires_host_selection(): + with Timer( + name="reset_selection_host", + msg="Host reset selection took:", + enable=DEBUG_TIMER_STEP, + time_unit="us", + ): + reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) + if reset_env_ids.numel() == 0: + return + if self._has_recorders: + self.recorder_manager.record_pre_reset(reset_env_ids) + self._reset_host_pre(reset_env_ids) + + with Timer( + name="reset_idx", + msg="Reset idx took:", + enable=DEBUG_TIMER_STEP, + time_unit="us", + ): + self._reset_idx(env_mask=reset_mask, env_ids=reset_env_ids) + + if reset_env_ids is not None and reset_env_ids.numel() > 0: + self.extras["log"].update(self._reset_host_post(reset_env_ids)) + if self.has_rtx_sensors and self.cfg.num_rerenders_on_reset > 0: + for _ in range(self.cfg.num_rerenders_on_reset): + self.sim.render() + if self._has_recorders: + self.recorder_manager.record_post_reset(reset_env_ids) + + def _reset_host_pre(self, env_ids: Sequence[int] | torch.Tensor) -> None: + """Run host-only pre-reset work for selected environments.""" + super()._reset_host_pre(env_ids) + + def _reset_host_post(self, env_ids: Sequence[int] | torch.Tensor) -> dict[str, Any]: + """Run host-only manager resets and return their logging values.""" + return super()._reset_host_post(env_ids) + + def _reset_requires_host_selection(self) -> bool: + """Return whether enabled reset features require a host-visible selection.""" + return bool( + self.curriculum_manager.requires_host_boundary + or self._has_recorders + or self.scene.surface_grippers + or (self.has_rtx_sensors and self.cfg.num_rerenders_on_reset > 0) + ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi index a6dadcfce5bc..804dbb689c7a 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi @@ -19,6 +19,12 @@ from .actions import ( # noqa: F401 JointPositionAction, JointPositionActionCfg, ) +from .commands import ( # noqa: F401 + UniformPoseCommand, + UniformPoseCommandCfg, + UniformVelocityCommand, + UniformVelocityCommandCfg, +) # Override stable terms with experimental Warp-first implementations. These leaf # modules are import-clean (no eager backend imports), so re-exporting them here diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.py new file mode 100644 index 000000000000..57c6a813798d --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.py @@ -0,0 +1,10 @@ +# 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 + +"""Warp-native command terms.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi new file mode 100644 index 000000000000..e30d38de19df --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi @@ -0,0 +1,15 @@ +# 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 + +__all__ = [ + "UniformPoseCommandCfg", + "UniformVelocityCommandCfg", + "UniformPoseCommand", + "UniformVelocityCommand", +] + +from .commands_cfg import UniformPoseCommandCfg, UniformVelocityCommandCfg +from .pose_command import UniformPoseCommand +from .velocity_command import UniformVelocityCommand diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/commands_cfg.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/commands_cfg.py new file mode 100644 index 000000000000..04d09e2a9948 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/commands_cfg.py @@ -0,0 +1,32 @@ +# 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 + +"""Configuration classes for Warp-native command terms.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from isaaclab.envs.mdp.commands.commands_cfg import UniformPoseCommandCfg as _UniformPoseCommandCfg +from isaaclab.envs.mdp.commands.commands_cfg import UniformVelocityCommandCfg as _UniformVelocityCommandCfg +from isaaclab.utils.configclass import configclass + +if TYPE_CHECKING: + from .pose_command import UniformPoseCommand + from .velocity_command import UniformVelocityCommand + + +@configclass +class UniformVelocityCommandCfg(_UniformVelocityCommandCfg): + """Configuration for the Warp-native uniform velocity command generator.""" + + class_type: type[UniformVelocityCommand] | str = "{DIR}.velocity_command:UniformVelocityCommand" + + +@configclass +class UniformPoseCommandCfg(_UniformPoseCommandCfg): + """Configuration for the Warp-native uniform pose command generator.""" + + class_type: type[UniformPoseCommand] | str = "{DIR}.pose_command:UniformPoseCommand" diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py new file mode 100644 index 000000000000..b88d3be3d546 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py @@ -0,0 +1,206 @@ +# 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 + +"""Warp-native pose command generator.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import warp as wp + +from isaaclab.assets import Articulation +from isaaclab.envs.mdp.commands._debug_vis import _PoseCommandDebugVis +from isaaclab.utils.leapp import POSE7_ELEMENT_NAMES + +from isaaclab_experimental.managers import CommandTerm + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from .commands_cfg import UniformPoseCommandCfg + + +@wp.kernel +def _resample_pose_command( + env_mask: wp.array(dtype=wp.bool), + rng_state: wp.array(dtype=wp.uint32), + command: wp.array(dtype=wp.float32, ndim=2), + pos_x_min: float, + pos_x_max: float, + pos_y_min: float, + pos_y_max: float, + pos_z_min: float, + pos_z_max: float, + roll_min: float, + roll_max: float, + pitch_min: float, + pitch_max: float, + yaw_min: float, + yaw_max: float, + make_quat_unique: bool, +): + env_id = wp.tid() + if env_mask[env_id]: + state = rng_state[env_id] + command[env_id, 0] = wp.randf(state, pos_x_min, pos_x_max) + command[env_id, 1] = wp.randf(state, pos_y_min, pos_y_max) + command[env_id, 2] = wp.randf(state, pos_z_min, pos_z_max) + + roll = wp.randf(state, roll_min, roll_max) + pitch = wp.randf(state, pitch_min, pitch_max) + yaw = wp.randf(state, yaw_min, yaw_max) + qx = wp.quat_from_axis_angle(wp.vec3f(1.0, 0.0, 0.0), roll) + qy = wp.quat_from_axis_angle(wp.vec3f(0.0, 1.0, 0.0), pitch) + qz = wp.quat_from_axis_angle(wp.vec3f(0.0, 0.0, 1.0), yaw) + quat = wp.mul(wp.mul(qz, qy), qx) + if make_quat_unique and quat[3] < 0.0: + quat = wp.quatf(-quat[0], -quat[1], -quat[2], -quat[3]) + command[env_id, 3] = quat[0] + command[env_id, 4] = quat[1] + command[env_id, 5] = quat[2] + command[env_id, 6] = quat[3] + rng_state[env_id] = state + + +@wp.kernel +def _update_pose_metrics( + command_b: wp.array(dtype=wp.float32, ndim=2), + root_pos_w: wp.array(dtype=wp.vec3f), + root_quat_w: wp.array(dtype=wp.quatf), + body_pos_w: wp.array(dtype=wp.vec3f, ndim=2), + body_quat_w: wp.array(dtype=wp.quatf, ndim=2), + body_idx: int, + command_w: wp.array(dtype=wp.float32, ndim=2), + position_error: wp.array(dtype=wp.float32), + orientation_error: wp.array(dtype=wp.float32), + success_rate: wp.array(dtype=wp.float32), + track_success: bool, + position_success_threshold: float, +): + env_id = wp.tid() + position_b = wp.vec3f(command_b[env_id, 0], command_b[env_id, 1], command_b[env_id, 2]) + orientation_b = wp.quatf(command_b[env_id, 3], command_b[env_id, 4], command_b[env_id, 5], command_b[env_id, 6]) + desired_position_w = root_pos_w[env_id] + wp.quat_rotate(root_quat_w[env_id], position_b) + desired_orientation_w = root_quat_w[env_id] * orientation_b + + command_w[env_id, 0] = desired_position_w[0] + command_w[env_id, 1] = desired_position_w[1] + command_w[env_id, 2] = desired_position_w[2] + command_w[env_id, 3] = desired_orientation_w[0] + command_w[env_id, 4] = desired_orientation_w[1] + command_w[env_id, 5] = desired_orientation_w[2] + command_w[env_id, 6] = desired_orientation_w[3] + + position_delta = body_pos_w[env_id, body_idx] - desired_position_w + position_error_value = wp.length(position_delta) + orientation_delta = body_quat_w[env_id, body_idx] * wp.quat_inverse(desired_orientation_w) + orientation_error_value = 2.0 * wp.acos(wp.clamp(wp.abs(orientation_delta[3]), 0.0, 1.0)) + position_error[env_id] = position_error_value + orientation_error[env_id] = orientation_error_value + if track_success and position_error_value < position_success_threshold: + success_rate[env_id] = 1.0 + + +class UniformPoseCommand(_PoseCommandDebugVis, CommandTerm): + """Generate a pose command from uniform position and Euler-angle distributions.""" + + cfg: UniformPoseCommandCfg + """Configuration for the command generator.""" + + def __init__(self, cfg: UniformPoseCommandCfg, env: ManagerBasedEnv): + """Initialize the command generator. + + Args: + cfg: Configuration for the command generator. + env: Environment containing the commanded articulation. + """ + super().__init__(cfg, env) + + self.robot: Articulation = env.scene[cfg.asset_name] + self.body_idx = self.robot.find_bodies(cfg.body_name)[0][0] + + self._pose_command_b_wp = wp.zeros((self.num_envs, 7), dtype=wp.float32, device=self.device) + self.pose_command_b = wp.to_torch(self._pose_command_b_wp) + self.pose_command_b[:, 6] = 1.0 + self._pose_command_w_wp = wp.zeros((self.num_envs, 7), dtype=wp.float32, device=self.device) + self.pose_command_w = wp.to_torch(self._pose_command_w_wp) + + self.metrics["position_error"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self.metrics["orientation_error"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._success_rate_wp = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._track_success = cfg.position_success_threshold is not None + if self._track_success: + self.metrics["success_rate"] = self._success_rate_wp + + self.cfg.cmd_kind = self.cfg.cmd_kind or "command/body/pose" + self.cfg.element_names = self.cfg.element_names or POSE7_ELEMENT_NAMES + + def __str__(self) -> str: + """Return a string representation of the command generator.""" + msg = "UniformPoseCommand:\n" + msg += f"\tCommand dimension: {tuple(self.command.shape[1:])}\n" + msg += f"\tResampling time range: {self.cfg.resampling_time_range}\n" + return msg + + @property + def command(self) -> torch.Tensor: + """Desired pose command [m, m, m, quaternion xyzw], shape ``(num_envs, 7)``.""" + return self.pose_command_b + + @property + def command_wp(self) -> wp.array(dtype=wp.float32, ndim=2): + """Pointer-stable desired pose command [m, m, m, quaternion xyzw].""" + return self._pose_command_b_wp + + def _update_metrics(self): + wp.launch( + kernel=_update_pose_metrics, + dim=self.num_envs, + inputs=[ + self._pose_command_b_wp, + self.robot.data.root_pos_w.warp, + self.robot.data.root_quat_w.warp, + self.robot.data.body_pos_w.warp, + self.robot.data.body_quat_w.warp, + self.body_idx, + self._pose_command_w_wp, + self.metrics["position_error"], + self.metrics["orientation_error"], + self._success_rate_wp, + self._track_success, + self.cfg.position_success_threshold or 0.0, + ], + device=self.device, + ) + + def _resample_command(self, env_mask: wp.array): + wp.launch( + kernel=_resample_pose_command, + dim=self.num_envs, + inputs=[ + env_mask, + self._env.rng_state_wp, + self._pose_command_b_wp, + self.cfg.ranges.pos_x[0], + self.cfg.ranges.pos_x[1], + self.cfg.ranges.pos_y[0], + self.cfg.ranges.pos_y[1], + self.cfg.ranges.pos_z[0], + self.cfg.ranges.pos_z[1], + self.cfg.ranges.roll[0], + self.cfg.ranges.roll[1], + self.cfg.ranges.pitch[0], + self.cfg.ranges.pitch[1], + self.cfg.ranges.yaw[0], + self.cfg.ranges.yaw[1], + self.cfg.make_quat_unique, + ], + device=self.device, + ) + + def _update_command(self): + pass diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py new file mode 100644 index 000000000000..57dae69d59cf --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py @@ -0,0 +1,299 @@ +# 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 + +"""Warp-native velocity command generator.""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch +import warp as wp + +from isaaclab.assets import Articulation +from isaaclab.envs.mdp.commands._debug_vis import _VelocityCommandDebugVis + +from isaaclab_experimental.managers import CommandTerm +from isaaclab_experimental.utils.warp import wrap_to_pi + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from .commands_cfg import UniformVelocityCommandCfg + +logger = logging.getLogger(__name__) + + +@wp.kernel +def _accumulate_velocity_metrics( + command: wp.array(dtype=wp.float32, ndim=2), + root_lin_vel_b: wp.array(dtype=wp.vec3f), + root_ang_vel_b: wp.array(dtype=wp.vec3f), + error_xy_sum: wp.array(dtype=wp.float32), + error_yaw_sum: wp.array(dtype=wp.float32), + step_count: wp.array(dtype=wp.float32), +): + env_id = wp.tid() + dx = command[env_id, 0] - root_lin_vel_b[env_id][0] + dy = command[env_id, 1] - root_lin_vel_b[env_id][1] + error_xy_sum[env_id] += wp.sqrt(dx * dx + dy * dy) + error_yaw_sum[env_id] += wp.abs(command[env_id, 2] - root_ang_vel_b[env_id][2]) + step_count[env_id] += 1.0 + + +@wp.kernel +def _finalize_velocity_metrics( + env_mask: wp.array(dtype=wp.bool), + error_xy_sum: wp.array(dtype=wp.float32), + error_yaw_sum: wp.array(dtype=wp.float32), + step_count: wp.array(dtype=wp.float32), + error_xy: wp.array(dtype=wp.float32), + error_yaw: wp.array(dtype=wp.float32), + success_rate: wp.array(dtype=wp.float32), + error_xy_threshold: float, + error_yaw_threshold: float, +): + env_id = wp.tid() + if env_mask[env_id]: + denominator = wp.max(step_count[env_id], 1.0) + mean_error_xy = error_xy_sum[env_id] / denominator + mean_error_yaw = error_yaw_sum[env_id] / denominator + error_xy[env_id] = mean_error_xy + error_yaw[env_id] = mean_error_yaw + success_rate[env_id] = wp.where( + (mean_error_xy < error_xy_threshold) and (mean_error_yaw < error_yaw_threshold), 1.0, 0.0 + ) + error_xy_sum[env_id] = 0.0 + error_yaw_sum[env_id] = 0.0 + step_count[env_id] = 0.0 + + +@wp.kernel +def _resample_velocity_command( + env_mask: wp.array(dtype=wp.bool), + rng_state: wp.array(dtype=wp.uint32), + command: wp.array(dtype=wp.float32, ndim=2), + heading_target: wp.array(dtype=wp.float32), + is_heading_env: wp.array(dtype=wp.bool), + is_standing_env: wp.array(dtype=wp.bool), + lin_vel_x_min: float, + lin_vel_x_max: float, + lin_vel_y_min: float, + lin_vel_y_max: float, + ang_vel_z_min: float, + ang_vel_z_max: float, + heading_min: float, + heading_max: float, + heading_command: bool, + rel_heading_envs: float, + rel_standing_envs: float, +): + env_id = wp.tid() + if env_mask[env_id]: + state = rng_state[env_id] + command[env_id, 0] = wp.randf(state, lin_vel_x_min, lin_vel_x_max) + command[env_id, 1] = wp.randf(state, lin_vel_y_min, lin_vel_y_max) + command[env_id, 2] = wp.randf(state, ang_vel_z_min, ang_vel_z_max) + if heading_command: + heading_target[env_id] = wp.randf(state, heading_min, heading_max) + is_heading_env[env_id] = wp.randf(state, 0.0, 1.0) <= rel_heading_envs + else: + is_heading_env[env_id] = False + is_standing_env[env_id] = wp.randf(state, 0.0, 1.0) <= rel_standing_envs + rng_state[env_id] = state + + +@wp.kernel +def _update_velocity_command( + command: wp.array(dtype=wp.float32, ndim=2), + heading_target: wp.array(dtype=wp.float32), + is_heading_env: wp.array(dtype=wp.bool), + is_standing_env: wp.array(dtype=wp.bool), + heading_w: wp.array(dtype=wp.float32), + heading_command: bool, + heading_control_stiffness: float, + ang_vel_z_min: float, + ang_vel_z_max: float, +): + env_id = wp.tid() + if heading_command and is_heading_env[env_id]: + heading_error = wrap_to_pi(heading_target[env_id] - heading_w[env_id]) + command[env_id, 2] = wp.clamp( + heading_control_stiffness * heading_error, + ang_vel_z_min, + ang_vel_z_max, + ) + if is_standing_env[env_id]: + command[env_id, 0] = 0.0 + command[env_id, 1] = 0.0 + command[env_id, 2] = 0.0 + + +class UniformVelocityCommand(_VelocityCommandDebugVis, CommandTerm): + """Generate an SE(2) velocity command from uniform distributions.""" + + cfg: UniformVelocityCommandCfg + """Configuration for the command generator.""" + + def __init__(self, cfg: UniformVelocityCommandCfg, env: ManagerBasedEnv): + """Initialize the command generator. + + Args: + cfg: Configuration for the command generator. + env: Environment containing the commanded articulation. + + Raises: + ValueError: If heading commands are enabled without a heading range. + """ + super().__init__(cfg, env) + + if cfg.heading_command and cfg.ranges.heading is None: + raise ValueError( + "The velocity command has heading commands active (heading_command=True) but the `ranges.heading`" + " parameter is set to None." + ) + if cfg.ranges.heading and not cfg.heading_command: + logger.warning( + f"The velocity command has the 'ranges.heading' attribute set to '{cfg.ranges.heading}'" + " but the heading command is not active. Consider setting the flag for the heading command to True." + ) + + self.robot: Articulation = env.scene[cfg.asset_name] + + self._vel_command_b_wp = wp.zeros((self.num_envs, 3), dtype=wp.float32, device=self.device) + self.vel_command_b = wp.to_torch(self._vel_command_b_wp) + self._heading_target_wp = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self.heading_target = wp.to_torch(self._heading_target_wp) + self._is_heading_env_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + self.is_heading_env = wp.to_torch(self._is_heading_env_wp) + self._is_standing_env_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + self.is_standing_env = wp.to_torch(self._is_standing_env_wp) + + self.metrics["error_vel_xy"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self.metrics["error_vel_yaw"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self.metrics["success_rate"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._error_xy_sum_wp = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._error_yaw_sum_wp = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._step_count_wp = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + + self.cfg.cmd_kind = self.cfg.cmd_kind or "command/body/velocity" + self.cfg.element_names = self.cfg.element_names or ["lin_vel_x", "lin_vel_y", "ang_vel_z"] + + def __str__(self) -> str: + """Return a string representation of the command generator.""" + msg = "UniformVelocityCommand:\n" + msg += f"\tCommand dimension: {tuple(self.command.shape[1:])}\n" + msg += f"\tResampling time range: {self.cfg.resampling_time_range}\n" + msg += f"\tHeading command: {self.cfg.heading_command}\n" + if self.cfg.heading_command: + msg += f"\tHeading probability: {self.cfg.rel_heading_envs}\n" + msg += f"\tStanding probability: {self.cfg.rel_standing_envs}" + return msg + + @property + def command(self) -> torch.Tensor: + """Desired base velocity command [m/s, m/s, rad/s], shape ``(num_envs, 3)``.""" + return self.vel_command_b + + @property + def command_wp(self) -> wp.array(dtype=wp.float32, ndim=2): + """Pointer-stable desired base velocity command [m/s, m/s, rad/s].""" + return self._vel_command_b_wp + + def reset( + self, + env_ids: Sequence[int] | torch.Tensor | None = None, + *, + env_mask: wp.array | None = None, + ) -> dict[str, torch.Tensor]: + """Finalize episode metrics and reset selected command state. + + Args: + env_ids: Environment indices used by compatibility call sites. + env_mask: Boolean Warp mask selecting environments to reset. Takes precedence over ``env_ids``. + + Returns: + Persistent scalar metric views for logging. + """ + env_mask = self._resolve_reset_mask(env_ids, env_mask) + wp.launch( + kernel=_finalize_velocity_metrics, + dim=self.num_envs, + inputs=[ + env_mask, + self._error_xy_sum_wp, + self._error_yaw_sum_wp, + self._step_count_wp, + self.metrics["error_vel_xy"], + self.metrics["error_vel_yaw"], + self.metrics["success_rate"], + self.cfg.vel_xy_success_threshold, + self.cfg.vel_yaw_success_threshold, + ], + device=self.device, + ) + return super().reset(env_mask=env_mask) + + def _update_metrics(self): + wp.launch( + kernel=_accumulate_velocity_metrics, + dim=self.num_envs, + inputs=[ + self._vel_command_b_wp, + self.robot.data.root_lin_vel_b.warp, + self.robot.data.root_ang_vel_b.warp, + self._error_xy_sum_wp, + self._error_yaw_sum_wp, + self._step_count_wp, + ], + device=self.device, + ) + + def _resample_command(self, env_mask: wp.array): + heading_range = self.cfg.ranges.heading if self.cfg.ranges.heading is not None else (0.0, 0.0) + wp.launch( + kernel=_resample_velocity_command, + dim=self.num_envs, + inputs=[ + env_mask, + self._env.rng_state_wp, + self._vel_command_b_wp, + self._heading_target_wp, + self._is_heading_env_wp, + self._is_standing_env_wp, + self.cfg.ranges.lin_vel_x[0], + self.cfg.ranges.lin_vel_x[1], + self.cfg.ranges.lin_vel_y[0], + self.cfg.ranges.lin_vel_y[1], + self.cfg.ranges.ang_vel_z[0], + self.cfg.ranges.ang_vel_z[1], + heading_range[0], + heading_range[1], + self.cfg.heading_command, + self.cfg.rel_heading_envs, + self.cfg.rel_standing_envs, + ], + device=self.device, + ) + + def _update_command(self): + wp.launch( + kernel=_update_velocity_command, + dim=self.num_envs, + inputs=[ + self._vel_command_b_wp, + self._heading_target_wp, + self._is_heading_env_wp, + self._is_standing_env_wp, + self.robot.data.heading_w.warp, + self.cfg.heading_command, + self.cfg.heading_control_stiffness, + self.cfg.ranges.ang_vel_z[0], + self.cfg.ranges.ang_vel_z[1], + ], + device=self.device, + ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py index d9e61041a7ad..0a4c3c274b9e 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py @@ -24,7 +24,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import warp as wp @@ -32,7 +32,138 @@ from isaaclab_experimental.utils.warp import WarpCapturable if TYPE_CHECKING: - from isaaclab.assets import Articulation + from isaaclab.assets import Articulation, RigidObject + from isaaclab.envs import ManagerBasedEnv + + +def _range_cache_key( + range_values: dict[str, tuple[float, float]] | tuple[float, float], +) -> tuple[object, ...]: + """Return an immutable snapshot of an event range for state-cache lookup.""" + if isinstance(range_values, dict): + return tuple((name, tuple(bounds)) for name, bounds in sorted(range_values.items())) + return tuple(range_values) + + +class _EventStateCache: + """Environment-owned storage for persistent Warp event buffers and constants. + + The cache lifetime matches the environment lifetime. Keys include the event + callable, asset, term configuration, and immutable range values so stale parsed + constants are not reused after a range changes. + """ + + _ENV_ATTRIBUTE: ClassVar[str] = "_warp_event_state_cache" + + @classmethod + def for_env(cls, env: ManagerBasedEnv) -> dict[tuple[object, ...], object]: + """Return the event-state cache owned by an environment.""" + cache = getattr(env, cls._ENV_ATTRIBUTE, None) + if cache is None: + cache = {} + setattr(env, cls._ENV_ATTRIBUTE, cache) + return cache + + +class _RandomizeRigidBodyComState: + """Persistent arguments for center-of-mass randomization.""" + + def __init__( + self, + asset: RigidObject | Articulation, + asset_cfg: SceneEntityCfg, + com_range: dict[str, tuple[float, float]], + ): + self.asset = asset + self.asset_cfg = asset_cfg + self.com_range = com_range + self.default_com = wp.clone(asset.data.body_com_pos_b.warp) + ranges = [com_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z")] + self.com_lo = wp.vec3f(ranges[0][0], ranges[1][0], ranges[2][0]) + self.com_hi = wp.vec3f(ranges[0][1], ranges[1][1], ranges[2][1]) + + +class _ApplyExternalForceTorqueState: + """Persistent output buffers for external wrench randomization.""" + + def __init__( + self, + env: ManagerBasedEnv, + asset: RigidObject | Articulation, + asset_cfg: SceneEntityCfg, + force_range: tuple[float, float], + torque_range: tuple[float, float], + ): + self.asset = asset + self.asset_cfg = asset_cfg + self.force_range = force_range + self.torque_range = torque_range + self.forces = wp.zeros((env.num_envs, asset.num_bodies), dtype=wp.vec3f, device=env.device) + self.torques = wp.zeros((env.num_envs, asset.num_bodies), dtype=wp.vec3f, device=env.device) + if asset_cfg.body_ids is None or asset_cfg.body_ids == slice(None): + body_ids = list(range(asset.num_bodies)) + elif isinstance(asset_cfg.body_ids, int): + body_ids = [asset_cfg.body_ids] + else: + body_ids = list(asset_cfg.body_ids) + body_mask = [False] * asset.num_bodies + for body_id in body_ids: + body_mask[body_id] = True + self.body_ids = wp.array(body_ids, dtype=wp.int32, device=env.device) + self.body_mask = wp.array(body_mask, dtype=wp.bool, device=env.device) + + +class _PushBySettingVelocityState: + """Persistent output buffer and ranges for root-velocity pushes.""" + + def __init__( + self, + env: ManagerBasedEnv, + asset: RigidObject | Articulation, + asset_cfg: SceneEntityCfg, + velocity_range: dict[str, tuple[float, float]], + ): + self.asset = asset + self.asset_cfg = asset_cfg + self.velocity_range = velocity_range + self.velocity = wp.zeros((env.num_envs,), dtype=wp.spatial_vectorf, device=env.device) + ranges = [velocity_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z", "roll", "pitch", "yaw")] + self.lin_lo = wp.vec3f(ranges[0][0], ranges[1][0], ranges[2][0]) + self.lin_hi = wp.vec3f(ranges[0][1], ranges[1][1], ranges[2][1]) + self.ang_lo = wp.vec3f(ranges[3][0], ranges[4][0], ranges[5][0]) + self.ang_hi = wp.vec3f(ranges[3][1], ranges[4][1], ranges[5][1]) + + +class _ResetRootStateUniformState: + """Persistent output buffers and ranges for uniform root-state reset.""" + + def __init__( + self, + env: ManagerBasedEnv, + asset: RigidObject | Articulation, + asset_cfg: SceneEntityCfg, + pose_range: dict[str, tuple[float, float]], + velocity_range: dict[str, tuple[float, float]], + ): + self.asset = asset + self.asset_cfg = asset_cfg + self.pose_range = pose_range + self.velocity_range = velocity_range + self.pose = wp.zeros((env.num_envs,), dtype=wp.transformf, device=env.device) + self.velocity = wp.zeros((env.num_envs,), dtype=wp.spatial_vectorf, device=env.device) + + pose_ranges = [pose_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z", "roll", "pitch", "yaw")] + self.pos_lo = wp.vec3f(pose_ranges[0][0], pose_ranges[1][0], pose_ranges[2][0]) + self.pos_hi = wp.vec3f(pose_ranges[0][1], pose_ranges[1][1], pose_ranges[2][1]) + self.rot_lo = wp.vec3f(pose_ranges[3][0], pose_ranges[4][0], pose_ranges[5][0]) + self.rot_hi = wp.vec3f(pose_ranges[3][1], pose_ranges[4][1], pose_ranges[5][1]) + + velocity_ranges = [velocity_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z", "roll", "pitch", "yaw")] + self.vel_lin_lo = wp.vec3f(velocity_ranges[0][0], velocity_ranges[1][0], velocity_ranges[2][0]) + self.vel_lin_hi = wp.vec3f(velocity_ranges[0][1], velocity_ranges[1][1], velocity_ranges[2][1]) + self.vel_ang_lo = wp.vec3f(velocity_ranges[3][0], velocity_ranges[4][0], velocity_ranges[5][0]) + self.vel_ang_hi = wp.vec3f(velocity_ranges[3][1], velocity_ranges[4][1], velocity_ranges[5][1]) + # --------------------------------------------------------------------------- # Randomize rigid body center of mass @@ -43,23 +174,24 @@ def _randomize_com_kernel( env_mask: wp.array(dtype=wp.bool), rng_state: wp.array(dtype=wp.uint32), + default_body_com_pos_b: wp.array(dtype=wp.vec3f, ndim=2), body_com_pos_b: wp.array(dtype=wp.vec3f, ndim=2), body_ids: wp.array(dtype=wp.int32), com_lo: wp.vec3f, com_hi: wp.vec3f, ): - """Add random offset to center of mass positions for selected bodies.""" + """Add random offsets to the default center-of-mass positions for selected bodies.""" env_id = wp.tid() if not env_mask[env_id]: return state = rng_state[env_id] + dx = wp.randf(state, com_lo[0], com_hi[0]) + dy = wp.randf(state, com_lo[1], com_hi[1]) + dz = wp.randf(state, com_lo[2], com_hi[2]) for k in range(body_ids.shape[0]): b = body_ids[k] - v = body_com_pos_b[env_id, b] - dx = wp.randf(state, com_lo[0], com_hi[0]) - dy = wp.randf(state, com_lo[1], com_hi[1]) - dz = wp.randf(state, com_lo[2], com_hi[2]) + v = default_body_com_pos_b[env_id, b] body_com_pos_b[env_id, b] = wp.vec3f(v[0] + dx, v[1] + dy, v[2] + dz) rng_state[env_id] = state @@ -78,14 +210,12 @@ def randomize_rigid_body_com( via :meth:`set_coms_mask` so it recomputes inertial properties. """ asset: Articulation = env.scene[asset_cfg.name] - - fn = randomize_rigid_body_com - if not getattr(fn, "_is_warmed_up", False) or fn._asset_name != asset_cfg.name: - fn._asset_name = asset_cfg.name - r = [com_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] - fn._com_lo = wp.vec3f(r[0][0], r[1][0], r[2][0]) - fn._com_hi = wp.vec3f(r[0][1], r[1][1], r[2][1]) - fn._is_warmed_up = True + cache = _EventStateCache.for_env(env) + cache_key = (randomize_rigid_body_com, id(asset), id(asset_cfg), _range_cache_key(com_range)) + state = cache.get(cache_key) + if state is None: + state = _RandomizeRigidBodyComState(asset, asset_cfg, com_range) + cache[cache_key] = state wp.launch( kernel=_randomize_com_kernel, @@ -93,10 +223,11 @@ def randomize_rigid_body_com( inputs=[ env_mask, env.rng_state_wp, + state.default_com, asset.data.body_com_pos_b.warp, - asset_cfg.body_ids_wp, - fn._com_lo, - fn._com_hi, + state.asset_cfg.body_ids_wp, + state.com_lo, + state.com_hi, ], device=env.device, ) @@ -114,6 +245,7 @@ def randomize_rigid_body_com( def _apply_external_force_torque_kernel( env_mask: wp.array(dtype=wp.bool), rng_state: wp.array(dtype=wp.uint32), + body_ids: wp.array(dtype=wp.int32), force_out: wp.array(dtype=wp.vec3f, ndim=2), torque_out: wp.array(dtype=wp.vec3f, ndim=2), force_lo: float, @@ -123,14 +255,11 @@ def _apply_external_force_torque_kernel( ): env_id = wp.tid() if not env_mask[env_id]: - # zero out unmasked envs so they don't accumulate stale forces - for b in range(force_out.shape[1]): - force_out[env_id, b] = wp.vec3f(0.0, 0.0, 0.0) - torque_out[env_id, b] = wp.vec3f(0.0, 0.0, 0.0) return state = rng_state[env_id] - for b in range(force_out.shape[1]): + for body_index in range(body_ids.shape[0]): + b = body_ids[body_index] force_out[env_id, b] = wp.vec3f( wp.randf(state, force_lo, force_hi), wp.randf(state, force_lo, force_hi), @@ -156,16 +285,18 @@ def apply_external_force_torque( Warp-first override of :func:`isaaclab.envs.mdp.events.apply_external_force_torque`. """ asset: Articulation = env.scene[asset_cfg.name] - - # First-call: allocate scratch and pre-convert constant arguments. - if not getattr(apply_external_force_torque, "_is_warmed_up", False): - apply_external_force_torque._scratch_forces = wp.zeros( - (env.num_envs, asset.num_bodies), dtype=wp.vec3f, device=env.device - ) - apply_external_force_torque._scratch_torques = wp.zeros( - (env.num_envs, asset.num_bodies), dtype=wp.vec3f, device=env.device - ) - apply_external_force_torque._is_warmed_up = True + cache = _EventStateCache.for_env(env) + cache_key = ( + apply_external_force_torque, + id(asset), + id(asset_cfg), + _range_cache_key(force_range), + _range_cache_key(torque_range), + ) + state = cache.get(cache_key) + if state is None: + state = _ApplyExternalForceTorqueState(env, asset, asset_cfg, force_range, torque_range) + cache[cache_key] = state wp.launch( kernel=_apply_external_force_torque_kernel, @@ -173,19 +304,21 @@ def apply_external_force_torque( inputs=[ env_mask, env.rng_state_wp, - apply_external_force_torque._scratch_forces, - apply_external_force_torque._scratch_torques, - force_range[0], - force_range[1], - torque_range[0], - torque_range[1], + state.body_ids, + state.forces, + state.torques, + state.force_range[0], + state.force_range[1], + state.torque_range[0], + state.torque_range[1], ], device=env.device, ) asset.permanent_wrench_composer.set_forces_and_torques_mask( - forces=apply_external_force_torque._scratch_forces, - torques=apply_external_force_torque._scratch_torques, + forces=state.forces, + torques=state.torques, + body_mask=state.body_mask, env_mask=env_mask, ) @@ -236,16 +369,12 @@ def push_by_setting_velocity( Warp-first override of :func:`isaaclab.envs.mdp.events.push_by_setting_velocity`. """ asset: Articulation = env.scene[asset_cfg.name] - - # First-call: allocate scratch and pre-parse constant range arguments. - if not getattr(push_by_setting_velocity, "_is_warmed_up", False): - push_by_setting_velocity._scratch_vel = wp.zeros((env.num_envs,), dtype=wp.spatial_vectorf, device=env.device) - r = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] - push_by_setting_velocity._lin_lo = wp.vec3f(r[0][0], r[1][0], r[2][0]) - push_by_setting_velocity._lin_hi = wp.vec3f(r[0][1], r[1][1], r[2][1]) - push_by_setting_velocity._ang_lo = wp.vec3f(r[3][0], r[4][0], r[5][0]) - push_by_setting_velocity._ang_hi = wp.vec3f(r[3][1], r[4][1], r[5][1]) - push_by_setting_velocity._is_warmed_up = True + cache = _EventStateCache.for_env(env) + cache_key = (push_by_setting_velocity, id(asset), id(asset_cfg), _range_cache_key(velocity_range)) + state = cache.get(cache_key) + if state is None: + state = _PushBySettingVelocityState(env, asset, asset_cfg, velocity_range) + cache[cache_key] = state wp.launch( kernel=_push_by_setting_velocity_kernel, @@ -254,16 +383,16 @@ def push_by_setting_velocity( env_mask, env.rng_state_wp, asset.data.root_vel_w.warp, - push_by_setting_velocity._scratch_vel, - push_by_setting_velocity._lin_lo, - push_by_setting_velocity._lin_hi, - push_by_setting_velocity._ang_lo, - push_by_setting_velocity._ang_hi, + state.velocity, + state.lin_lo, + state.lin_hi, + state.ang_lo, + state.ang_hi, ], device=env.device, ) - asset.write_root_velocity_to_sim_mask(root_velocity=push_by_setting_velocity._scratch_vel, env_mask=env_mask) + asset.write_root_velocity_to_sim_mask(root_velocity=state.velocity, env_mask=env_mask) # --------------------------------------------------------------------------- @@ -347,24 +476,18 @@ def reset_root_state_uniform( Warp-first override of :func:`isaaclab.envs.mdp.events.reset_root_state_uniform`. """ asset: Articulation = env.scene[asset_cfg.name] - - # First-call: allocate scratch and pre-parse range dicts. - if not getattr(reset_root_state_uniform, "_is_warmed_up", False): - reset_root_state_uniform._scratch_pose = wp.zeros((env.num_envs,), dtype=wp.transformf, device=env.device) - reset_root_state_uniform._scratch_vel = wp.zeros((env.num_envs,), dtype=wp.spatial_vectorf, device=env.device) - # Pre-parse pose_range dict - p = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] - reset_root_state_uniform._pos_lo = wp.vec3f(p[0][0], p[1][0], p[2][0]) - reset_root_state_uniform._pos_hi = wp.vec3f(p[0][1], p[1][1], p[2][1]) - reset_root_state_uniform._rot_lo = wp.vec3f(p[3][0], p[4][0], p[5][0]) - reset_root_state_uniform._rot_hi = wp.vec3f(p[3][1], p[4][1], p[5][1]) - # Pre-parse velocity_range dict - v = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] - reset_root_state_uniform._vel_lin_lo = wp.vec3f(v[0][0], v[1][0], v[2][0]) - reset_root_state_uniform._vel_lin_hi = wp.vec3f(v[0][1], v[1][1], v[2][1]) - reset_root_state_uniform._vel_ang_lo = wp.vec3f(v[3][0], v[4][0], v[5][0]) - reset_root_state_uniform._vel_ang_hi = wp.vec3f(v[3][1], v[4][1], v[5][1]) - reset_root_state_uniform._is_warmed_up = True + cache = _EventStateCache.for_env(env) + cache_key = ( + reset_root_state_uniform, + id(asset), + id(asset_cfg), + _range_cache_key(pose_range), + _range_cache_key(velocity_range), + ) + state = cache.get(cache_key) + if state is None: + state = _ResetRootStateUniformState(env, asset, asset_cfg, pose_range, velocity_range) + cache[cache_key] = state wp.launch( kernel=_reset_root_state_uniform_kernel, @@ -375,22 +498,22 @@ def reset_root_state_uniform( asset.data.default_root_pose.warp, asset.data.default_root_vel.warp, env.env_origins_wp, - reset_root_state_uniform._scratch_pose, - reset_root_state_uniform._scratch_vel, - reset_root_state_uniform._pos_lo, - reset_root_state_uniform._pos_hi, - reset_root_state_uniform._rot_lo, - reset_root_state_uniform._rot_hi, - reset_root_state_uniform._vel_lin_lo, - reset_root_state_uniform._vel_lin_hi, - reset_root_state_uniform._vel_ang_lo, - reset_root_state_uniform._vel_ang_hi, + state.pose, + state.velocity, + state.pos_lo, + state.pos_hi, + state.rot_lo, + state.rot_hi, + state.vel_lin_lo, + state.vel_lin_hi, + state.vel_ang_lo, + state.vel_ang_hi, ], device=env.device, ) - asset.write_root_pose_to_sim_mask(root_pose=reset_root_state_uniform._scratch_pose, env_mask=env_mask) - asset.write_root_velocity_to_sim_mask(root_velocity=reset_root_state_uniform._scratch_vel, env_mask=env_mask) + asset.write_root_pose_to_sim_mask(root_pose=state.pose, env_mask=env_mask) + asset.write_root_velocity_to_sim_mask(root_velocity=state.velocity, env_mask=env_mask) # --------------------------------------------------------------------------- diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py index 808231c8faf6..fb0320251a9c 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py @@ -353,17 +353,6 @@ def generated_commands(env: ManagerBasedEnv, out, command_name: str) -> None: """The generated command from the command manager. Writes into ``out``. Warp-first override of :func:`isaaclab.envs.mdp.observations.generated_commands`. - Uses ``wp.from_torch`` to create a zero-copy warp view of the command tensor on first call. + Reads the command manager's pointer-stable Warp storage directly. """ - # TODO(warp-migration): Cross-manager access (observation → command). Replace with direct - # warp getter once all managers are guaranteed to be warp-native. - fn = generated_commands - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - if isinstance(cmd, wp.array): - fn._cmd_wp = cmd - else: - fn._cmd_wp = wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True - wp.copy(out, fn._cmd_wp) + wp.copy(out, env.command_manager.get_command_wp(command_name)) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py index 338d27c99c5e..c4bde2e7b5a4 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py @@ -415,24 +415,13 @@ def track_lin_vel_xy_exp( Warp-first override of :func:`isaaclab.envs.mdp.rewards.track_lin_vel_xy_exp`. """ asset: Articulation = env.scene[asset_cfg.name] - # cache the warp view of the command tensor on first call (zero-copy) - # TODO(warp-migration): Cross-manager access (reward → command). Replace with direct - # warp getter once all managers are guaranteed to be warp-native. - if not getattr(track_lin_vel_xy_exp, "_is_warmed_up", False) or track_lin_vel_xy_exp._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - if isinstance(cmd, wp.array): - track_lin_vel_xy_exp._cmd_wp = cmd - else: - track_lin_vel_xy_exp._cmd_wp = wp.from_torch(cmd) - track_lin_vel_xy_exp._cmd_name = command_name - track_lin_vel_xy_exp._is_warmed_up = True wp.launch( kernel=_track_lin_vel_xy_exp_kernel, dim=env.num_envs, inputs=[ asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, - track_lin_vel_xy_exp._cmd_wp, + env.command_manager.get_command_wp(command_name), 1.0 / (std * std), out, ], @@ -466,23 +455,13 @@ def track_ang_vel_z_exp( Warp-first override of :func:`isaaclab.envs.mdp.rewards.track_ang_vel_z_exp`. """ asset: Articulation = env.scene[asset_cfg.name] - # TODO(warp-migration): Cross-manager access (reward → command). Replace with direct - # warp getter once all managers are guaranteed to be warp-native. - if not getattr(track_ang_vel_z_exp, "_is_warmed_up", False) or track_ang_vel_z_exp._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - if isinstance(cmd, wp.array): - track_ang_vel_z_exp._cmd_wp = cmd - else: - track_ang_vel_z_exp._cmd_wp = wp.from_torch(cmd) - track_ang_vel_z_exp._cmd_name = command_name - track_ang_vel_z_exp._is_warmed_up = True wp.launch( kernel=_track_ang_vel_z_exp_kernel, dim=env.num_envs, inputs=[ asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, - track_ang_vel_z_exp._cmd_wp, + env.command_manager.get_command_wp(command_name), 2, 1.0 / (std * std), out, diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/managers/__init__.pyi index 18431a745b31..29c0dfa1b494 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/__init__.pyi +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/__init__.pyi @@ -8,6 +8,7 @@ __all__ = [ "ActionTerm", "CommandManager", "CommandTerm", + "CurriculumManager", "EventManager", "ManagerBase", "ManagerTermBase", @@ -31,6 +32,7 @@ from isaaclab.managers import * # noqa: F401, F403 from .action_manager import ActionManager, ActionTerm from .command_manager import CommandManager, CommandTerm +from .curriculum_manager import CurriculumManager from .event_manager import EventManager from .manager_base import ManagerBase, ManagerTermBase from .manager_term_cfg import ( diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py index 672af82bc8ab..0ccee0e6d381 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py @@ -373,32 +373,21 @@ def reset( Returns: An empty dictionary. """ - # Mask-first path: captured callers must provide env_mask. - if env_mask is None or not isinstance(env_mask, wp.array): - if wp.get_device().is_capturing: - raise RuntimeError( - "ActionManager.reset requires env_mask(wp.array[bool]) during capture. " - "Do not pass env_ids on captured paths." - ) - env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + env_mask = self._resolve_reset_mask(env_ids, env_mask) # reset the action history - if env_mask is None: - self._prev_action.fill_(0.0) - self._action.fill_(0.0) - else: - wp.launch( - kernel=_zero_masked_2d, - dim=(self.num_envs, self.total_action_dim), - inputs=[env_mask, self._prev_action], - device=self.device, - ) - wp.launch( - kernel=_zero_masked_2d, - dim=(self.num_envs, self.total_action_dim), - inputs=[env_mask, self._action], - device=self.device, - ) + wp.launch( + kernel=_zero_masked_2d, + dim=(self.num_envs, self.total_action_dim), + inputs=[env_mask, self._prev_action], + device=self.device, + ) + wp.launch( + kernel=_zero_masked_2d, + dim=(self.num_envs, self.total_action_dim), + inputs=[env_mask, self._action], + device=self.device, + ) # reset all action terms for term in self._terms.values(): diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py index 8b4e8f83dbe7..496bc17b87fd 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py @@ -8,7 +8,6 @@ from __future__ import annotations import inspect -import weakref from abc import abstractmethod from collections.abc import Sequence from typing import TYPE_CHECKING @@ -124,9 +123,11 @@ def __init__(self, cfg: CommandTermCfg, env: ManagerBasedRLEnv): def __del__(self): """Unsubscribe from the callbacks.""" - if self._debug_vis_handle: - self._debug_vis_handle.unsubscribe() - self._debug_vis_handle = None + env = getattr(self, "_env", None) + sim = getattr(env, "sim", None) + registry = getattr(sim, "vis_marker_registry", None) + if registry is not None: + registry.clear_debug_vis_callback(self) """ Properties @@ -138,6 +139,21 @@ def command(self) -> torch.Tensor | wp.array: """The command tensor. Shape is (num_envs, command_dim).""" raise NotImplementedError + @property + def command_wp(self) -> wp.array(dtype=wp.float32, ndim=2): + """Pointer-stable Warp command storage with shape ``(num_envs, command_dim)``. + + Existing custom terms that only implement :attr:`command` receive a + lazily cached zero-copy Warp view. Warp-native terms may override this + property to return their owner-held array directly. + """ + command = self.command + if isinstance(command, wp.array): + return command + if not hasattr(self, "_command_wp_compat"): + self._command_wp_compat = wp.from_torch(command, dtype=wp.float32) + return self._command_wp_compat + @property def has_debug_vis_implementation(self) -> bool: """Whether the command generator has a debug visualization implemented.""" @@ -171,25 +187,10 @@ def set_debug_vis(self, debug_vis: bool) -> bool: self._set_debug_vis_impl(debug_vis) # toggle debug visualization handles if debug_vis: - # only enable debug_vis if omniverse is available - from isaaclab.sim.simulation_context import SimulationContext - - sim_context = SimulationContext.instance() - if not sim_context.has_omniverse_visualizer(): - return False - # create a subscriber for the post update event if it doesn't exist if self._debug_vis_handle is None: - import omni.kit.app - - app_interface = omni.kit.app.get_app_interface() - self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop( - lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event) - ) + self._debug_vis_handle = self._env.sim.vis_marker_registry.add_debug_vis_callback(self) else: - # remove the subscriber if it exists - if self._debug_vis_handle is not None: - self._debug_vis_handle.unsubscribe() - self._debug_vis_handle = None + self._env.sim.vis_marker_registry.clear_debug_vis_callback(self) # return success return True @@ -213,14 +214,7 @@ def reset( Returns: A dictionary containing the information to log under the "{name}" key. """ - # Mask-first path: captured callers must provide env_mask. - if env_mask is None or not isinstance(env_mask, wp.array): - if wp.get_device().is_capturing: - raise RuntimeError( - "CommandTerm.reset requires env_mask(wp.array[bool]) during capture. " - "Do not pass env_ids on captured paths." - ) - env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + env_mask = self._resolve_reset_mask(env_ids, env_mask) # compute selected count and reset scale self._reset_count_wp.zero_() @@ -410,9 +404,14 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): # reset logging extras (persistent holder for orchestrator aggregation) self._reset_extras: dict[str, torch.Tensor] = {} + success_term_count = sum("success_rate" in term.reset_extras for term in self._terms.values()) for term_name, term in self._terms.items(): for metric_name, metric_value in term.reset_extras.items(): - self._reset_extras[f"Metrics/{term_name}/{metric_name}"] = metric_value + if metric_name == "success_rate" and success_term_count == 1: + metric_key = "Metrics/success_rate" + else: + metric_key = f"Metrics/{term_name}/{metric_name}" + self._reset_extras[metric_key] = metric_value def __str__(self) -> str: """Returns: A string representation for the command manager.""" @@ -513,14 +512,7 @@ def reset( Returns: A dictionary containing the information to log under the "Metrics/{term_name}/{metric_name}" key. """ - # Mask-first path: captured callers must provide env_mask. - if env_mask is None or not isinstance(env_mask, wp.array): - if wp.get_device().is_capturing: - raise RuntimeError( - "CommandManager.reset requires env_mask(wp.array[bool]) during capture. " - "Do not pass env_ids on captured paths." - ) - env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + env_mask = self._resolve_reset_mask(env_ids, env_mask) for term in self._terms.values(): # reset the command term @@ -556,6 +548,17 @@ def get_command(self, name: str) -> torch.Tensor: return wp.to_torch(command) return command + def get_command_wp(self, name: str) -> wp.array(dtype=wp.float32, ndim=2): + """Return pointer-stable Warp storage for a command term. + + Args: + name: The name of the command term. + + Returns: + The command array with shape ``(num_envs, command_dim)``. + """ + return self._terms[name].command_wp + def get_term(self, name: str) -> CommandTerm: """Returns the command term with the specified name. diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py new file mode 100644 index 000000000000..611246aeb55d --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py @@ -0,0 +1,245 @@ +# 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 + +"""Warp-first curriculum manager with mask-native term dispatch.""" + +from __future__ import annotations + +import inspect +from collections.abc import Sequence +from numbers import Real +from typing import TYPE_CHECKING + +import torch +import warp as wp +from prettytable import PrettyTable + +from isaaclab.managers import CurriculumTermCfg as StableCurriculumTermCfg +from isaaclab.managers import ManagerTermBase as StableManagerTermBase +from isaaclab.utils import string_to_callable + +from .manager_base import ManagerBase, ManagerTermBase + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +class CurriculumManager(ManagerBase): + """Manage curriculum terms that write state into persistent Warp scalar outputs.""" + + _env: ManagerBasedRLEnv + """The environment instance.""" + + def __init__(self, cfg: object, env: ManagerBasedRLEnv): + """Initialize the curriculum manager. + + Args: + cfg: Configuration object or dictionary of curriculum terms. + env: Environment instance. + """ + self._term_names: list[str] = [] + self._term_cfgs: list[StableCurriculumTermCfg] = [] + self._term_modes: list[str] = [] + + super().__init__(cfg, env) + + num_terms = len(self._term_names) + self._term_states_wp = wp.zeros(num_terms, dtype=wp.float32, device=self.device) + self._term_state_views_wp: list[wp.array] = [] + if num_terms > 0: + stride = self._term_states_wp.strides[0] + for term_idx, term_cfg in enumerate(self._term_cfgs): + out_view = wp.array( + ptr=self._term_states_wp.ptr + term_idx * stride, + dtype=wp.float32, + shape=(1,), + strides=(stride,), + device=self.device, + ) + self._term_state_views_wp.append(out_view) + term_cfg.out = out_view + + self._term_states = wp.to_torch(self._term_states_wp) + self._reset_extras = { + f"Curriculum/{term_name}": self._term_states[term_idx] + for term_idx, term_name in enumerate(self._term_names) + } + + def __str__(self) -> str: + """Return a string representation of the curriculum manager.""" + msg = f" contains {len(self._term_names)} active terms.\n" + table = PrettyTable() + table.title = "Active Curriculum Terms" + table.field_names = ["Index", "Name"] + table.align["Name"] = "l" + for index, name in enumerate(self._term_names): + table.add_row([index, name]) + return msg + table.get_string() + "\n" + + @property + def active_terms(self) -> list[str]: + """Names of active curriculum terms.""" + return self._term_names + + @property + def reset_extras(self) -> dict[str, torch.Tensor]: + """Persistent scalar views containing the latest curriculum states.""" + return self._reset_extras + + def reset( + self, + env_mask: wp.array(dtype=wp.bool), + *, + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Reset selected class terms and return persistent curriculum logging outputs. + + Args: + env_mask: Boolean Warp mask selecting environments to reset. + env_ids: Compact environment IDs for legacy class terms that explicitly require them. + + Returns: + Persistent scalar curriculum states keyed by their logging paths. + """ + env_mask = self._resolve_reset_mask(None, env_mask) + for term_cfg, mode in zip(self._term_cfgs, self._term_modes): + if isinstance(term_cfg.func, ManagerTermBase): + term_cfg.func.reset(env_mask=env_mask) + elif isinstance(term_cfg.func, StableManagerTermBase): + if mode == "legacy_ids": + if env_ids is None: + raise RuntimeError("Legacy curriculum reset requires compact env_ids.") + legacy_env_ids = env_ids + else: + legacy_env_ids = slice(None) + term_cfg.func.reset(env_ids=legacy_env_ids) + return self._reset_extras + + @property + def requires_host_ids(self) -> bool: + """Whether any active legacy term genuinely consumes compact environment IDs.""" + return "legacy_ids" in self._term_modes + + @property + def requires_host_boundary(self) -> bool: + """Whether any active legacy term must run only when a reset occurs.""" + return any(mode != "mask" for mode in self._term_modes) + + def compute( + self, + env_mask: wp.array(dtype=wp.bool), + *, + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Update curriculum terms for selected environments. + + Args: + env_mask: Boolean Warp mask selecting environments to update. + env_ids: Compact environment IDs for legacy terms that explicitly require them. + + Raises: + RuntimeError: If a legacy ID term is active and ``env_ids`` is not provided. + """ + env_mask = self._resolve_reset_mask(None, env_mask) + self._term_states_wp.zero_() + for term_idx, (term_cfg, mode) in enumerate(zip(self._term_cfgs, self._term_modes)): + if mode == "mask": + term_cfg.func(self._env, env_mask, term_cfg.out, **term_cfg.params) + continue + if mode == "legacy_ids": + if env_ids is None: + raise RuntimeError(f"Curriculum term '{self._term_names[term_idx]}' requires compact env_ids.") + legacy_env_ids = env_ids + else: + legacy_env_ids = slice(None) + state = term_cfg.func(self._env, legacy_env_ids, **term_cfg.params) + if state is not None: + if isinstance(state, torch.Tensor) and state.numel() != 1: + raise TypeError( + f"Curriculum term '{self._term_names[term_idx]}' must return a scalar state;" + f" received tensor shape {tuple(state.shape)}." + ) + if not isinstance(state, (Real, torch.Tensor)): + raise TypeError( + f"Curriculum term '{self._term_names[term_idx]}' returned {type(state).__name__}." + " Warp CurriculumManager supports scalar logging states only." + ) + self._term_states[term_idx] = state + + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + """Return curriculum states for debug inspection. + + Args: + env_idx: Unused environment index retained for the manager interface. + + Returns: + Curriculum term names and their scalar states. + """ + del env_idx + states = self._term_states.detach().cpu().tolist() + return [(name, [states[index]]) for index, name in enumerate(self._term_names)] + + def _prepare_terms(self): + cfg_items = self.cfg.items() if isinstance(self.cfg, dict) else self.cfg.__dict__.items() + for term_name, term_cfg in cfg_items: + if term_cfg is None: + continue + if not isinstance(term_cfg, StableCurriculumTermCfg): + raise TypeError( + f"Configuration for the term '{term_name}' is not of type CurriculumTermCfg." + f" Received: '{type(term_cfg)}'." + ) + if isinstance(term_cfg.func, str): + term_cfg.func = string_to_callable(term_cfg.func) + if self._is_mask_term(term_cfg.func): + mode = "mask" + self._resolve_common_term_cfg(term_name, term_cfg, min_argc=3) + else: + mode = "legacy_global" if getattr(term_cfg, "requires_host_ids", None) is False else "legacy_ids" + self._resolve_legacy_term_cfg(term_name, term_cfg) + self._term_names.append(term_name) + self._term_cfgs.append(term_cfg) + self._term_modes.append(mode) + + @staticmethod + def _is_mask_term(func) -> bool: + """Return whether a term follows ``(env, env_mask, out, ...)``.""" + func_static = func.__call__ if inspect.isclass(func) else func + parameters = list(inspect.signature(func_static).parameters) + if inspect.isclass(func): + parameters = parameters[1:] + return parameters[:3] == ["env", "env_mask", "out"] + + def _resolve_legacy_term_cfg(self, term_name: str, term_cfg: StableCurriculumTermCfg) -> None: + """Validate and initialize a stable ``(env, env_ids, ...)`` curriculum term.""" + if not callable(term_cfg.func): + raise AttributeError(f"The term '{term_name}' is not callable. Received: {term_cfg.func}") + is_class = inspect.isclass(term_cfg.func) + func_static = term_cfg.func.__call__ if is_class else term_cfg.func + min_argc = 3 if is_class else 2 + signature = inspect.signature(func_static) + parameters = list(signature.parameters.values()) + for parameter in parameters[min_argc:]: + if ( + parameter.default is not inspect.Parameter.empty + and parameter.name not in term_cfg.params + and hasattr(parameter.default, "__dataclass_fields__") + ): + term_cfg.params[parameter.name] = parameter.default.copy() + required = { + parameter.name for parameter in parameters[min_argc:] if parameter.default is inspect.Parameter.empty + } + accepted = {parameter.name for parameter in parameters[min_argc:]} + provided = set(term_cfg.params) + if not required.issubset(provided) or not provided.issubset(accepted): + raise ValueError( + f"The legacy curriculum term '{term_name}' expects parameters {sorted(accepted)}," + f" but received {sorted(provided)}." + ) + switch = getattr(self._env, "_manager_call_switch", None) + if switch is not None: + switch.register_manager_capturability(type(self).__name__, False) + if self._env.sim.is_playing(): + self._process_term_cfg_at_play(term_name, term_cfg) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py index ab8a0264ea9b..b745d54e26cd 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py @@ -243,15 +243,7 @@ def reset( *, env_mask: wp.array | torch.Tensor | None = None, ) -> dict[str, float]: - # Mask-first path: captured callers must provide env_mask. - if env_mask is None or not isinstance(env_mask, wp.array): - # Keep all id->mask resolution strictly outside capture. - if wp.get_device().is_capturing: - raise RuntimeError( - "EventManager.reset requires env_mask(wp.array[bool]) during capture. " - "Do not pass env_ids on captured paths." - ) - env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + env_mask = self._resolve_reset_mask(env_ids, env_mask) # reset class terms (mask-based) for mode_cfg in self._mode_class_term_cfgs.values(): diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py index 6caea3eebcbb..31bfceb62516 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py @@ -43,6 +43,29 @@ logger = logging.getLogger(__name__) +def _resolve_reset_mask( + owner: ManagerBase | ManagerTermBase, + env_ids: Sequence[int] | None, + env_mask: wp.array(dtype=wp.bool) | None, +) -> wp.array(dtype=wp.bool): + """Resolve reset input once while keeping ID conversion outside capture.""" + if isinstance(env_mask, wp.array): + if env_mask.dtype != wp.bool or env_mask.ndim != 1 or env_mask.shape[0] != owner.num_envs: + raise ValueError( + f"env_mask must be a Warp boolean array with shape ({owner.num_envs},), received {env_mask}." + ) + expected_device = wp.get_device(owner.device) + if env_mask.device != expected_device: + raise ValueError(f"env_mask must be on {expected_device}, received {env_mask.device}.") + return env_mask + if wp.get_device(owner.device).is_capturing: + raise RuntimeError( + f"{type(owner).__name__}.reset requires env_mask(wp.array[bool]) during capture. " + "Resolve environment IDs before entering the captured stage." + ) + return owner._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + + class ManagerTermBase(ABC): """Base class for manager terms. @@ -116,6 +139,14 @@ def reset(self, env_mask: wp.array | None = None) -> None: """ pass + def _resolve_reset_mask( + self, + env_ids: Sequence[int] | None, + env_mask: wp.array(dtype=wp.bool) | None, + ) -> wp.array(dtype=wp.bool): + """Resolve a legacy reset selection to the canonical Warp mask.""" + return _resolve_reset_mask(self, env_ids, env_mask) + def serialize(self) -> dict: """General serialization call. Includes the configuration dict.""" return {"cfg": class_to_dict(self.cfg)} @@ -226,7 +257,11 @@ def active_terms(self) -> list[str] | dict[str, list[str]]: Operations. """ - def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> dict[str, float]: + def reset( + self, + env_ids: Sequence[int] | None = None, + env_mask: wp.array(dtype=wp.bool) | None = None, + ) -> dict[str, float]: """Resets the manager and returns logging information for the current time-step. Args: @@ -238,6 +273,14 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None """ return {} + def _resolve_reset_mask( + self, + env_ids: Sequence[int] | None, + env_mask: wp.array(dtype=wp.bool) | None, + ) -> wp.array(dtype=wp.bool): + """Resolve a legacy reset selection to the canonical Warp mask.""" + return _resolve_reset_mask(self, env_ids, env_mask) + def find_terms(self, name_keys: str | Sequence[str]) -> list[str]: """Find terms in the manager based on the names. diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py index ce354be66a2a..dd25fac1b7e7 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py @@ -8,6 +8,7 @@ This module is a passthrough to :mod:`isaaclab.managers.manager_term_cfg` except for the following term configs which are overridden for Warp-first execution: +- :class:`CurriculumTermCfg` - :class:`ObservationTermCfg` - :class:`RewardTermCfg` - :class:`TerminationTermCfg` @@ -19,10 +20,23 @@ from dataclasses import MISSING from isaaclab.managers.manager_term_cfg import * # noqa: F401,F403 +from isaaclab.managers.manager_term_cfg import CurriculumTermCfg as _CurriculumTermCfg from isaaclab.managers.manager_term_cfg import ManagerTermBaseCfg as _ManagerTermBaseCfg from isaaclab.utils.configclass import configclass +@configclass +class CurriculumTermCfg(_CurriculumTermCfg): + """Configuration for a Warp-mask or legacy curriculum term.""" + + requires_host_ids: bool | None = None + """Whether a legacy term consumes compact environment IDs. + + ``None`` selects the manager default: mask-native terms do not require IDs, while legacy terms do. + Set this to ``False`` for global legacy terms that ignore their ``env_ids`` argument. + """ + + @configclass class RewardTermCfg(_ManagerTermBaseCfg): """Configuration for a reward term. diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index 74d0a9a955be..0825e917e2f2 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -382,15 +382,7 @@ def reset( *, env_mask: wp.array | None = None, ) -> dict[str, float]: - # Mask-first path: captured callers must provide env_mask. - if env_mask is None or not isinstance(env_mask, wp.array): - # Keep all id->mask resolution strictly outside capture. - if wp.get_device().is_capturing: - raise RuntimeError( - "ObservationManager.reset requires env_mask(wp.array[bool]) during capture. " - "Do not pass env_ids on captured paths." - ) - env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + env_mask = self._resolve_reset_mask(env_ids, env_mask) # call all terms that are classes for group_name, group_cfg in self._group_obs_class_term_cfgs.items(): diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py index 67dcbc055c2f..5cba3efcb8e6 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py @@ -127,6 +127,7 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): # call the base class constructor (this will parse the terms config) super().__init__(cfg, env) self._term_name_to_term_idx = {name: i for i, name in enumerate(self._term_names)} + self._device_weight_terms: set[str] = set() num_terms = len(self._term_names) self._num_terms = num_terms @@ -176,6 +177,17 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): self._term_weights_wp = wp.array( [float(term_cfg.weight) for term_cfg in self._term_cfgs], dtype=wp.float32, device=self.device ) + self._term_weight_views_wp: dict[str, wp.array] = {} + if num_terms > 0: + stride = self._term_weights_wp.strides[0] + for term_idx, term_name in enumerate(self._term_names): + self._term_weight_views_wp[term_name] = wp.array( + ptr=self._term_weights_wp.ptr + term_idx * stride, + dtype=wp.float32, + shape=(1,), + strides=(stride,), + device=self.device, + ) # persistent reset-time logging buffers (warp buffers) self._episode_sum_avg_wp = wp.zeros((num_terms,), dtype=wp.float32, device=self.device) @@ -242,14 +254,7 @@ def reset( Returns: A dictionary containing the information to log under the "Reward/{term_name}" key. """ - # Mask-first path: captured callers must provide env_mask. - if env_mask is None or not isinstance(env_mask, wp.array): - if wp.get_device().is_capturing: - raise RuntimeError( - "RewardManager.reset requires env_mask(wp.array[bool]) during capture. " - "Do not pass env_ids on captured paths." - ) - env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + env_mask = self._resolve_reset_mask(env_ids, env_mask) self._episode_sum_avg_wp.zero_() self._reset_count_wp.zero_() @@ -298,9 +303,10 @@ def compute(self, dt: float) -> torch.Tensor: device=self.device, ) # iterate over all the reward terms (Python loop; per-term math is warp) - for term_cfg in self._term_cfgs: - # skip if weight is zero (kind of a micro-optimization) - if term_cfg.weight == 0.0: + for term_name, term_cfg in zip(self._term_names, self._term_cfgs): + # Static zero-weight terms remain free. Terms whose weight is owned + # on-device must still run so a curriculum can enable them later. + if term_cfg.weight == 0.0 and term_name not in self._device_weight_terms: continue # compute term into the persistent warp buffer (raw, unweighted) # NOTE: `out` is pre-zeroed every step by `_reward_pre_compute_reset`. @@ -362,8 +368,30 @@ def get_term_cfg(self, term_name: str) -> RewardTermCfg: """ if term_name not in self._term_names: raise ValueError(f"Reward term '{term_name}' not found.") - # return the configuration - return self._term_cfgs[self._term_names.index(term_name)] + term_idx = self._term_names.index(term_name) + term_cfg = self._term_cfgs[term_idx] + if term_name in self._device_weight_terms: + # Explicit config inspection is a control-plane operation, so it is + # the appropriate place to synchronize a device-owned scalar. + term_cfg.weight = float(self._term_weights_tensor_view[term_idx].item()) + return term_cfg + + def get_term_weight_wp(self, term_name: str) -> wp.array(dtype=wp.float32): + """Return the pointer-stable Warp view of a reward term's scalar weight. + + Args: + term_name: Name of the reward term. + + Returns: + One-element Warp array containing the reward weight. + + Raises: + ValueError: If the term name is not found. + """ + if term_name not in self._term_weight_views_wp: + raise ValueError(f"Reward term '{term_name}' not found.") + self._device_weight_terms.add(term_name) + return self._term_weight_views_wp[term_name] def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: """Returns the active terms as iterable sequence of tuples. diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py index 8a768651b292..cc58dbb296c2 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py @@ -22,6 +22,8 @@ import warp as wp from prettytable import PrettyTable +from isaaclab_experimental.utils.warp.kernels import compute_reset_scale, count_masked + from .manager_base import ManagerBase, ManagerTermBase from .manager_term_cfg import TerminationTermCfg @@ -92,15 +94,16 @@ def _termination_finalize( # TODO(jichuanh): Look into wp.tile for better performance @wp.kernel -def _termination_reset_mean_all_2d( +def _termination_reset_mean_masked_2d( + env_mask: wp.array(dtype=wp.bool), last_episode_dones: wp.array(dtype=wp.bool, ndim=2), + scale: wp.array(dtype=wp.float32), term_done_avg: wp.array(dtype=wp.float32), ): - """Compute mean(done) per term with 2D parallel accumulation.""" + """Compute mean(done) per term over selected environments.""" env_id, term_idx = wp.tid() - num_envs = last_episode_dones.shape[0] - if num_envs > 0 and last_episode_dones[env_id, term_idx]: - wp.atomic_add(term_done_avg, term_idx, 1.0 / float(num_envs)) + if env_mask[env_id] and last_episode_dones[env_id, term_idx]: + wp.atomic_add(term_done_avg, term_idx, scale[0]) class TerminationManager(ManagerBase): @@ -129,6 +132,8 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): num_terms = len(self._term_names) self._term_dones_wp = wp.zeros((self.num_envs, num_terms), dtype=wp.bool, device=self.device) self._term_done_avg_wp = wp.zeros((num_terms,), dtype=wp.float32, device=self.device) + self._reset_count_wp = wp.zeros((1,), dtype=wp.int32, device=self.device) + self._reset_scale_wp = wp.zeros((1,), dtype=wp.float32, device=self.device) self._last_episode_dones_wp = wp.zeros((self.num_envs, num_terms), dtype=wp.bool, device=self.device) self._truncated_wp = wp.zeros((self.num_envs,), dtype=wp.bool, device=self.device) self._terminated_wp = wp.zeros((self.num_envs,), dtype=wp.bool, device=self.device) @@ -249,20 +254,27 @@ def reset( Returns: A dictionary containing the information to log under the "Termination/{term_name}" key. """ - # Mask-first path: captured callers must provide env_mask. - if env_mask is None or not isinstance(env_mask, wp.array): - if wp.get_device().is_capturing: - raise RuntimeError( - "TerminationManager.reset requires env_mask(wp.array[bool]) during capture. " - "Do not pass env_ids on captured paths." - ) - env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + env_mask = self._resolve_reset_mask(env_ids, env_mask) if len(self._term_names) > 0: self._term_done_avg_wp.zero_() + self._reset_count_wp.zero_() + self._reset_scale_wp.zero_() + wp.launch( + kernel=count_masked, + dim=self.num_envs, + inputs=[env_mask, self._reset_count_wp], + device=self.device, + ) + wp.launch( + kernel=compute_reset_scale, + dim=1, + inputs=[self._reset_count_wp, 1.0, self._reset_scale_wp], + device=self.device, + ) wp.launch( - kernel=_termination_reset_mean_all_2d, + kernel=_termination_reset_mean_masked_2d, dim=(self.num_envs, len(self._term_names)), - inputs=[self._last_episode_dones_wp, self._term_done_avg_wp], + inputs=[env_mask, self._last_episode_dones_wp, self._reset_scale_wp, self._term_done_avg_wp], device=self.device, ) for term_cfg in self._class_term_cfgs: diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py b/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py index 89924c789bc3..5296a96a6bb7 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py @@ -10,6 +10,7 @@ import importlib import json import os +from collections.abc import Callable from enum import IntEnum from typing import Any @@ -22,6 +23,7 @@ class ManagerCallMode(IntEnum): """Execution mode for manager stage calls. * ``STABLE`` (0): Call stable Python manager implementations from :mod:`isaaclab.managers`. + This mode is deprecated inside ``ManagerBasedEnvWarp``; use the stable environment instead. * ``WARP_NOT_CAPTURED`` (1): Call Warp-compatible implementations without CUDA graph capture. * ``WARP_CAPTURED`` (2): Call Warp implementations with CUDA graph capture/replay. """ @@ -39,7 +41,9 @@ class ManagerCallSwitch: wraps each call in a :class:`Timer` context for profiling. """ - DEFAULT_CONFIG: dict[str, int] = {"default": 2} + # Warp eager is the correctness-first default. Capture remains an explicit + # optimization while stage state and pointer contracts are validated. + DEFAULT_CONFIG: dict[str, int] = {"default": 1} DEFAULT_KEY = "default" MANAGER_NAMES: tuple[str, ...] = ( "ActionManager", @@ -52,9 +56,8 @@ class ManagerCallSwitch: "CurriculumManager", "Scene", ) - # FIXME: Scene_write_data_to_sim calls articulation._apply_actuator_model which - # uses wp.to_torch + torch indexing -- not capture-safe on this branch. - # Cap Scene stages to WARP_NOT_CAPTURED until the articulation layer is capture-ready. + # Scene stages remain eager until scene, sensor, and actuator graphability is + # validated together. Warp-first execution does not depend on that later step. MAX_MODE_OVERRIDES: dict[str, int] = {"Scene": ManagerCallMode.WARP_NOT_CAPTURED} ENV_VAR = "MANAGER_CALL_CONFIG" @@ -62,7 +65,7 @@ class ManagerCallSwitch: Example usage:: - MANAGER_CALL_CONFIG='{"RewardManager": 0, "default": 2}' python train.py ... + MANAGER_CALL_CONFIG='{"RewardManager": 0, "default": 1}' python train.py ... """ def __init__( @@ -71,7 +74,8 @@ def __init__( *, max_modes: dict[str, int] | None = None, ): - self._graph_cache = WarpGraphCache() + # The graph cache is reached only for an explicit WARP_CAPTURED mode. + self._graph_cache = WarpGraphCache(enabled=True) # Merge caller-supplied max_modes with the class-level MAX_MODE_OVERRIDES. self._max_modes = dict(self.MAX_MODE_OVERRIDES) if max_modes is not None: @@ -100,6 +104,40 @@ def invalidate_graphs(self) -> None: # Stage dispatch # ------------------------------------------------------------------ + def call( + self, + stage: str, + fn: Callable[..., Any], + /, + *args: Any, + _output: Callable[[Any], Any] | None = None, + _timer: bool = False, + **kwargs: Any, + ) -> Any: + """Run a Warp frontend stage eagerly or through its cached CUDA graph. + + The call site always supplies the same callable and arguments. The + configured manager mode only controls how that call is executed. + + Args: + stage: Stage identifier in the form ``"ManagerName_function_name"``. + fn: Callable implementing the stage. + *args: Positional arguments forwarded to :paramref:`fn`. + _output: Optional transform applied to the stage result after execution. + _timer: Whether to time the stage. + **kwargs: Keyword arguments forwarded to :paramref:`fn`. + + Returns: + The stage result, optionally transformed by :paramref:`_output`. + """ + with Timer(name=stage, msg=f"{stage} took:", enable=_timer, time_unit="us"): + mode = self.get_mode_for_manager(self._manager_name_from_stage(stage)) + if mode == ManagerCallMode.WARP_CAPTURED: + result = self._graph_cache.capture_or_replay(stage, fn, args=args, kwargs=kwargs) + else: + result = fn(*args, **kwargs) + return _output(result) if _output is not None else result + def call_stage( self, *, diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py index 2d071b823a46..27e117d43478 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py @@ -5,11 +5,19 @@ """Warp utility functions and shared kernels for isaaclab_experimental.""" -from .kernels import compute_reset_scale, count_masked +from isaaclab.utils.warp.utils import resolve_1d_mask + +from .kernels import ( + compute_reset_scale, + count_masked, + increment_all_int32, + increment_all_int64, + zero_masked_int32, + zero_masked_int64, +) from .utils import ( WarpCapturable, is_warp_capturable, - resolve_1d_mask, warp_capturable, wrap_to_pi, zero_masked_2d, diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py index 8d3f9e65d49a..dbfd3740a61b 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py @@ -10,6 +10,36 @@ import warp as wp +@wp.kernel +def increment_all_int32(values: wp.array(dtype=wp.int32), increment: wp.int32): + """Increment every value in a one-dimensional int32 array.""" + index = wp.tid() + values[index] = values[index] + increment + + +@wp.kernel +def increment_all_int64(values: wp.array(dtype=wp.int64), increment: wp.int64): + """Increment every value in a one-dimensional int64 array.""" + index = wp.tid() + values[index] = values[index] + increment + + +@wp.kernel +def zero_masked_int32(mask: wp.array(dtype=wp.bool), values: wp.array(dtype=wp.int32)): + """Zero selected values in a one-dimensional int32 array.""" + index = wp.tid() + if mask[index]: + values[index] = 0 + + +@wp.kernel +def zero_masked_int64(mask: wp.array(dtype=wp.bool), values: wp.array(dtype=wp.int64)): + """Zero selected values in a one-dimensional int64 array.""" + index = wp.tid() + if mask[index]: + values[index] = wp.int64(0) + + @wp.kernel def count_masked( mask: wp.array(dtype=wp.bool), diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py index a90f6b2d6fc6..6b8aafc415f8 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py @@ -5,97 +5,9 @@ from __future__ import annotations -from collections.abc import Sequence - -import torch import warp as wp -@wp.kernel -def _set_mask_from_ids( - mask: wp.array(dtype=wp.bool), - ids: wp.array(dtype=wp.int32), -): - """Set ``mask[ids[i]] = True`` for each thread *i*.""" - i = wp.tid() - mask[ids[i]] = True - - -def resolve_1d_mask( - *, - ids: Sequence[int] | slice | wp.array | torch.Tensor | None, - mask: wp.array | torch.Tensor | None, - all_mask: wp.array, - scratch_mask: wp.array, - device: str, -) -> wp.array: - """Resolve ids/mask into a warp boolean mask. - - Matches the contract of ``ArticulationData._resolve_1d_mask`` on dev/newton. - Callers must provide pre-allocated ``all_mask`` (all-True) and ``scratch_mask`` - (reusable working buffer). No allocations happen inside this function. - - Args: - ids: Indices to set to ``True``. ``None`` or ``slice(None)`` means all. - mask: Explicit boolean mask. If provided, returned directly (after - torch->warp normalization if needed). Takes precedence over *ids*. - all_mask: Pre-allocated all-True mask of shape ``(size,)``, returned - when both *ids* and *mask* are ``None``. - scratch_mask: Pre-allocated scratch mask of shape ``(size,)``, filled - in-place when *ids* are provided. - device: Warp device string. - - Returns: - A ``wp.array(dtype=wp.bool)`` -- ``mask``, ``all_mask``, or ``scratch_mask``. - """ - # Fast path: explicit mask provided. - if mask is not None: - if isinstance(mask, torch.Tensor): - if mask.dtype != torch.bool: - mask = mask.to(dtype=torch.bool) - if str(mask.device) != device: - mask = mask.to(device) - return wp.from_torch(mask, dtype=wp.bool) - return mask - - # Fast path: all ids. - if ids is None or (isinstance(ids, slice) and ids == slice(None)): - return all_mask - - # Normalize slice into explicit indices. - if isinstance(ids, slice): - start, stop, step = ids.indices(scratch_mask.shape[0]) - ids = list(range(start, stop, step)) - elif not isinstance(ids, (torch.Tensor, wp.array)): - ids = list(ids) - - # Prepare output mask. - scratch_mask.fill_(False) - - # Normalize ids to wp.int32 array and launch kernel. - if isinstance(ids, torch.Tensor): - if ids.numel() == 0: - return scratch_mask - if str(ids.device) != device: - ids = ids.to(device) - if ids.dtype != torch.int32: - ids = ids.to(dtype=torch.int32) - if not ids.is_contiguous(): - ids = ids.contiguous() - ids_wp = wp.from_torch(ids, dtype=wp.int32) - elif isinstance(ids, wp.array): - if ids.shape[0] == 0: - return scratch_mask - ids_wp = ids - else: - if len(ids) == 0: - return scratch_mask - ids_wp = wp.array(ids, dtype=wp.int32, device=device) - - wp.launch(kernel=_set_mask_from_ids, dim=ids_wp.shape[0], inputs=[scratch_mask, ids_wp], device=device) - return scratch_mask - - def warp_capturable(capturable: bool): """Annotate an MDP term's CUDA-graph capturability. diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py index d446be28d099..047e0a3fd0e5 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py @@ -12,7 +12,7 @@ class WarpGraphCache: - """Caches Warp CUDA graphs by stage name: captures on first call, replays after. + """Run Warp stages eagerly or cache CUDA graphs by stage name. On the very first call for a given stage, an **eager warm-up** run executes *before* graph capture. This lets one-time initialisation @@ -32,7 +32,13 @@ class WarpGraphCache: result2 = cache.capture_or_replay("my_stage_post", my_other_function) """ - def __init__(self): + def __init__(self, *, enabled: bool = True): + """Initialize the execution cache. + + Args: + enabled: Whether to capture stages. When false, stages execute eagerly. + """ + self._enabled = enabled self._graphs: dict[str, Any] = {} self._results: dict[str, Any] = {} @@ -54,10 +60,12 @@ def capture_or_replay( kwargs: Keyword arguments forwarded to *fn*. Defaults to ``None``. Returns: - The cached return value from the first (capture) invocation. + The eager result when disabled, otherwise the cached result from capture. """ if kwargs is None: kwargs = {} + if not self._enabled: + return fn(*args, **kwargs) graph = self._graphs.get(stage) if graph is not None: wp.capture_launch(graph) diff --git a/source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py b/source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py new file mode 100644 index 000000000000..f6dca451673e --- /dev/null +++ b/source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py @@ -0,0 +1,302 @@ +# 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 + +"""Tests for Warp-native registered command terms.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import torch +import warp as wp +from isaaclab_experimental.envs import mdp as experimental_mdp +from isaaclab_experimental.envs.mdp.commands import ( + UniformPoseCommand, + UniformPoseCommandCfg, + UniformVelocityCommand, + UniformVelocityCommandCfg, +) +from isaaclab_experimental.managers import CommandManager, CommandTerm + +from isaaclab.envs.mdp.commands import UniformPoseCommand as StableUniformPoseCommand +from isaaclab.envs.mdp.commands import UniformPoseCommandCfg as StableUniformPoseCommandCfg +from isaaclab.envs.mdp.commands import UniformVelocityCommand as StableUniformVelocityCommand +from isaaclab.envs.mdp.commands import UniformVelocityCommandCfg as StableUniformVelocityCommandCfg +from isaaclab.utils.math import combine_frame_transforms, compute_pose_error, quat_from_euler_xyz, quat_unique + + +class _ArrayProxy: + """Minimal backend-neutral data proxy used by command terms.""" + + def __init__(self, values: np.ndarray, dtype): + self.warp = wp.array(values, dtype=dtype, device="cpu") + self.torch = wp.to_torch(self.warp) + + +class _Robot: + """Minimal articulation data needed by velocity and pose commands.""" + + def __init__(self, num_envs: int): + identity_quat = np.zeros((num_envs, 4), dtype=np.float32) + identity_quat[:, 3] = 1.0 + self.data = SimpleNamespace( + root_pos_w=_ArrayProxy(np.zeros((num_envs, 3), dtype=np.float32), wp.vec3f), + root_quat_w=_ArrayProxy(identity_quat, wp.quatf), + root_lin_vel_b=_ArrayProxy(np.zeros((num_envs, 3), dtype=np.float32), wp.vec3f), + root_ang_vel_b=_ArrayProxy(np.zeros((num_envs, 3), dtype=np.float32), wp.vec3f), + heading_w=_ArrayProxy(np.zeros(num_envs, dtype=np.float32), wp.float32), + body_pos_w=_ArrayProxy(np.zeros((num_envs, 1, 3), dtype=np.float32), wp.vec3f), + body_quat_w=_ArrayProxy(identity_quat[:, None, :], wp.quatf), + ) + self.is_initialized = True + + def find_bodies(self, body_name: str) -> tuple[list[int], list[str]]: + return [0], [body_name] + + +class _Env: + """Minimal manager-based environment used by command terms.""" + + def __init__(self, num_envs: int = 4): + self.num_envs = num_envs + self.device = "cpu" + self.scene = {"robot": _Robot(num_envs)} + self.rng_state_wp = wp.array(np.arange(num_envs, dtype=np.uint32) + 41, device=self.device) + self.extras = {} + self.sim = SimpleNamespace( + vis_marker_registry=SimpleNamespace( + add_debug_vis_callback=lambda term: None, + clear_debug_vis_callback=lambda term: None, + ) + ) + + +def _velocity_cfg() -> UniformVelocityCommandCfg: + return UniformVelocityCommandCfg( + asset_name="robot", + resampling_time_range=(2.0, 2.0), + heading_command=True, + heading_control_stiffness=2.0, + rel_heading_envs=1.0, + rel_standing_envs=0.0, + vel_xy_success_threshold=0.5, + vel_yaw_success_threshold=0.4, + ranges=UniformVelocityCommandCfg.Ranges( + lin_vel_x=(1.0, 1.0), + lin_vel_y=(2.0, 2.0), + ang_vel_z=(0.3, 0.3), + heading=(0.7, 0.7), + ), + debug_vis=False, + ) + + +def _pose_cfg() -> UniformPoseCommandCfg: + return UniformPoseCommandCfg( + asset_name="robot", + body_name="tool", + resampling_time_range=(3.0, 3.0), + make_quat_unique=True, + position_success_threshold=0.25, + ranges=UniformPoseCommandCfg.Ranges( + pos_x=(1.0, 1.0), + pos_y=(2.0, 2.0), + pos_z=(3.0, 3.0), + roll=(0.2, 0.2), + pitch=(-0.3, -0.3), + yaw=(4.0, 4.0), + ), + debug_vis=False, + ) + + +def test_configs_exports_and_command_accessors_preserve_shared_contracts(): + """Configs should inherit stable fields and terms should expose shared debug UI and persistent storage.""" + assert issubclass(UniformVelocityCommandCfg, StableUniformVelocityCommandCfg) + assert issubclass(UniformPoseCommandCfg, StableUniformPoseCommandCfg) + assert experimental_mdp.UniformVelocityCommandCfg is UniformVelocityCommandCfg + assert experimental_mdp.UniformPoseCommandCfg is UniformPoseCommandCfg + assert str(_velocity_cfg().class_type).endswith(".commands.velocity_command:UniformVelocityCommand") + assert str(_pose_cfg().class_type).endswith(".commands.pose_command:UniformPoseCommand") + assert UniformVelocityCommand._set_debug_vis_impl is StableUniformVelocityCommand._set_debug_vis_impl + assert UniformPoseCommand._set_debug_vis_impl is StableUniformPoseCommand._set_debug_vis_impl + + env = _Env() + velocity_term = UniformVelocityCommand(_velocity_cfg(), env) + pose_term = UniformPoseCommand(_pose_cfg(), env) + stable_pose_term = StableUniformPoseCommand(_pose_cfg(), env) + manager = CommandManager.__new__(CommandManager) + manager._terms = {"velocity": velocity_term, "pose": pose_term} + + assert velocity_term.command_wp is velocity_term.command_wp + assert pose_term.command_wp is pose_term.command_wp + assert manager.get_command_wp("velocity") is velocity_term.command_wp + assert manager.get_command_wp("pose") is pose_term.command_wp + assert manager.get_command("velocity") is velocity_term.command + assert manager.get_command("pose") is pose_term.command + torch.testing.assert_close(pose_term.command[:, :6], torch.zeros(4, 6)) + torch.testing.assert_close(pose_term.command[:, 6], torch.ones(4)) + torch.testing.assert_close(stable_pose_term.command[:, :6], torch.zeros(4, 6)) + torch.testing.assert_close(stable_pose_term.command[:, 6], torch.ones(4)) + + +def test_debug_visualization_uses_current_simulation_registry(): + """Command terms should register visualization through the current simulation API.""" + registry = SimpleNamespace(add_debug_vis_callback=Mock(return_value="handle"), clear_debug_vis_callback=Mock()) + term = SimpleNamespace( + has_debug_vis_implementation=True, + _set_debug_vis_impl=Mock(), + _debug_vis_handle=None, + _env=SimpleNamespace(sim=SimpleNamespace(vis_marker_registry=registry)), + ) + + assert CommandTerm.set_debug_vis(term, True) + assert term._debug_vis_handle == "handle" + registry.add_debug_vis_callback.assert_called_once_with(term) + + assert CommandTerm.set_debug_vis(term, False) + registry.clear_debug_vis_callback.assert_called_once_with(term) + + CommandTerm.__del__(term) + assert registry.clear_debug_vis_callback.call_count == 2 + + +def test_uniform_velocity_command_updates_and_resets_selected_rows(): + """Velocity commands should update in Warp and preserve unselected rows during sparse reset.""" + env = _Env() + term = UniformVelocityCommand(_velocity_cfg(), env) + term._prepare_reset_extras() + command_pointer = term.command.data_ptr() + command_wp = term.command_wp + + term.command[:] = torch.tensor([[1.0, 2.0, 0.3], [2.0, -1.0, -0.2], [0.0, 0.5, 0.1], [-1.0, -2.0, 0.0]]) + env.scene["robot"].data.root_lin_vel_b.torch[:] = torch.tensor( + [[0.0, 0.0, 0.0], [1.0, -1.0, 0.0], [0.0, 0.0, 0.0], [-2.0, -2.0, 0.0]] + ) + env.scene["robot"].data.root_ang_vel_b.torch[:, 2] = torch.tensor([0.1, -0.1, 0.4, 0.0]) + term._update_metrics() + wp.synchronize() + torch.testing.assert_close( + wp.to_torch(term._error_xy_sum_wp), + torch.tensor([np.sqrt(5.0), 1.0, 0.5, 1.0], dtype=torch.float32), + ) + torch.testing.assert_close(wp.to_torch(term._error_yaw_sum_wp), torch.tensor([0.2, 0.1, 0.3, 0.0])) + + term.command.fill_(-9.0) + wp.to_torch(term._error_xy_sum_wp)[:] = torch.tensor([0.2, 9.0, 2.0, 8.0]) + wp.to_torch(term._error_yaw_sum_wp)[:] = torch.tensor([0.1, 9.0, 2.0, 8.0]) + wp.to_torch(term._step_count_wp)[:] = torch.tensor([1.0, 1.0, 2.0, 1.0]) + wp.to_torch(term.time_left_wp)[:] = torch.tensor([-1.0, 6.0, -1.0, 8.0]) + wp.to_torch(term.command_counter_wp)[:] = torch.tensor([5, 6, 7, 8], dtype=torch.int32) + env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + + extras = term.reset(env_mask=env_mask) + wp.synchronize() + + assert term.command_wp is command_wp + assert term.command.data_ptr() == command_pointer + torch.testing.assert_close(term.command[[0, 2]], torch.tensor([[1.0, 2.0, 0.3], [1.0, 2.0, 0.3]])) + torch.testing.assert_close(term.command[[1, 3]], torch.full((2, 3), -9.0)) + torch.testing.assert_close(wp.to_torch(term.time_left_wp), torch.tensor([2.0, 6.0, 2.0, 8.0])) + torch.testing.assert_close(wp.to_torch(term.command_counter_wp), torch.tensor([1, 6, 1, 8], dtype=torch.int32)) + torch.testing.assert_close(extras["error_vel_xy"], torch.tensor(0.6)) + torch.testing.assert_close(extras["error_vel_yaw"], torch.tensor(0.55)) + torch.testing.assert_close(extras["success_rate"], torch.tensor(0.5)) + torch.testing.assert_close(wp.to_torch(term._error_xy_sum_wp), torch.tensor([0.0, 9.0, 0.0, 8.0])) + torch.testing.assert_close(wp.to_torch(term._step_count_wp), torch.tensor([0.0, 1.0, 0.0, 1.0])) + + term.command[:] = torch.tensor([[1.0, 2.0, 0.4]]).repeat(4, 1) + term.heading_target[:] = torch.tensor([np.pi / 2, 0.0, np.pi / 2, 0.0]) + term.is_heading_env[:] = torch.tensor([True, False, True, False]) + term.is_standing_env[:] = torch.tensor([False, False, True, True]) + env.scene["robot"].data.heading_w.torch.zero_() + term._update_command() + wp.synchronize() + torch.testing.assert_close( + term.command, + torch.tensor([[1.0, 2.0, 0.3], [1.0, 2.0, 0.4], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]), + ) + + +def test_uniform_pose_command_matches_stable_math_and_tracks_sticky_success(): + """Pose commands should use xyzw math, update metrics, and reset sticky success by mask.""" + env = _Env() + term = UniformPoseCommand(_pose_cfg(), env) + term._prepare_reset_extras() + command_pointer = term.command.data_ptr() + command_wp = term.command_wp + term.command.fill_(-9.0) + wp.to_torch(term.time_left_wp)[:] = torch.tensor([-1.0, 6.0, -1.0, 8.0]) + wp.to_torch(term.command_counter_wp)[:] = torch.tensor([5, 6, 7, 8], dtype=torch.int32) + env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + + term.reset(env_mask=env_mask) + wp.synchronize() + + expected_quat = quat_unique(quat_from_euler_xyz(torch.tensor([0.2]), torch.tensor([-0.3]), torch.tensor([4.0])))[0] + assert term.command_wp is command_wp + assert term.command.data_ptr() == command_pointer + torch.testing.assert_close(term.command[[0, 2], :3], torch.tensor([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]])) + torch.testing.assert_close(term.command[[0, 2], 3:], expected_quat.repeat(2, 1), atol=1.0e-6, rtol=1.0e-6) + torch.testing.assert_close(term.command[[1, 3]], torch.full((2, 7), -9.0)) + torch.testing.assert_close(wp.to_torch(term.time_left_wp), torch.tensor([3.0, 6.0, 3.0, 8.0])) + torch.testing.assert_close(wp.to_torch(term.command_counter_wp), torch.tensor([1, 6, 1, 8], dtype=torch.int32)) + + term.command[[1, 3]] = torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]) + + roll = torch.tensor([0.0, 0.1, -0.2, 0.3]) + pitch = torch.tensor([0.0, -0.1, 0.2, -0.3]) + yaw = torch.tensor([0.5, -0.4, 0.3, -0.2]) + root_quat = quat_from_euler_xyz(roll, pitch, yaw) + env.scene["robot"].data.root_pos_w.torch[:] = torch.tensor( + [[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [-5.0, 1.0, 2.0], [1.0, 2.0, 3.0]] + ) + env.scene["robot"].data.root_quat_w.torch[:] = root_quat + expected_position_w, expected_orientation_w = combine_frame_transforms( + env.scene["robot"].data.root_pos_w.torch, + root_quat, + term.command[:, :3], + term.command[:, 3:], + ) + offsets = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.1, 0.0, 0.0], [0.5, 0.0, 0.0]]) + env.scene["robot"].data.body_pos_w.torch[:, 0] = expected_position_w + offsets + identity_quat = torch.zeros(4, 4) + identity_quat[:, 3] = 1.0 + env.scene["robot"].data.body_quat_w.torch[:, 0] = identity_quat + + term._update_metrics() + wp.synchronize() + + expected_position_delta, expected_orientation_delta = compute_pose_error( + expected_position_w, + expected_orientation_w, + env.scene["robot"].data.body_pos_w.torch[:, 0], + env.scene["robot"].data.body_quat_w.torch[:, 0], + ) + torch.testing.assert_close(term.pose_command_w[:, :3], expected_position_w, atol=1.0e-5, rtol=1.0e-5) + torch.testing.assert_close(term.pose_command_w[:, 3:], expected_orientation_w, atol=1.0e-5, rtol=1.0e-5) + torch.testing.assert_close( + wp.to_torch(term.metrics["position_error"]), torch.linalg.norm(expected_position_delta, dim=-1) + ) + torch.testing.assert_close( + wp.to_torch(term.metrics["orientation_error"]), + torch.linalg.norm(expected_orientation_delta, dim=-1), + atol=1.0e-5, + rtol=1.0e-5, + ) + torch.testing.assert_close(wp.to_torch(term._success_rate_wp), torch.tensor([1.0, 0.0, 1.0, 0.0])) + + env.scene["robot"].data.body_pos_w.torch[0, 0] += 2.0 + term._update_metrics() + wp.synchronize() + torch.testing.assert_close(wp.to_torch(term._success_rate_wp), torch.tensor([1.0, 0.0, 1.0, 0.0])) + + reset_mask = wp.array([True, False, False, True], dtype=wp.bool, device="cpu") + extras = term.reset(env_mask=reset_mask) + wp.synchronize() + torch.testing.assert_close(extras["success_rate"], torch.tensor(0.5)) + torch.testing.assert_close(wp.to_torch(term._success_rate_wp), torch.tensor([0.0, 0.0, 1.0, 0.0])) diff --git a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py index 4ceec2ecb60b..a549592a9ada 100644 --- a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py +++ b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py @@ -533,11 +533,15 @@ class MockCommandManager: def __init__(self, command_tensor: torch.Tensor, cmd_term: MockCommandTerm): self._cmd = command_tensor + self._cmd_wp = wp.from_torch(command_tensor, dtype=wp.float32) self._term = cmd_term def get_command(self, name: str) -> torch.Tensor: return self._cmd + def get_command_wp(self, name: str) -> wp.array(dtype=wp.float32, ndim=2): + return self._cmd_wp + def get_term(self, name: str): return self._term diff --git a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py index 24b9eb022665..ecdfb67549c4 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py @@ -17,6 +17,7 @@ pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") import isaaclab_experimental.envs.mdp.events as warp_evt +from isaaclab_experimental.managers import SceneEntityCfg from parity_helpers import ( DEVICE, NUM_ACTIONS, @@ -37,25 +38,6 @@ # ============================================================================ -@pytest.fixture(autouse=True) -def _clear_function_caches(): - """Clear first-call caches on warp MDP functions so each test starts fresh. - - Functions that cache warp views via the ``hasattr`` pattern need clearing - between tests to avoid stale references from prior fixtures. - """ - yield - for fn in ( - warp_evt.push_by_setting_velocity, - warp_evt.apply_external_force_torque, - warp_evt.reset_root_state_uniform, - warp_evt.randomize_rigid_body_com, - ): - for attr in list(vars(fn)): - if attr.startswith("_"): - delattr(fn, attr) - - @pytest.fixture() def art_data(): return MockArticulationData(NUM_ENVS, NUM_JOINTS, DEVICE) @@ -130,6 +112,24 @@ def all_joints_cfg(): return MockSceneEntityCfg("robot", list(range(NUM_JOINTS)), NUM_JOINTS, DEVICE) +def _make_event_env(seed: int, *, num_bodies: int = 1): + """Create an independent Warp event environment for ownership tests.""" + + class _Env: + pass + + data = MockArticulationData(NUM_ENVS, NUM_JOINTS, DEVICE, seed=seed, num_bodies=num_bodies) + asset = MockArticulation(data, num_bodies=num_bodies) + origins = wp.zeros(NUM_ENVS, dtype=wp.vec3f, device=DEVICE) + env = _Env() + env.scene = MockScene({"robot": asset}, origins) + env.num_envs = NUM_ENVS + env.device = DEVICE + env.env_origins_wp = origins + env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + seed, device=DEVICE) + return env, data, asset + + # ============================================================================ # Event parity tests: deterministic (zero-width range) warp vs stable # ============================================================================ @@ -268,6 +268,8 @@ def test_reset_joints_by_scale(self, warp_env, art_data, all_joints_cfg): def test_push_by_setting_velocity(self, warp_env, art_data, all_joints_cfg): """With zero-width velocity range, scratch == root_vel_w. Mutate root_vel_w -> scratch tracks.""" mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + captured = {} + warp_env.scene["robot"].write_root_velocity_to_sim_mask = lambda **kwargs: captured.update(kwargs) zero_range = { "x": (0.0, 0.0), "y": (0.0, 0.0), @@ -288,7 +290,7 @@ def test_push_by_setting_velocity(self, warp_env, art_data, all_joints_cfg): wp.capture_launch(cap.graph) wp.synchronize() - scratch = wp.to_torch(warp_evt.push_by_setting_velocity._scratch_vel) + scratch = wp.to_torch(captured["root_velocity"]) expected = torch.tensor([1.0, 2.0, 3.0, 0.1, 0.2, 0.3], device=DEVICE).expand(NUM_ENVS, -1) assert_close(scratch, expected) @@ -297,16 +299,21 @@ def test_push_by_setting_velocity(self, warp_env, art_data, all_joints_cfg): def test_apply_external_force_torque(self, warp_env, art_data, all_joints_cfg): """With zero-width ranges, forces/torques are zero. Non-zero ranges produce non-zero output.""" mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + captured = {} + warp_env.scene["robot"].permanent_wrench_composer.set_forces_and_torques_mask = ( + lambda **kwargs: captured.update(kwargs) + ) + zero_range = (0.0, 0.0) # Zero-range: forces and torques should be zero - warp_evt.apply_external_force_torque(warp_env, mask, force_range=(0.0, 0.0), torque_range=(0.0, 0.0)) + warp_evt.apply_external_force_torque(warp_env, mask, force_range=zero_range, torque_range=zero_range) with wp.ScopedCapture() as cap: - warp_evt.apply_external_force_torque(warp_env, mask, force_range=(0.0, 0.0), torque_range=(0.0, 0.0)) + warp_evt.apply_external_force_torque(warp_env, mask, force_range=zero_range, torque_range=zero_range) wp.capture_launch(cap.graph) wp.synchronize() - forces = wp.to_torch(warp_evt.apply_external_force_torque._scratch_forces) - torques = wp.to_torch(warp_evt.apply_external_force_torque._scratch_torques) + forces = wp.to_torch(captured["forces"]) + torques = wp.to_torch(captured["torques"]) assert_close(forces, torch.zeros_like(forces)) assert_close(torques, torch.zeros_like(torques)) @@ -338,3 +345,218 @@ def test_reset_joints_mask_selectivity(self, warp_env, art_data, all_joints_cfg) assert_close(result[: NUM_ENVS // 2], torch.zeros(NUM_ENVS // 2, NUM_JOINTS, device=DEVICE)) # Unmasked envs: still 999.0 assert_close(result[NUM_ENVS // 2 :], torch.full((NUM_ENVS // 2, NUM_JOINTS), 999.0, device=DEVICE)) + + +class TestEventStateOwnership: + """Verify event scratch and parsed ranges are owned by one environment configuration.""" + + def test_push_state_is_not_shared_between_environments(self): + env_a, data_a, asset_a = _make_event_env(101) + env_b, data_b, asset_b = _make_event_env(202) + copy_np_to_wp(data_a.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + copy_np_to_wp(data_b.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + captured_a = {} + captured_b = {} + asset_a.write_root_velocity_to_sim_mask = lambda **kwargs: captured_a.update(kwargs) + asset_b.write_root_velocity_to_sim_mask = lambda **kwargs: captured_b.update(kwargs) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + range_a = {"x": (1.0, 1.0)} + range_b = {"x": (2.0, 2.0)} + + warp_evt.push_by_setting_velocity(env_a, env_mask, velocity_range=range_a) + warp_evt.push_by_setting_velocity(env_b, env_mask, velocity_range=range_b) + wp.synchronize() + + velocity_a = captured_a["root_velocity"] + velocity_b = captured_b["root_velocity"] + assert velocity_a.ptr != velocity_b.ptr + assert_close(wp.to_torch(velocity_a)[:, 0], torch.ones(NUM_ENVS, device=DEVICE)) + assert_close(wp.to_torch(velocity_b)[:, 0], torch.full((NUM_ENVS,), 2.0, device=DEVICE)) + + def test_push_state_is_not_shared_between_configurations(self): + env, data, asset = _make_event_env(252) + copy_np_to_wp(data.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + captured = [] + asset.write_root_velocity_to_sim_mask = lambda **kwargs: captured.append(kwargs) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + range_a = {"x": (1.0, 1.0)} + range_b = {"x": (2.0, 2.0)} + + warp_evt.push_by_setting_velocity(env, env_mask, velocity_range=range_a) + warp_evt.push_by_setting_velocity(env, env_mask, velocity_range=range_b) + wp.synchronize() + + velocity_a = captured[0]["root_velocity"] + velocity_b = captured[1]["root_velocity"] + assert velocity_a.ptr != velocity_b.ptr + assert_close(wp.to_torch(velocity_a)[:, 0], torch.ones(NUM_ENVS, device=DEVICE)) + assert_close(wp.to_torch(velocity_b)[:, 0], torch.full((NUM_ENVS,), 2.0, device=DEVICE)) + + def test_push_state_updates_after_in_place_range_mutation(self): + env, data, asset = _make_event_env(277) + copy_np_to_wp(data.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + captured = {} + asset.write_root_velocity_to_sim_mask = lambda **kwargs: captured.update(kwargs) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + velocity_range = {"x": (1.0, 1.0)} + + warp_evt.push_by_setting_velocity(env, env_mask, velocity_range=velocity_range) + wp.synchronize() + assert_close(wp.to_torch(captured["root_velocity"])[:, 0], torch.ones(NUM_ENVS, device=DEVICE)) + + velocity_range["x"] = (4.0, 4.0) + warp_evt.push_by_setting_velocity(env, env_mask, velocity_range=velocity_range) + wp.synchronize() + + assert_close(wp.to_torch(captured["root_velocity"])[:, 0], torch.full((NUM_ENVS,), 4.0, device=DEVICE)) + + def test_push_state_respects_sparse_mask(self): + env, data, asset = _make_event_env(303) + copy_np_to_wp(data.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + captured = {} + asset.write_root_velocity_to_sim_mask = lambda **kwargs: captured.update(kwargs) + mask_np = np.arange(NUM_ENVS) % 3 == 0 + env_mask = wp.array(mask_np, dtype=wp.bool, device=DEVICE) + + warp_evt.push_by_setting_velocity(env, env_mask, velocity_range={"x": (3.0, 3.0)}) + wp.synchronize() + + velocity = wp.to_torch(captured["root_velocity"]) + mask = torch.from_numpy(mask_np).to(device=DEVICE) + assert_close(velocity[mask, 0], torch.full((int(mask_np.sum()),), 3.0, device=DEVICE)) + assert_close(velocity[~mask], torch.zeros((int((~mask_np).sum()), 6), device=DEVICE)) + + def test_external_wrench_state_is_not_shared_between_environments(self): + env_a, _, asset_a = _make_event_env(404, num_bodies=2) + env_b, _, asset_b = _make_event_env(505, num_bodies=2) + captured_a = {} + captured_b = {} + asset_a.permanent_wrench_composer.set_forces_and_torques_mask = lambda **kwargs: captured_a.update(kwargs) + asset_b.permanent_wrench_composer.set_forces_and_torques_mask = lambda **kwargs: captured_b.update(kwargs) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + force_range_a = (1.0, 1.0) + force_range_b = (2.0, 2.0) + torque_range = (0.0, 0.0) + + warp_evt.apply_external_force_torque(env_a, env_mask, force_range=force_range_a, torque_range=torque_range) + warp_evt.apply_external_force_torque(env_b, env_mask, force_range=force_range_b, torque_range=torque_range) + wp.synchronize() + + forces_a = captured_a["forces"] + forces_b = captured_b["forces"] + assert forces_a.ptr != forces_b.ptr + assert_close(wp.to_torch(forces_a), torch.ones((NUM_ENVS, 2, 3), device=DEVICE)) + assert_close(wp.to_torch(forces_b), torch.full((NUM_ENVS, 2, 3), 2.0, device=DEVICE)) + + def test_external_wrench_respects_body_selection(self): + env, _, asset = _make_event_env(550, num_bodies=3) + captured = {} + asset.permanent_wrench_composer.set_forces_and_torques_mask = lambda **kwargs: captured.update(kwargs) + asset_cfg = SceneEntityCfg("robot", body_ids=[1]) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + + warp_evt.apply_external_force_torque( + env, + env_mask, + force_range=(2.0, 2.0), + torque_range=(0.0, 0.0), + asset_cfg=asset_cfg, + ) + wp.synchronize() + + expected_body_mask = torch.tensor([False, True, False], dtype=torch.bool, device=DEVICE) + assert torch.equal(wp.to_torch(captured["body_mask"]), expected_body_mask) + forces = wp.to_torch(captured["forces"]) + assert_close(forces[:, 0], torch.zeros((NUM_ENVS, 3), device=DEVICE)) + assert_close(forces[:, 1], torch.full((NUM_ENVS, 3), 2.0, device=DEVICE)) + assert_close(forces[:, 2], torch.zeros((NUM_ENVS, 3), device=DEVICE)) + + def test_root_reset_state_is_not_shared_between_environments(self): + env_a, data_a, asset_a = _make_event_env(606) + env_b, data_b, asset_b = _make_event_env(707) + default_pose = np.zeros((NUM_ENVS, 7), dtype=np.float32) + default_pose[:, 6] = 1.0 + copy_np_to_wp(data_a.default_root_pose, default_pose) + copy_np_to_wp(data_b.default_root_pose, default_pose) + copy_np_to_wp(data_a.default_root_vel, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + copy_np_to_wp(data_b.default_root_vel, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + captured_a = {} + captured_b = {} + asset_a.write_root_pose_to_sim_mask = lambda **kwargs: captured_a.update(kwargs) + asset_b.write_root_pose_to_sim_mask = lambda **kwargs: captured_b.update(kwargs) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + pose_range_a = {"x": (1.0, 1.0)} + pose_range_b = {"x": (2.0, 2.0)} + velocity_range = {} + + warp_evt.reset_root_state_uniform(env_a, env_mask, pose_range=pose_range_a, velocity_range=velocity_range) + warp_evt.reset_root_state_uniform(env_b, env_mask, pose_range=pose_range_b, velocity_range=velocity_range) + wp.synchronize() + + pose_a = captured_a["root_pose"] + pose_b = captured_b["root_pose"] + assert pose_a.ptr != pose_b.ptr + assert_close(wp.to_torch(pose_a)[:, 0], torch.ones(NUM_ENVS, device=DEVICE)) + assert_close(wp.to_torch(pose_b)[:, 0], torch.full((NUM_ENVS,), 2.0, device=DEVICE)) + + def test_com_ranges_are_not_shared_between_environments(self): + env_a, data_a, asset_a = _make_event_env(808, num_bodies=2) + env_b, data_b, asset_b = _make_event_env(909, num_bodies=2) + copy_np_to_wp(data_a.body_com_pos_b, np.zeros((NUM_ENVS, 2, 3), dtype=np.float32)) + copy_np_to_wp(data_b.body_com_pos_b, np.zeros((NUM_ENVS, 2, 3), dtype=np.float32)) + asset_a.set_coms_mask = lambda **kwargs: None + asset_b.set_coms_mask = lambda **kwargs: None + asset_cfg_a = SceneEntityCfg("robot") + asset_cfg_b = SceneEntityCfg("robot") + asset_cfg_a.body_ids_wp = wp.array([0, 1], dtype=wp.int32, device=DEVICE) + asset_cfg_b.body_ids_wp = wp.array([0, 1], dtype=wp.int32, device=DEVICE) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + com_range_a = {"x": (1.0, 1.0)} + com_range_b = {"x": (2.0, 2.0)} + + warp_evt.randomize_rigid_body_com(env_a, env_mask, com_range=com_range_a, asset_cfg=asset_cfg_a) + warp_evt.randomize_rigid_body_com(env_b, env_mask, com_range=com_range_b, asset_cfg=asset_cfg_b) + wp.synchronize() + + assert env_a._warp_event_state_cache is not env_b._warp_event_state_cache + assert_close(data_a.body_com_pos_b.torch[..., 0], torch.ones((NUM_ENVS, 2), device=DEVICE)) + assert_close(data_b.body_com_pos_b.torch[..., 0], torch.full((NUM_ENVS, 2), 2.0, device=DEVICE)) + + def test_com_randomization_uses_cached_baseline(self): + env, data, asset = _make_event_env(1001, num_bodies=2) + baseline = np.zeros((NUM_ENVS, 2, 3), dtype=np.float32) + baseline[..., 0] = 0.25 + copy_np_to_wp(data.body_com_pos_b, baseline) + asset.set_coms_mask = lambda **kwargs: None + asset_cfg = SceneEntityCfg("robot", body_ids=[0, 1]) + asset_cfg.body_ids_wp = wp.array([0, 1], dtype=wp.int32, device=DEVICE) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + com_range = {"x": (1.0, 1.0)} + + warp_evt.randomize_rigid_body_com(env, env_mask, com_range=com_range, asset_cfg=asset_cfg) + wp.synchronize() + assert_close(data.body_com_pos_b.torch[..., 0], torch.full((NUM_ENVS, 2), 1.25, device=DEVICE)) + + warp_evt.randomize_rigid_body_com(env, env_mask, com_range=com_range, asset_cfg=asset_cfg) + wp.synchronize() + assert_close(data.body_com_pos_b.torch[..., 0], torch.full((NUM_ENVS, 2), 1.25, device=DEVICE)) + + def test_com_randomization_broadcasts_one_offset_across_selected_bodies(self): + env, data, asset = _make_event_env(1111, num_bodies=3) + baseline = np.zeros((NUM_ENVS, 3, 3), dtype=np.float32) + copy_np_to_wp(data.body_com_pos_b, baseline) + asset.set_coms_mask = lambda **kwargs: None + asset_cfg = SceneEntityCfg("robot", body_ids=[0, 2]) + asset_cfg.body_ids_wp = wp.array([0, 2], dtype=wp.int32, device=DEVICE) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + + warp_evt.randomize_rigid_body_com( + env, + env_mask, + com_range={"x": (-1.0, 1.0), "y": (-2.0, 2.0), "z": (-3.0, 3.0)}, + asset_cfg=asset_cfg, + ) + wp.synchronize() + + assert_close(data.body_com_pos_b.torch[:, 0], data.body_com_pos_b.torch[:, 2]) + assert_close(data.body_com_pos_b.torch[:, 1], torch.zeros((NUM_ENVS, 3), device=DEVICE)) diff --git a/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py index 9f1eea23bf8d..8654949bbed5 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py @@ -47,15 +47,6 @@ # ============================================================================ -@pytest.fixture(autouse=True) -def _clear_caches(): - yield - for fn in [warp_obs.generated_commands]: - for attr in list(vars(fn)): - if attr.startswith("_"): - delattr(fn, attr) - - @pytest.fixture() def art_data(): return MockArticulationData(NUM_ENVS, NUM_JOINTS, DEVICE) diff --git a/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py index 69cabcd14a36..10c9be10ec6b 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py @@ -55,7 +55,7 @@ @pytest.fixture(autouse=True) def _clear_caches(): yield - for fn in [warp_rew.track_lin_vel_xy_exp, warp_rew.track_ang_vel_z_exp, warp_rew.undesired_contacts]: + for fn in [warp_rew.undesired_contacts]: for attr in list(vars(fn)): if attr.startswith("_"): delattr(fn, attr) diff --git a/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py b/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py new file mode 100644 index 000000000000..dbd56df36754 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py @@ -0,0 +1,64 @@ +# 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 + +"""Tests for mask-first scene reset dispatch.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import warp as wp +from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp + + +class TestInteractiveSceneWarp: + """Tests for :class:`InteractiveSceneWarp`.""" + + def test_mask_reset_stays_mask_based_for_supported_entities(self): + """Mask-based reset should not pass environment IDs to Warp-capable entities.""" + scene = InteractiveSceneWarp.__new__(InteractiveSceneWarp) + articulation = Mock() + deformable = Mock() + rigid_object = Mock() + rigid_collection = Mock() + sensor = Mock() + surface_gripper = Mock() + scene._articulations = {"articulation": articulation} + scene._deformable_objects = {"deformable": deformable} + scene._rigid_objects = {"rigid_object": rigid_object} + scene._rigid_object_collections = {"collection": rigid_collection} + scene._sensors = {"sensor": sensor} + scene._surface_grippers = {"gripper": surface_gripper} + scene.cfg = SimpleNamespace(num_envs=3) + scene.sim = SimpleNamespace(device="cpu") + env_mask = wp.array([True, False, True], dtype=wp.bool, device="cpu") + + scene.reset(env_mask=env_mask) + + for entity in (articulation, deformable, rigid_object, rigid_collection, sensor): + entity.reset.assert_called_once_with(env_mask=env_mask) + surface_gripper.reset.assert_not_called() + + def test_surface_gripper_reset_is_an_explicit_id_boundary(self): + """Surface grippers should be reset separately because their API is ID-based.""" + scene = InteractiveSceneWarp.__new__(InteractiveSceneWarp) + surface_gripper = Mock() + scene._surface_grippers = {"gripper": surface_gripper} + + scene.reset_host([0, 2]) + + surface_gripper.reset.assert_called_once_with([0, 2]) + + def test_mask_shape_is_validated_before_entity_dispatch(self): + """A malformed mask should fail before reaching scene entities.""" + scene = InteractiveSceneWarp.__new__(InteractiveSceneWarp) + scene.cfg = SimpleNamespace(num_envs=3) + scene.sim = SimpleNamespace(device="cpu") + env_mask = wp.array([True, False], dtype=wp.bool, device="cpu") + + with pytest.raises(ValueError, match="shape"): + scene.reset(env_mask=env_mask) diff --git a/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py new file mode 100644 index 000000000000..00d4c09222be --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py @@ -0,0 +1,192 @@ +# 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 + +"""Tests for mask-first manager-based RL reset orchestration.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +import warp as wp +from isaaclab_experimental.envs.manager_based_rl_env_warp import ManagerBasedRLEnvWarp +from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode + + +def test_warp_environment_deprecates_and_preserves_stable_manager_profile() -> None: + """The legacy stable-manager profile should warn while retaining its config bridge.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + env.cfg = SimpleNamespace(rewards="warp_rewards", actions="warp_actions") + env._manager_call_switch = SimpleNamespace( + get_mode_for_manager=lambda name: ( + ManagerCallMode.STABLE if name == "RewardManager" else ManagerCallMode.WARP_NOT_CAPTURED + ), + ) + env._resolve_stable_cfg_counterpart = Mock( + return_value=SimpleNamespace(rewards="stable_rewards", actions="stable_actions") + ) + + with pytest.warns(DeprecationWarning, match="STABLE managers"): + env._apply_manager_term_cfg_profile() + + assert env.cfg.rewards == "stable_rewards" + assert env.cfg.actions == "warp_actions" + + +def test_reset_without_host_consumers_does_not_compact_ids() -> None: + """The normal Warp reset path should not call :meth:`torch.Tensor.nonzero`.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") + env.termination_manager = SimpleNamespace(dones_wp=reset_mask) + env.reset_buf = Mock() + env.reset_buf.nonzero.side_effect = AssertionError("reset IDs should not be materialized") + env._reset_requires_host_selection = Mock(return_value=False) + env._reset_idx = Mock() + env.extras = {"log": {}} + + env._reset_terminated_envs() + + env.reset_buf.nonzero.assert_not_called() + env._reset_idx.assert_called_once_with(env_mask=reset_mask, env_ids=None) + + +def test_public_reset_selection_stays_as_mask_without_host_consumers() -> None: + """Mask-native reset selection should not be compacted preemptively.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + env._reset_requires_host_selection = Mock(return_value=False) + reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") + + result = env._resolve_reset_host_ids(env_ids=None, env_mask=reset_mask) + + assert result is None + + +def test_public_reset_compacts_mask_only_for_host_consumer() -> None: + """A mask should cross to IDs only at an enabled host reset boundary.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + env._reset_requires_host_selection = Mock(return_value=True) + reset_mask = wp.array([False, True, False, True], dtype=wp.bool, device="cpu") + + result = env._resolve_reset_host_ids(env_ids=None, env_mask=reset_mask) + + torch.testing.assert_close(result, torch.tensor([1, 3])) + + +def test_curriculum_uses_mask_before_scene_reset() -> None: + """Curriculum updates should stay mask-native and precede scene reset.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + stages: list[tuple[str, dict]] = [] + + def call(stage, fn, /, *args, **kwargs): + del fn, args + stages.append((stage, kwargs)) + return None if stage.endswith("_compute") or stage == "Scene_reset" else {} + + env._manager_call_switch = SimpleNamespace(call=call) + env.sim = SimpleNamespace(device="cpu") + env.reset_mask_wp = wp.zeros(3, dtype=wp.bool, device="cpu") + env.scene = SimpleNamespace(num_envs=3, reset=Mock(), surface_grippers={}) + env.curriculum_manager = SimpleNamespace(compute=Mock(), reset=Mock()) + env.event_manager = SimpleNamespace(available_modes=[], reset=Mock()) + env.observation_manager = SimpleNamespace(reset=Mock()) + env.action_manager = SimpleNamespace(reset=Mock()) + env.reward_manager = SimpleNamespace(reset=Mock()) + env.command_manager = SimpleNamespace(reset=Mock()) + env.termination_manager = SimpleNamespace(reset=Mock()) + env._episode_length_buf_wp = wp.ones(3, dtype=wp.int64, device="cpu") + env.extras = {} + reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") + reset_ids = torch.tensor([1]) + + env._reset_idx(env_mask=reset_mask, env_ids=reset_ids) + + assert [stage for stage, _ in stages[:2]] == ["CurriculumManager_compute", "Scene_reset"] + assert stages[0][1]["env_mask"] is env.reset_mask_wp + np.testing.assert_array_equal(stages[0][1]["env_mask"].numpy(), reset_mask.numpy()) + assert stages[0][1]["env_ids"] is reset_ids + assert "CurriculumManager_reset" in [stage for stage, _ in stages] + + +def test_reset_stages_reuse_owner_mask_for_sparse_and_empty_selections() -> None: + """Every reset dispatch should receive the stable owner-held mask pointer.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + seen: list[tuple[wp.array, np.ndarray]] = [] + + def call(stage, fn, /, *args, **kwargs): + del fn, args + if stage == "CurriculumManager_compute": + mask = kwargs["env_mask"] + seen.append((mask, mask.numpy().copy())) + return None if stage.endswith("_compute") or stage == "Scene_reset" else {} + + env._manager_call_switch = SimpleNamespace(call=call) + env.sim = SimpleNamespace(device="cpu") + env.reset_mask_wp = wp.zeros(3, dtype=wp.bool, device="cpu") + env.scene = SimpleNamespace(num_envs=3, reset=Mock(), surface_grippers={}) + env.curriculum_manager = SimpleNamespace(compute=Mock(), reset=Mock()) + env.event_manager = SimpleNamespace(available_modes=[], reset=Mock()) + env.observation_manager = SimpleNamespace(reset=Mock()) + env.action_manager = SimpleNamespace(reset=Mock()) + env.reward_manager = SimpleNamespace(reset=Mock()) + env.command_manager = SimpleNamespace(reset=Mock()) + env.termination_manager = SimpleNamespace(reset=Mock()) + env._episode_length_buf_wp = wp.ones(3, dtype=wp.int64, device="cpu") + env.extras = {} + + env._reset_idx(env_mask=wp.array([False, True, False], dtype=wp.bool, device="cpu")) + env._reset_idx(env_mask=wp.zeros(3, dtype=wp.bool, device="cpu")) + + assert seen[0][0] is env.reset_mask_wp + assert seen[1][0] is env.reset_mask_wp + np.testing.assert_array_equal(seen[0][1], [False, True, False]) + np.testing.assert_array_equal(seen[1][1], [False, False, False]) + + +def test_reset_compacts_ids_for_enabled_host_consumer() -> None: + """Host-only reset hooks should receive IDs while the Warp stage keeps its mask.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") + env.termination_manager = SimpleNamespace(dones_wp=reset_mask) + env.reset_buf = torch.tensor([False, True, False]) + env._reset_requires_host_selection = Mock(return_value=True) + env._reset_idx = Mock() + env._reset_host_pre = Mock() + env._reset_host_post = Mock(return_value={"host": 1.0}) + env.recorder_manager = Mock() + env._has_recorders = False + env.has_rtx_sensors = False + env.cfg = SimpleNamespace(num_rerenders_on_reset=0) + env.extras = {"log": {}} + + env._reset_terminated_envs() + + reset_ids = env._reset_host_pre.call_args.args[0] + torch.testing.assert_close(reset_ids, torch.tensor([1])) + env._reset_idx.assert_called_once_with(env_mask=reset_mask, env_ids=reset_ids) + assert env.extras["log"]["host"] == 1.0 + + +def test_empty_reset_skips_legacy_host_boundary_and_preserves_logs() -> None: + """An empty host-visible selection should not execute reset side effects.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + reset_mask = wp.zeros(3, dtype=wp.bool, device="cpu") + env.termination_manager = SimpleNamespace(dones_wp=reset_mask) + env.reset_buf = torch.zeros(3, dtype=torch.bool) + env._reset_requires_host_selection = Mock(return_value=True) + env._reset_idx = Mock() + env._reset_host_pre = Mock() + env._reset_host_post = Mock() + env._has_recorders = False + env.extras = {"log": {"previous": 1.0}} + + env._reset_terminated_envs() + + env._reset_idx.assert_not_called() + env._reset_host_pre.assert_not_called() + env._reset_host_post.assert_not_called() + assert env.extras["log"] == {"previous": 1.0} diff --git a/source/isaaclab_experimental/test/managers/test_curriculum_manager.py b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py new file mode 100644 index 000000000000..dcd9dbac6cf3 --- /dev/null +++ b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py @@ -0,0 +1,322 @@ +# 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 + +"""Tests for the Warp-first curriculum manager and terrain-level update path.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +import warp as wp +from isaaclab_experimental.managers import CurriculumManager, CurriculumTermCfg +from isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp import TerrainLevelsVel, terrain_levels_vel +from isaaclab_tasks_experimental.manager_based.manipulation.reach.reach_env_cfg import ( + CurriculumCfg as ReachCurriculumCfg, +) + +from isaaclab.managers import ManagerTermBase as StableManagerTermBase +from isaaclab.terrains import TerrainImporter + + +class _GlobalCurriculumTerm(StableManagerTermBase): + """Legacy global term that intentionally ignores compact environment IDs.""" + + def __call__(self, env, env_ids, value: float) -> float: + assert env_ids == slice(None) + return value + + +def _legacy_id_count(env, env_ids, scale: float) -> float: + """Legacy curriculum state that genuinely depends on compact environment IDs.""" + return len(env_ids) * scale + + +def _structured_state(env, env_ids) -> dict[str, float]: + """Legacy curriculum state unsupported by the scalar Warp logging buffer.""" + del env, env_ids + return {"value": 1.0} + + +class _ArrayProxy: + """Minimal Torch/Warp array proxy.""" + + def __init__(self, values: np.ndarray, dtype): + self.warp = wp.array(values, dtype=dtype, device="cpu") + self.torch = wp.to_torch(self.warp) + + +class _Robot: + """Minimal robot exposing root positions.""" + + def __init__(self, root_positions: np.ndarray): + self.cfg = SimpleNamespace() + self.data = SimpleNamespace(root_pos_w=_ArrayProxy(root_positions, wp.vec3f)) + + +class _CommandManager: + """Minimal command manager exposing persistent Warp storage.""" + + def __init__(self, commands: np.ndarray): + self.command_wp = wp.array(commands, dtype=wp.float32, device="cpu") + + def get_command_wp(self, name: str) -> wp.array(dtype=wp.float32, ndim=2): + assert name == "base_velocity" + return self.command_wp + + +class _RewardManager: + """Minimal reward manager used by the registered Reach curricula.""" + + def __init__(self): + self._term_names = ["action_rate", "joint_vel"] + self._cfgs = { + "action_rate": SimpleNamespace(weight=-0.01), + "joint_vel": SimpleNamespace(weight=-0.0001), + } + self._weights_wp = wp.array([-0.01, -0.0001], dtype=wp.float32, device="cpu") + stride = self._weights_wp.strides[0] + self._weight_views_wp = { + name: wp.array( + ptr=self._weights_wp.ptr + term_idx * stride, + dtype=wp.float32, + shape=(1,), + strides=(stride,), + device="cpu", + ) + for term_idx, name in enumerate(self._term_names) + } + + def get_term_cfg(self, name: str): + return self._cfgs[name] + + def set_term_cfg(self, name: str, cfg) -> None: + self._cfgs[name] = cfg + + def get_term_weight_wp(self, name: str) -> wp.array(dtype=wp.float32): + return self._weight_views_wp[name] + + +class _Scene: + """Minimal scene mapping with terrain-backed environment origins.""" + + def __init__(self, robot: _Robot, terrain: TerrainImporter): + self._entities = {"robot": robot} + self.terrain = terrain + self.device = "cpu" + + @property + def env_origins(self) -> torch.Tensor: + return self.terrain.env_origins + + def __getitem__(self, name: str): + return self._entities[name] + + def keys(self): + return self._entities.keys() + + +class _Env: + """Minimal environment for curriculum manager execution.""" + + def __init__(self, terrain: TerrainImporter, root_positions: np.ndarray, commands: np.ndarray): + self.num_envs = root_positions.shape[0] + self.device = "cpu" + self.scene = _Scene(_Robot(root_positions), terrain) + self.command_manager = _CommandManager(commands) + self.env_origins_wp = wp.from_torch(terrain.env_origins, dtype=wp.vec3f) + self.rng_state_wp = wp.array(np.arange(self.num_envs, dtype=np.uint32) + 101, device=self.device) + self.max_episode_length_s = 10.0 + self.sim = SimpleNamespace(is_playing=lambda: True) + + +def _make_terrain(levels: list[int]) -> TerrainImporter: + num_envs = len(levels) + terrain = TerrainImporter.__new__(TerrainImporter) + terrain.device = "cpu" + terrain.cfg = SimpleNamespace(terrain_generator=SimpleNamespace(size=(8.0, 8.0))) + terrain.max_terrain_level = 3 + terrain.terrain_levels = torch.tensor(levels, dtype=torch.int64) + terrain.terrain_types = torch.tensor([index % 2 for index in range(num_envs)], dtype=torch.int64) + terrain.terrain_origins = torch.zeros(3, 2, 3, dtype=torch.float32) + for level in range(3): + for terrain_type in range(2): + terrain.terrain_origins[level, terrain_type] = torch.tensor( + [100.0 * level + 10.0 * terrain_type, float(level), float(terrain_type)] + ) + terrain.env_origins = terrain.terrain_origins[terrain.terrain_levels, terrain.terrain_types].clone() + terrain._invalidate_warp_origin_views() + return terrain + + +def test_terrain_importer_mask_update_preserves_sparse_and_stable_semantics(): + """Mask updates should handle up/down/clamp/wrap while stable IDs remain supported.""" + terrain = _make_terrain([0, 1, 2, 2, 1, 0]) + env_mask = wp.array([True, True, True, True, False, False], dtype=wp.bool, device="cpu") + move_up = wp.array([False, True, True, False, True, False], dtype=wp.bool, device="cpu") + move_down = wp.array([True, False, False, True, False, True], dtype=wp.bool, device="cpu") + rng_state = wp.array(np.arange(6, dtype=np.uint32) + 71, device="cpu") + rng_before = rng_state.numpy().copy() + origins_before = terrain.env_origins.clone() + levels_wp = terrain.terrain_levels_wp + + terrain.update_env_origins_mask(env_mask, move_up, move_down, rng_state) + wp.synchronize() + + levels = terrain.terrain_levels.tolist() + assert levels[0] == 0 + assert levels[1] == 2 + assert 0 <= levels[2] < terrain.max_terrain_level + assert levels[3] == 1 + assert levels[4:] == [1, 0] + torch.testing.assert_close( + terrain.env_origins[:4], terrain.terrain_origins[terrain.terrain_levels[:4], terrain.terrain_types[:4]] + ) + torch.testing.assert_close(terrain.env_origins[4:], origins_before[4:]) + assert np.all(rng_state.numpy()[:4] != rng_before[:4]) + np.testing.assert_array_equal(rng_state.numpy()[4:], rng_before[4:]) + assert terrain.terrain_levels_wp is levels_wp + + levels_before_stable_update = terrain.terrain_levels.clone() + env_ids = torch.tensor([1, 3], dtype=torch.int64) + terrain.update_env_origins( + env_ids, + move_up=torch.tensor([False, False]), + move_down=torch.tensor([True, True]), + ) + expected_levels = levels_before_stable_update.clone() + expected_levels[env_ids] -= 1 + torch.testing.assert_close(terrain.terrain_levels, expected_levels) + assert terrain.terrain_levels_wp is levels_wp + + +def test_curriculum_manager_updates_masked_levels_and_logs_without_host_compaction(monkeypatch: pytest.MonkeyPatch): + """Manager compute/reset should remain mask-native and expose persistent scalar logging.""" + assert terrain_levels_vel is not TerrainLevelsVel + terrain = _make_terrain([0, 1, 2, 2, 1, 0]) + root_positions = terrain.env_origins.numpy().copy() + root_positions[:, 0] += np.array([5.0, 1.0, 3.0, 5.0, 9.0, 0.0], dtype=np.float32) + commands = np.array( + [[0.1, 0.0, 0.0], [1.0, 0.0, 0.0], [0.2, 0.0, 0.0], [2.0, 0.0, 0.0], [0.1, 0.0, 0.0], [0.1, 0.0, 0.0]], + dtype=np.float32, + ) + env = _Env(terrain, root_positions, commands) + manager = CurriculumManager( + { + "terrain_levels": CurriculumTermCfg(func=TerrainLevelsVel), + "global_state": CurriculumTermCfg( + func=_GlobalCurriculumTerm, + params={"value": 7.0}, + requires_host_ids=False, + ), + }, + env, + ) + term = manager._term_cfgs[0].func + env_mask = wp.array([True, True, True, True, False, False], dtype=wp.bool, device="cpu") + move_up_ptr = term._move_up_wp.ptr + move_down_ptr = term._move_down_wp.ptr + state_ptr = manager._term_states_wp.ptr + extras_ref = manager.reset_extras + assert not manager.requires_host_ids + assert manager.requires_host_boundary + + def _fail_host_compaction(*args, **kwargs): + raise AssertionError("Torch host compaction/scalar extraction reached the Warp curriculum path") + + with monkeypatch.context() as context: + context.setattr(torch.Tensor, "nonzero", _fail_host_compaction) + context.setattr(torch.Tensor, "item", _fail_host_compaction) + manager.compute(env_mask) + extras = manager.reset(env_mask) + wp.synchronize() + + levels = terrain.terrain_levels.tolist() + assert levels[0] == 1 + assert levels[1] == 0 + assert levels[2] == 2 + assert 0 <= levels[3] < terrain.max_terrain_level + assert levels[4:] == [1, 0] + np.testing.assert_array_equal(term._move_up_wp.numpy(), [True, False, False, True, False, False]) + np.testing.assert_array_equal(term._move_down_wp.numpy(), [False, True, False, False, False, False]) + assert extras is extras_ref + assert extras["Curriculum/terrain_levels"] is manager.reset_extras["Curriculum/terrain_levels"] + torch.testing.assert_close(extras["Curriculum/terrain_levels"], terrain.terrain_levels.float().mean()) + torch.testing.assert_close(extras["Curriculum/global_state"], torch.tensor(7.0)) + assert term._move_up_wp.ptr == move_up_ptr + assert term._move_down_wp.ptr == move_down_ptr + assert manager._term_states_wp.ptr == state_ptr + + +def test_curriculum_manager_requests_ids_only_for_legacy_id_terms(): + """Legacy terms should require host IDs only when their config says they consume them.""" + terrain = _make_terrain([0, 1, 2, 2]) + env = _Env( + terrain, + terrain.env_origins.numpy().copy(), + np.zeros((4, 3), dtype=np.float32), + ) + manager = CurriculumManager( + {"id_count": CurriculumTermCfg(func=_legacy_id_count, params={"scale": 2.0})}, + env, + ) + env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + + assert manager.requires_host_ids + assert manager.requires_host_boundary + with pytest.raises(RuntimeError, match="requires compact env_ids"): + manager.compute(env_mask) + + manager.compute(env_mask, env_ids=[0, 2]) + extras = manager.reset(env_mask, env_ids=[0, 2]) + torch.testing.assert_close(extras["Curriculum/id_count"], torch.tensor(4.0)) + + +def test_registered_reach_curricula_update_weights_without_a_host_boundary(): + """Registered Reach reward schedules should update only on a device-selected reset.""" + terrain = _make_terrain([0, 1, 2, 2]) + env = _Env( + terrain, + terrain.env_origins.numpy().copy(), + np.zeros((4, 3), dtype=np.float32), + ) + env.reward_manager = _RewardManager() + env.common_step_counter = 0 + manager = CurriculumManager(ReachCurriculumCfg(), env) + env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + + assert not manager.requires_host_ids + assert not manager.requires_host_boundary + manager.compute(env_mask) + extras = manager.reset(env_mask) + torch.testing.assert_close(extras["Curriculum/action_rate"], torch.tensor(-0.01)) + torch.testing.assert_close(extras["Curriculum/joint_vel"], torch.tensor(-0.0001)) + + env.common_step_counter = 4501 + empty_mask = wp.zeros(4, dtype=wp.bool, device="cpu") + manager.compute(empty_mask) + np.testing.assert_allclose(env.reward_manager._weights_wp.numpy(), [-0.01, -0.0001]) + + manager.compute(env_mask) + extras = manager.reset(env_mask) + torch.testing.assert_close(extras["Curriculum/action_rate"], torch.tensor(-0.005)) + torch.testing.assert_close(extras["Curriculum/joint_vel"], torch.tensor(-0.001)) + np.testing.assert_allclose(env.reward_manager._weights_wp.numpy(), [-0.005, -0.001]) + + +def test_curriculum_manager_rejects_structured_legacy_state() -> None: + """Unsupported fallback logging should fail explicitly instead of narrowing silently.""" + terrain = _make_terrain([0, 1]) + env = _Env(terrain, terrain.env_origins.numpy().copy(), np.zeros((2, 3), dtype=np.float32)) + manager = CurriculumManager( + {"structured": CurriculumTermCfg(func=_structured_state, requires_host_ids=False)}, + env, + ) + env_mask = wp.array([True, False], dtype=wp.bool, device="cpu") + + with pytest.raises(TypeError, match="scalar logging states"): + manager.compute(env_mask) diff --git a/source/isaaclab_experimental/test/managers/test_manager_base.py b/source/isaaclab_experimental/test/managers/test_manager_base.py new file mode 100644 index 000000000000..d2cdef5e2553 --- /dev/null +++ b/source/isaaclab_experimental/test/managers/test_manager_base.py @@ -0,0 +1,23 @@ +# 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 + +"""Tests for Warp manager reset-mask contracts.""" + +from types import SimpleNamespace + +import pytest +import warp as wp +from isaaclab_experimental.managers.manager_base import _resolve_reset_mask + + +def test_reset_mask_rejects_wrong_device(monkeypatch: pytest.MonkeyPatch) -> None: + """Manager kernels should fail before receiving a mask from another device.""" + owner = SimpleNamespace(num_envs=2, device="expected") + env_mask = wp.zeros(2, dtype=wp.bool, device="cpu") + expected_device = SimpleNamespace(alias="expected") + monkeypatch.setattr(wp, "get_device", lambda device: expected_device) + + with pytest.raises(ValueError, match="env_mask must be on"): + _resolve_reset_mask(owner, None, env_mask) diff --git a/source/isaaclab_experimental/test/managers/test_reward_manager.py b/source/isaaclab_experimental/test/managers/test_reward_manager.py new file mode 100644 index 000000000000..569465822377 --- /dev/null +++ b/source/isaaclab_experimental/test/managers/test_reward_manager.py @@ -0,0 +1,41 @@ +# 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 + +"""Tests for device-owned Warp reward weights.""" + +from types import SimpleNamespace + +import torch +import warp as wp +from isaaclab_experimental.managers import RewardManager, RewardTermCfg + + +@wp.kernel +def _fill_unit_reward(out: wp.array(dtype=wp.float32)): + out[wp.tid()] = 1.0 + + +def _unit_reward(env, out: wp.array(dtype=wp.float32)) -> None: + env.term_call_count += 1 + wp.launch(_fill_unit_reward, dim=env.num_envs, inputs=[out], device=env.device) + + +def test_device_owned_weight_can_enable_zero_weight_term() -> None: + """A curriculum-owned weight should be authoritative over the initial config scalar.""" + env = SimpleNamespace( + num_envs=4, + device="cpu", + sim=SimpleNamespace(is_playing=lambda: False), + term_call_count=0, + ) + manager = RewardManager({"scheduled": RewardTermCfg(func=_unit_reward, weight=0.0)}, env) + + torch.testing.assert_close(manager.compute(dt=1.0), torch.zeros(4)) + assert env.term_call_count == 0 + + manager.get_term_weight_wp("scheduled").fill_(2.0) + torch.testing.assert_close(manager.compute(dt=1.0), torch.full((4,), 2.0)) + assert env.term_call_count == 1 + assert manager.get_term_cfg("scheduled").weight == 2.0 diff --git a/source/isaaclab_experimental/test/managers/test_termination_manager.py b/source/isaaclab_experimental/test/managers/test_termination_manager.py new file mode 100644 index 000000000000..07605af1d853 --- /dev/null +++ b/source/isaaclab_experimental/test/managers/test_termination_manager.py @@ -0,0 +1,37 @@ +# 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 + +"""Tests for the Warp-first termination manager.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import warp as wp +from isaaclab_experimental.managers.termination_manager import TerminationManager + + +class TestTerminationManager: + """Tests for :class:`TerminationManager`.""" + + def test_reset_statistics_only_include_selected_environments(self): + """Partial reset statistics should be averaged over selected environments only.""" + manager = TerminationManager.__new__(TerminationManager) + manager._env = SimpleNamespace(num_envs=4, device="cpu") + manager._term_names = ["fall", "timeout"] + manager._last_episode_dones_wp = wp.array( + [[True, False], [False, True], [True, True], [False, False]], dtype=wp.bool, device="cpu" + ) + manager._term_done_avg_wp = wp.zeros(2, dtype=wp.float32, device="cpu") + manager._reset_count_wp = wp.zeros(1, dtype=wp.int32, device="cpu") + manager._reset_scale_wp = wp.zeros(1, dtype=wp.float32, device="cpu") + manager._class_term_cfgs = [] + manager._reset_extras = {} + env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + + manager.reset(env_mask=env_mask) + + np.testing.assert_allclose(manager._term_done_avg_wp.numpy(), [1.0, 0.5]) diff --git a/source/isaaclab_experimental/test/utils/test_manager_call_switch.py b/source/isaaclab_experimental/test/utils/test_manager_call_switch.py index b4745939eaa0..029c2eb236cf 100644 --- a/source/isaaclab_experimental/test/utils/test_manager_call_switch.py +++ b/source/isaaclab_experimental/test/utils/test_manager_call_switch.py @@ -44,7 +44,7 @@ class TestConfigLoading(unittest.TestCase): def test_none_uses_default(self): switch = ManagerCallSwitch(cfg_source=None) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_CAPTURED) + self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_NOT_CAPTURED) def test_dict_config(self): switch = ManagerCallSwitch(cfg_source={"default": 1}) @@ -58,8 +58,8 @@ def test_dict_per_manager_override(self): def test_dict_without_default_key(self): """A dict missing 'default' should inherit from DEFAULT_CONFIG.""" switch = ManagerCallSwitch(cfg_source={"RewardManager": 0}) - # default should be 2 (from DEFAULT_CONFIG) - self.assertEqual(switch.get_mode_for_manager("ActionManager"), ManagerCallMode.WARP_CAPTURED) + # default should be Warp eager (from DEFAULT_CONFIG) + self.assertEqual(switch.get_mode_for_manager("ActionManager"), ManagerCallMode.WARP_NOT_CAPTURED) self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.STABLE) def test_json_string_config(self): @@ -70,7 +70,7 @@ def test_json_string_config(self): def test_empty_string_uses_default(self): switch = ManagerCallSwitch(cfg_source="") - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_CAPTURED) + self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_NOT_CAPTURED) def test_env_var_fallback(self): """When cfg_source is None, should read from MANAGER_CALL_CONFIG env var.""" @@ -190,6 +190,20 @@ def test_mode_0_not_affected_by_cap(self): class TestStageDispatch(unittest.TestCase): """Tests for call_stage routing through stable / warp-eager / warp-captured paths.""" + def test_call_forwards_normal_arguments_and_output(self): + """The concise frontend call should not require a call-spec dictionary.""" + switch = ManagerCallSwitch(cfg_source={"default": 1}) + + result = switch.call( + "RewardManager_compute", + lambda value, *, scale: value * scale, + 3, + scale=4, + _output=lambda value: value + 1, + ) + + self.assertEqual(result, 13) + def test_stable_mode_calls_stable_fn(self): switch = ManagerCallSwitch(cfg_source={"default": 0}) called = {"stable": False, "warp": False} @@ -303,6 +317,24 @@ def warp_fn(): self.assertIs(result, result2, "Replay must return same reference") self.assertAlmostEqual(result2.numpy()[0], 11.0, places=5) + def test_call_uses_the_same_signature_when_captured(self): + """The concise call should capture without changing the function signature.""" + switch = ManagerCallSwitch(cfg_source={"default": 2}) + src = wp.full(4, value=2.0, dtype=wp.float32, device=self.device) + dst = wp.zeros(4, dtype=wp.float32, device=self.device) + + def warp_fn(source, destination): + wp.launch(_add_one, dim=4, inputs=[source, destination], device=self.device) + return destination + + result = switch.call("RewardManager_compute", warp_fn, src, dst) + self.assertAlmostEqual(result.numpy()[0], 3.0, places=5) + + wp.copy(src, wp.full(4, value=8.0, dtype=wp.float32, device=self.device)) + replay = switch.call("RewardManager_compute", warp_fn, src, dst) + self.assertIs(result, replay) + self.assertAlmostEqual(replay.numpy()[0], 9.0, places=5) + def test_captured_warmup_call_count(self): """WARP_CAPTURED first call should invoke fn exactly 2 times (warm-up + capture).""" switch = ManagerCallSwitch(cfg_source={"default": 2}) diff --git a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py index 1b9c10a76895..abe94920bc33 100644 --- a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py +++ b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py @@ -27,6 +27,20 @@ def setUp(self): self.device = "cuda:0" self.cache = WarpGraphCache() + def test_disabled_cache_runs_eager_once_per_call(self): + """A disabled cache should preserve eager execution semantics.""" + cache = WarpGraphCache(enabled=False) + call_count = 0 + + def counted_call(value): + nonlocal call_count + call_count += 1 + return value + + self.assertEqual(cache.capture_or_replay("stage", counted_call, args=(1,)), 1) + self.assertEqual(cache.capture_or_replay("stage", counted_call, args=(2,)), 2) + self.assertEqual(call_count, 2) + # ------------------------------------------------------------------ # Warm-up # ------------------------------------------------------------------ diff --git a/source/isaaclab_newton/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab_newton/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst new file mode 100644 index 000000000000..7ed70c412633 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -0,0 +1,5 @@ +Added +^^^^^ + +* Added boolean environment-mask resets for Newton actuator and external-wrench + state, avoiding host index conversion in Warp environments. diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py index 689232b81e60..b8cc862de278 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py @@ -26,11 +26,12 @@ import warp as wp from newton.actuators import Actuator, Clamping, Delay +from isaaclab.utils.warp.utils import resolve_1d_mask + from .kernels import ( build_implicit_dof_mask, build_per_dof_env_mask_kernel, scatter_gain_kernel, - set_mask_kernel, zero_at_indices_kernel, ) @@ -112,6 +113,8 @@ def __init__( self._states_a = [act.state() for act in actuators] self._states_b = [act.state() for act in actuators] + self._reset_env_mask = wp.zeros(num_envs, dtype=wp.bool, device=device) + self._reset_dof_masks = [wp.zeros(act.indices.shape[0], dtype=wp.bool, device=device) for act in actuators] # Pre-clamp computed effort buffer. Each Newton actuator scatter-adds # its raw controller output to ``sim_control.joint_computed_f`` when @@ -166,14 +169,20 @@ def step(self, sim_state: Any, sim_control: Any, dt: float) -> None: act.step(sim_state, sim_control, sa, sb, dt=dt) self._states_a, self._states_b = self._states_b, self._states_a - def reset(self, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: + def reset( + self, + env_ids: Sequence[int] | slice | torch.Tensor | wp.array | None = None, + env_mask: wp.array(dtype=wp.bool) | None = None, + ) -> None: """Reset actuator states for the given environments. Args: - env_ids: Environment indices to reset. ``None`` (or - ``slice(None)``, which IsaacLab callers sometimes pass) - resets all environments. Otherwise expects a torch tensor - or sequence of int indices. + env_ids: Environment indices to reset. Accepts a sequence, slice, + Torch tensor, or ``wp.int32`` array. ``None`` and + ``slice(None)`` reset all environments. + env_mask: Boolean Warp mask selecting environments to reset. When + provided, it takes precedence over :paramref:`env_ids`. Shape + is ``(num_envs,)`` on the adapter device. Newton's :meth:`Actuator.State.reset` expects a per-DOF boolean mask of length ``num_actuators`` (= ``num_envs * dofs_per_actuator``), @@ -182,7 +191,7 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: etc.). We therefore build a per-actuator per-DOF mask from the env mask before delegating to each state. """ - if env_ids is None or env_ids == slice(None): + if env_mask is None and (env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None))): for sa, sb in zip(self._states_a, self._states_b): if sa is not None: sa.reset(None) @@ -190,19 +199,22 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: sb.reset(None) return - if isinstance(env_ids, torch.Tensor): - if env_ids.numel() == 0: - return - idx = wp.from_torch(env_ids.to(device=self._device).contiguous().to(torch.int32), dtype=wp.int32) - else: - if len(env_ids) == 0: - return - idx = wp.array(list(env_ids), dtype=wp.int32, device=self._device) - env_mask = wp.zeros(self._num_envs, dtype=wp.bool, device=self._device) - wp.launch(set_mask_kernel, dim=idx.shape[0], inputs=[env_mask, idx], device=self._device) - - for act, sa, sb in zip(self.actuators, self._states_a, self._states_b): - per_dof_mask = wp.zeros(act.indices.shape[0], dtype=wp.bool, device=self._device) + if env_mask is None: + env_mask = resolve_1d_mask( + ids=env_ids, + all_mask=self._reset_env_mask, + scratch_mask=self._reset_env_mask, + device=self._device, + ) + if ( + env_mask.dtype != wp.bool + or env_mask.ndim != 1 + or env_mask.shape[0] != self._num_envs + or env_mask.device != wp.get_device(self._device) + ): + raise ValueError(f"env_mask must have shape ({self._num_envs},) and dtype wp.bool, received {env_mask}.") + + for act, sa, sb, per_dof_mask in zip(self.actuators, self._states_a, self._states_b, self._reset_dof_masks): wp.launch( build_per_dof_env_mask_kernel, dim=act.indices.shape[0], diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index e0c26072c1da..ec66a907cc1a 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -270,9 +270,15 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None # use ellipses object to skip initial indices. if (env_ids is None) or (env_ids == slice(None)): env_ids = slice(None) + actuator_env_ids = env_ids + if self.actuators and env_mask is not None and not getattr(self, "_has_newton_actuators", False): + # Explicit boundary for the legacy Torch actuator fallback. The Warp + # frontend enables Newton-native actuators, whose state resets below + # consume the mask directly without materializing compact IDs. + actuator_env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) # reset Lab actuators registered on this articulation for actuator in self.actuators.values(): - actuator.reset(env_ids) + actuator.reset(actuator_env_ids) # reset the global Newton actuator adapter (its ``_states_a/_b`` buffers # carry per-env state — delay queues, neural hidden states — that must # be cleared for the resetting envs). The adapter spans the whole model, @@ -281,7 +287,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None # ``getattr`` guards subclasses (e.g. ``Multirotor``) that override # ``_process_actuators_cfg`` and never initialize ``_has_newton_actuators``. if getattr(self, "_has_newton_actuators", False) and SimulationManager._adapter is not None: - SimulationManager._adapter.reset(env_ids) + SimulationManager._adapter.reset(env_ids, env_mask=env_mask) # reset external wrenches. self._instantaneous_wrench_composer.reset(env_ids, env_mask) self._permanent_wrench_composer.reset(env_ids, env_mask) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py index 3265ad2896dd..3f6eac323887 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py @@ -136,8 +136,8 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None if (env_ids is None) or (env_ids == slice(None)): env_ids = slice(None) # reset external wrench - self._instantaneous_wrench_composer.reset(env_ids) - self._permanent_wrench_composer.reset(env_ids) + self._instantaneous_wrench_composer.reset(env_ids, env_mask) + self._permanent_wrench_composer.reset(env_ids, env_mask) def write_data_to_sim(self) -> None: """Write external wrench to the simulation. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py index 57512ce1b33e..168c3b9aa347 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py @@ -182,8 +182,8 @@ def reset( if object_ids is None: object_ids = self._ALL_BODY_INDICES # reset external wrench - self._instantaneous_wrench_composer.reset(env_ids) - self._permanent_wrench_composer.reset(env_ids) + self._instantaneous_wrench_composer.reset(env_ids, env_mask) + self._permanent_wrench_composer.reset(env_ids, env_mask) def write_data_to_sim(self) -> None: """Write external wrench to the simulation. diff --git a/source/isaaclab_newton/test/actuators/test_adapter_mask_reset.py b/source/isaaclab_newton/test/actuators/test_adapter_mask_reset.py new file mode 100644 index 000000000000..e33235ba28c7 --- /dev/null +++ b/source/isaaclab_newton/test/actuators/test_adapter_mask_reset.py @@ -0,0 +1,65 @@ +# 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 + +"""Tests for mask-based Newton actuator adapter resets.""" + +from __future__ import annotations + +from unittest.mock import Mock + +import numpy as np +import warp as wp +from isaaclab_newton.actuators.adapter import NewtonActuatorAdapter + + +def test_adapter_expands_environment_mask_without_allocating() -> None: + """Adapter state masks should select every actuator DOF in masked environments.""" + adapter = NewtonActuatorAdapter.__new__(NewtonActuatorAdapter) + actuator = Mock() + actuator.indices = wp.array([0, 1, 2, 3], dtype=wp.uint32, device="cpu") + state_a = Mock() + state_b = Mock() + adapter.actuators = [actuator] + adapter._states_a = [state_a] + adapter._states_b = [state_b] + adapter._num_envs = 2 + adapter._dof_offset = 0 + adapter.num_joints = 2 + adapter._device = "cpu" + adapter._reset_env_mask = wp.zeros(2, dtype=wp.bool, device="cpu") + adapter._reset_dof_masks = [wp.zeros(4, dtype=wp.bool, device="cpu")] + env_mask = wp.array([False, True], dtype=wp.bool, device="cpu") + + adapter.reset(env_mask=env_mask) + + state_a.reset.assert_called_once() + state_b.reset.assert_called_once() + np.testing.assert_array_equal(state_a.reset.call_args.args[0].numpy(), [False, False, True, True]) + assert state_a.reset.call_args.args[0] is adapter._reset_dof_masks[0] + + +def test_adapter_reuses_id_mask_without_stale_selection() -> None: + """Sequential ID resets should clear the reusable environment mask.""" + adapter = NewtonActuatorAdapter.__new__(NewtonActuatorAdapter) + actuator = Mock() + actuator.indices = wp.array([0, 1, 2, 3], dtype=wp.uint32, device="cpu") + state_a = Mock() + adapter.actuators = [actuator] + adapter._states_a = [state_a] + adapter._states_b = [None] + adapter._num_envs = 2 + adapter._dof_offset = 0 + adapter.num_joints = 2 + adapter._device = "cpu" + adapter._reset_env_mask = wp.zeros(2, dtype=wp.bool, device="cpu") + adapter._reset_dof_masks = [wp.zeros(4, dtype=wp.bool, device="cpu")] + + adapter.reset(env_ids=[0]) + first_mask = state_a.reset.call_args.args[0].numpy().copy() + adapter.reset(env_ids=[1]) + second_mask = state_a.reset.call_args.args[0].numpy().copy() + + np.testing.assert_array_equal(first_mask, [True, True, False, False]) + np.testing.assert_array_equal(second_mask, [False, False, True, True]) diff --git a/source/isaaclab_newton/test/assets/test_newton_mask_reset_forwarding.py b/source/isaaclab_newton/test/assets/test_newton_mask_reset_forwarding.py new file mode 100644 index 000000000000..414234293114 --- /dev/null +++ b/source/isaaclab_newton/test/assets/test_newton_mask_reset_forwarding.py @@ -0,0 +1,63 @@ +# 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 + +"""Tests for Warp reset-mask forwarding to Newton asset wrench storage.""" + +from unittest.mock import Mock + +import warp as wp +from isaaclab_newton.assets.articulation.articulation import Articulation +from isaaclab_newton.assets.rigid_object.rigid_object import RigidObject +from isaaclab_newton.assets.rigid_object_collection.rigid_object_collection import RigidObjectCollection + + +def test_rigid_object_forwards_reset_mask_to_wrench_composers() -> None: + """Rigid-object resets should preserve mask selection at wrench storage.""" + asset = RigidObject.__new__(RigidObject) + asset._initialize_handle = None + asset._invalidate_initialize_handle = None + asset._prim_deletion_handle = None + asset._instantaneous_wrench_composer = Mock() + asset._permanent_wrench_composer = Mock() + env_mask = wp.array([False, True], dtype=wp.bool, device="cpu") + + asset.reset(env_mask=env_mask) + + asset._instantaneous_wrench_composer.reset.assert_called_once_with(slice(None), env_mask) + asset._permanent_wrench_composer.reset.assert_called_once_with(slice(None), env_mask) + + +def test_rigid_object_collection_forwards_reset_mask_to_wrench_composers() -> None: + """Collection resets should preserve mask selection at wrench storage.""" + asset = RigidObjectCollection.__new__(RigidObjectCollection) + asset._initialize_handle = None + asset._invalidate_initialize_handle = None + asset._prim_deletion_handle = None + asset._instantaneous_wrench_composer = Mock() + asset._permanent_wrench_composer = Mock() + env_ids = [0] + env_mask = wp.array([True, False], dtype=wp.bool, device="cpu") + + asset.reset(env_ids=env_ids, object_ids=slice(None), env_mask=env_mask) + + asset._instantaneous_wrench_composer.reset.assert_called_once_with(env_ids, env_mask) + asset._permanent_wrench_composer.reset.assert_called_once_with(env_ids, env_mask) + + +def test_articulation_without_lab_actuators_keeps_mask_on_device() -> None: + """Mask resets should not materialize IDs when no Torch actuator consumes them.""" + asset = Articulation.__new__(Articulation) + asset._initialize_handle = None + asset._invalidate_initialize_handle = None + asset._prim_deletion_handle = None + asset.actuators = {} + asset._has_newton_actuators = False + asset._instantaneous_wrench_composer = Mock() + asset._permanent_wrench_composer = Mock() + env_mask = object() + + asset.reset(env_mask=env_mask) + + asset._instantaneous_wrench_composer.reset.assert_called_once_with(slice(None), env_mask) diff --git a/source/isaaclab_physx/changelog.d/jichuanh-warp-frontend-cleanup.rst b/source/isaaclab_physx/changelog.d/jichuanh-warp-frontend-cleanup.rst new file mode 100644 index 000000000000..c56453d2cc8e --- /dev/null +++ b/source/isaaclab_physx/changelog.d/jichuanh-warp-frontend-cleanup.rst @@ -0,0 +1,5 @@ +Fixed +^^^^^ + +* Fixed PhysX asset resets to forward boolean environment masks to Newton + actuator and external-wrench state. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 0bb66b68543e..6e43c71b78f9 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -236,15 +236,21 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None # use ellipses object to skip initial indices. if (env_ids is None) or (env_ids == slice(None)): env_ids = slice(None) - # reset actuators (including Newton-native adapter which owns its states) + actuator_env_ids = env_ids + if self.actuators and env_mask is not None and not getattr(self, "_has_newton_actuators", False): + # Explicit boundary for the legacy Torch actuator fallback. The Warp + # frontend enables Newton-native actuators, whose state resets below + # consume the mask directly without materializing compact IDs. + actuator_env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) + # reset Lab actuators for actuator in self.actuators.values(): - actuator.reset(env_ids) + actuator.reset(actuator_env_ids) # Reset Newton-actuator per-env states (delay queues, neural hidden state, etc.). # The adapter is per-articulation on PhysX and is not part of ``self.actuators``. # ``getattr`` guards subclasses (e.g. ``Multirotor``) that override # ``_process_actuators_cfg`` and never initialize these attributes. if getattr(self, "_has_newton_actuators", False) and getattr(self, "newton_actuator_adapter", None) is not None: - self.newton_actuator_adapter.reset(env_ids) + self.newton_actuator_adapter.reset(env_ids, env_mask=env_mask) # reset external wrenches. self._instantaneous_wrench_composer.reset(env_ids, env_mask) self._permanent_wrench_composer.reset(env_ids, env_mask) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py index 16c9801cb47c..4c42628313fd 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py @@ -141,8 +141,8 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None if (env_ids is None) or (env_ids == slice(None)): env_ids = slice(None) # reset external wrench - self._instantaneous_wrench_composer.reset(env_ids) - self._permanent_wrench_composer.reset(env_ids) + self._instantaneous_wrench_composer.reset(env_ids, env_mask) + self._permanent_wrench_composer.reset(env_ids, env_mask) def write_data_to_sim(self) -> None: """Write external wrench to the simulation. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py index d03bfb78012f..e2b36424774d 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py @@ -170,6 +170,8 @@ def reset( Args: env_ids: Environment indices. If None, then all indices are used. object_ids: Object indices. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + object_mask: Object mask. Not used currently. """ # resolve all indices if env_ids is None: @@ -177,8 +179,8 @@ def reset( if object_ids is None: object_ids = self._ALL_BODY_INDICES # reset external wrench - self._instantaneous_wrench_composer.reset(env_ids) - self._permanent_wrench_composer.reset(env_ids) + self._instantaneous_wrench_composer.reset(env_ids, env_mask) + self._permanent_wrench_composer.reset(env_ids, env_mask) def write_data_to_sim(self) -> None: """Write external wrench to the simulation. diff --git a/source/isaaclab_physx/test/assets/test_physx_mask_reset_forwarding.py b/source/isaaclab_physx/test/assets/test_physx_mask_reset_forwarding.py new file mode 100644 index 000000000000..63e616b11b76 --- /dev/null +++ b/source/isaaclab_physx/test/assets/test_physx_mask_reset_forwarding.py @@ -0,0 +1,70 @@ +# 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 + +"""Tests for Warp reset-mask forwarding to PhysX asset wrench storage.""" + +from isaaclab.app import AppLauncher +from isaaclab.test.utils import resolve_test_sim_device + +simulation_app = AppLauncher(headless=True, device=resolve_test_sim_device()).app + +"""Everything else follows.""" + +from unittest.mock import Mock + +import warp as wp +from isaaclab_physx.assets.articulation.articulation import Articulation +from isaaclab_physx.assets.rigid_object.rigid_object import RigidObject +from isaaclab_physx.assets.rigid_object_collection.rigid_object_collection import RigidObjectCollection + + +def test_rigid_object_forwards_reset_mask_to_wrench_composers() -> None: + """Rigid-object resets should preserve mask selection at wrench storage.""" + asset = RigidObject.__new__(RigidObject) + asset._initialize_handle = None + asset._invalidate_initialize_handle = None + asset._prim_deletion_handle = None + asset._instantaneous_wrench_composer = Mock() + asset._permanent_wrench_composer = Mock() + env_mask = wp.array([False, True], dtype=wp.bool, device="cpu") + + asset.reset(env_mask=env_mask) + + asset._instantaneous_wrench_composer.reset.assert_called_once_with(slice(None), env_mask) + asset._permanent_wrench_composer.reset.assert_called_once_with(slice(None), env_mask) + + +def test_rigid_object_collection_forwards_reset_mask_to_wrench_composers() -> None: + """Collection resets should preserve mask selection at wrench storage.""" + asset = RigidObjectCollection.__new__(RigidObjectCollection) + asset._initialize_handle = None + asset._invalidate_initialize_handle = None + asset._prim_deletion_handle = None + asset._instantaneous_wrench_composer = Mock() + asset._permanent_wrench_composer = Mock() + env_ids = [0] + env_mask = wp.array([True, False], dtype=wp.bool, device="cpu") + + asset.reset(env_ids=env_ids, object_ids=slice(None), env_mask=env_mask) + + asset._instantaneous_wrench_composer.reset.assert_called_once_with(env_ids, env_mask) + asset._permanent_wrench_composer.reset.assert_called_once_with(env_ids, env_mask) + + +def test_articulation_without_lab_actuators_keeps_mask_on_device() -> None: + """Mask resets should not materialize IDs when no Torch actuator consumes them.""" + asset = Articulation.__new__(Articulation) + asset._initialize_handle = None + asset._invalidate_initialize_handle = None + asset._prim_deletion_handle = None + asset.actuators = {} + asset._has_newton_actuators = False + asset._instantaneous_wrench_composer = Mock() + asset._permanent_wrench_composer = Mock() + env_mask = object() + + asset.reset(env_mask=env_mask) + + asset._instantaneous_wrench_composer.reset.assert_called_once_with(slice(None), env_mask) diff --git a/source/isaaclab_tasks/changelog.d/jichuanh-warp-frontend-cleanup.rst b/source/isaaclab_tasks/changelog.d/jichuanh-warp-frontend-cleanup.rst new file mode 100644 index 000000000000..7d2f89fbbf9d --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/jichuanh-warp-frontend-cleanup.rst @@ -0,0 +1,4 @@ +Fixed +^^^^^ + +* Fixed identity quaternion initialization for Dexsuite uniform pose commands. diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/commands/pose_commands.py b/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/commands/pose_commands.py index b0165e452ee3..c7014ba7b304 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/commands/pose_commands.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/commands/pose_commands.py @@ -73,7 +73,7 @@ def __init__(self, cfg: dex_cmd_cfgs.ObjectUniformPoseCommandCfg, env: ManagerBa # create buffers # -- commands: (x, y, z, qx, qy, qz, qw) in root frame self.pose_command_b = torch.zeros(self.num_envs, 7, device=self.device) - self.pose_command_b[:, 3] = 1.0 + self.pose_command_b[:, 6] = 1.0 self.pose_command_w = torch.zeros_like(self.pose_command_b) # -- metrics self.metrics["position_error"] = torch.zeros(self.num_envs, device=self.device) diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst new file mode 100644 index 000000000000..cd2905af15b2 --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -0,0 +1,5 @@ +Added +^^^^^ + +* Added Warp-native terrain and reward-weight curricula and connected locomotion + and reach hot paths to pointer-stable Warp command arrays. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py index 2082d470ff9b..d2f1be03dd17 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py @@ -667,6 +667,7 @@ def __init__(self, cfg: AllegroHandWarpEnvCfg, render_mode: str | None = None, * self.torch_reset_terminated = wp.to_torch(self.reset_terminated) self.torch_reset_time_outs = wp.to_torch(self.reset_time_outs) self.torch_episode_length_buf = self.episode_length_buf # already a torch tensor via wp.to_torch + self.torch_consecutive_successes = wp.to_torch(self.consecutive_successes) def _setup_scene(self): # add hand, in-hand object, and goal object @@ -771,7 +772,7 @@ def _get_rewards(self) -> None: if "log" not in self.extras: self.extras["log"] = dict() # .mean() cannot be called here as it causes problems on stream - self.extras["log"]["consecutive_successes"] = wp.to_torch(self.consecutive_successes) + self.extras["log"]["consecutive_successes"] = self.torch_consecutive_successes # Reset goals for envs that reached the target (mask is `reset_goal_buf`). # This avoids Torch-side index extraction and keeps the step graphable. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.pyi index ae4b6d7af547..463c665cd67a 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.pyi +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.pyi @@ -5,6 +5,7 @@ __all__ = [ "terrain_levels_vel", + "TerrainLevelsVel", "feet_air_time", "feet_air_time_positive_biped", "feet_slide", @@ -18,7 +19,7 @@ __all__ = [ # observations, ...) lazily, then override with locomotion-specific terms below. from isaaclab_experimental.envs.mdp import * # noqa: F401, F403 -from .curriculums import terrain_levels_vel +from .curriculums import TerrainLevelsVel, terrain_levels_vel from .rewards import ( feet_air_time, feet_air_time_positive_biped, diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py index c1aa55114ab7..f55af9254efa 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py @@ -3,11 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Curriculum functions for the velocity locomotion environment. - -Curriculum terms are not warp-managed (they run at reset time, not per-step), -so they remain torch-based. -""" +"""Warp-first curriculum terms for velocity locomotion environments.""" from __future__ import annotations @@ -15,8 +11,8 @@ from typing import TYPE_CHECKING import torch - -from isaaclab.managers import SceneEntityCfg +import warp as wp +from isaaclab_experimental.managers import CurriculumTermCfg, ManagerTermBase, SceneEntityCfg if TYPE_CHECKING: from isaaclab.assets import Articulation @@ -24,10 +20,119 @@ from isaaclab.terrains import TerrainImporter +@wp.kernel +def _compute_terrain_level_updates( + env_mask: wp.array(dtype=wp.bool), + root_pos_w: wp.array(dtype=wp.vec3f), + env_origins: wp.array(dtype=wp.vec3f), + command: wp.array(dtype=wp.float32, ndim=2), + terrain_half_length: float, + max_episode_length_s: float, + move_up: wp.array(dtype=wp.bool), + move_down: wp.array(dtype=wp.bool), +): + env_id = wp.tid() + if env_mask[env_id]: + position_delta = root_pos_w[env_id] - env_origins[env_id] + distance = wp.sqrt(position_delta[0] * position_delta[0] + position_delta[1] * position_delta[1]) + command_distance = ( + wp.sqrt(command[env_id, 0] * command[env_id, 0] + command[env_id, 1] * command[env_id, 1]) + * max_episode_length_s + * 0.5 + ) + should_move_up = distance > terrain_half_length + move_up[env_id] = should_move_up + move_down[env_id] = (distance < command_distance) and (not should_move_up) + else: + move_up[env_id] = False + move_down[env_id] = False + + +@wp.kernel +def _compute_terrain_level_mean( + terrain_levels: wp.array(dtype=wp.int64), + scale: float, + out: wp.array(dtype=wp.float32), +): + env_id = wp.tid() + wp.atomic_add(out, 0, wp.float32(terrain_levels[env_id]) * scale) + + +class TerrainLevelsVel(ManagerTermBase): + """Update terrain levels from commanded locomotion progress using a boolean environment mask.""" + + _curriculum_mask_native = True + + def __init__(self, cfg: CurriculumTermCfg, env: ManagerBasedRLEnv): + """Initialize persistent views and decision buffers. + + Args: + cfg: Curriculum term configuration. + env: Environment containing the robot and terrain. + """ + super().__init__(cfg, env) + asset_cfg = cfg.params.get("asset_cfg", SceneEntityCfg("robot")) + self._asset: Articulation = env.scene[asset_cfg.name] + self._terrain: TerrainImporter = env.scene.terrain + self._root_pos_w_wp = self._asset.data.root_pos_w.warp + self._env_origins_wp = env.env_origins_wp + self._command_wp = env.command_manager.get_command_wp("base_velocity") + self._terrain_levels_wp = self._terrain.terrain_levels_wp + self._move_up_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + self._move_down_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + self._terrain_half_length = float(self._terrain.cfg.terrain_generator.size[0]) * 0.5 + self._max_episode_length_s = float(env.max_episode_length_s) + self._mean_scale = 1.0 / self.num_envs + + def __call__( + self, + env: ManagerBasedRLEnv, + env_mask: wp.array(dtype=wp.bool), + out: wp.array(dtype=wp.float32), + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + ) -> None: + """Update selected terrain levels and write the global mean level. + + Args: + env: Environment containing the per-environment random state. + env_mask: Boolean Warp mask selecting environments to update. + out: Persistent scalar output for the mean terrain level. + asset_cfg: Robot scene entity configuration. + """ + del asset_cfg + wp.launch( + kernel=_compute_terrain_level_updates, + dim=self.num_envs, + inputs=[ + env_mask, + self._root_pos_w_wp, + self._env_origins_wp, + self._command_wp, + self._terrain_half_length, + self._max_episode_length_s, + self._move_up_wp, + self._move_down_wp, + ], + device=self.device, + ) + self._terrain.update_env_origins_mask( + env_mask, + self._move_up_wp, + self._move_down_wp, + env.rng_state_wp, + ) + wp.launch( + kernel=_compute_terrain_level_mean, + dim=self.num_envs, + inputs=[self._terrain_levels_wp, self._mean_scale, out], + device=self.device, + ) + + def terrain_levels_vel( env: ManagerBasedRLEnv, env_ids: Sequence[int], asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> torch.Tensor: - """Curriculum based on the distance the robot walked when commanded to move at a desired velocity.""" + """Update terrain levels through the legacy ID-based public API.""" asset: Articulation = env.scene[asset_cfg.name] terrain: TerrainImporter = env.scene.terrain command = env.command_manager.get_command("base_velocity") diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py index d7e11c517c7d..1b1744c4078b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py @@ -6,8 +6,6 @@ """Warp-first reward functions for the velocity locomotion environment. All functions follow the ``func(env, out, **params) -> None`` signature. -Cross-manager torch tensors (contact sensor, commands) are cached as zero-copy -warp views on first call via ``wp.from_torch``. """ from __future__ import annotations @@ -50,14 +48,7 @@ def _feet_air_time_kernel( def feet_air_time(env: ManagerBasedRLEnv, out, command_name: str, sensor_cfg: SceneEntityCfg, threshold: float) -> None: """Reward long steps taken by the feet using L2-kernel.""" - fn = feet_air_time contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] - # Cache command bridge (persistent pointer) - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True # Newton contact sensor returns persistent wp.arrays — use directly, no wp.from_torch needed first_contact = contact_sensor.compute_first_contact(env.step_dt) wp.launch( @@ -67,7 +58,7 @@ def feet_air_time(env: ManagerBasedRLEnv, out, command_name: str, sensor_cfg: Sc contact_sensor.data.last_air_time.warp, first_contact, sensor_cfg.body_ids_wp, - fn._cmd_wp, + env.command_manager.get_command_wp(command_name), threshold, out, ], @@ -114,13 +105,7 @@ def _feet_air_time_positive_biped_kernel( def feet_air_time_positive_biped(env, out, command_name: str, threshold: float, sensor_cfg: SceneEntityCfg) -> None: """Reward long steps taken by the feet for bipeds.""" - fn = feet_air_time_positive_biped contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True wp.launch( kernel=_feet_air_time_positive_biped_kernel, dim=env.num_envs, @@ -128,7 +113,7 @@ def feet_air_time_positive_biped(env, out, command_name: str, threshold: float, contact_sensor.data.current_air_time.warp, contact_sensor.data.current_contact_time.warp, sensor_cfg.body_ids_wp, - fn._cmd_wp, + env.command_manager.get_command_wp(command_name), threshold, out, ], @@ -224,17 +209,17 @@ def track_lin_vel_xy_yaw_frame_exp( env, out, std: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> None: """Reward tracking of linear velocity commands (xy axes) in the gravity aligned robot frame.""" - fn = track_lin_vel_xy_yaw_frame_exp asset = env.scene[asset_cfg.name] - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True wp.launch( kernel=_track_lin_vel_xy_yaw_frame_exp_kernel, dim=env.num_envs, - inputs=[asset.data.root_quat_w.warp, asset.data.root_lin_vel_w.warp, fn._cmd_wp, 1.0 / (std * std), out], + inputs=[ + asset.data.root_quat_w.warp, + asset.data.root_lin_vel_w.warp, + env.command_manager.get_command_wp(command_name), + 1.0 / (std * std), + out, + ], device=env.device, ) @@ -260,17 +245,16 @@ def track_ang_vel_z_world_exp( env, out, command_name: str, std: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> None: """Reward tracking of angular velocity commands (yaw) in world frame.""" - fn = track_ang_vel_z_world_exp asset = env.scene[asset_cfg.name] - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True wp.launch( kernel=_track_ang_vel_z_world_exp_kernel, dim=env.num_envs, - inputs=[asset.data.root_ang_vel_w.warp, fn._cmd_wp, 1.0 / (std * std), out], + inputs=[ + asset.data.root_ang_vel_w.warp, + env.command_manager.get_command_wp(command_name), + 1.0 / (std * std), + out, + ], device=env.device, ) @@ -304,16 +288,16 @@ def stand_still_joint_deviation_l1( env, out, command_name: str, command_threshold: float = 0.06, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> None: """Penalize offsets from the default joint positions when the command is very small.""" - fn = stand_still_joint_deviation_l1 asset = env.scene[asset_cfg.name] - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True wp.launch( kernel=_stand_still_joint_deviation_l1_kernel, dim=env.num_envs, - inputs=[asset.data.joint_pos.warp, asset.data.default_joint_pos.warp, fn._cmd_wp, command_threshold, out], + inputs=[ + asset.data.joint_pos.warp, + asset.data.default_joint_pos.warp, + env.command_manager.get_command_wp(command_name), + command_threshold, + out, + ], device=env.device, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py index 52d53c3db268..2807bf2c5a31 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py @@ -6,6 +6,7 @@ import math from dataclasses import MISSING +from isaaclab_experimental.managers import CurriculumTermCfg as CurrTerm from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm from isaaclab_experimental.managers import RewardTermCfg as RewTerm from isaaclab_experimental.managers import SceneEntityCfg @@ -14,7 +15,6 @@ import isaaclab.sim as sim_utils from isaaclab.assets import ArticulationCfg, AssetBaseCfg from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import CurriculumTermCfg as CurrTerm from isaaclab.managers import EventTermCfg as EventTerm from isaaclab.managers import ObservationGroupCfg as ObsGroup from isaaclab.scene import InteractiveSceneCfg @@ -251,7 +251,7 @@ class TerminationsCfg: class CurriculumCfg: """Curriculum terms for the MDP.""" - terrain_levels = CurrTerm(func=mdp.terrain_levels_vel) + terrain_levels = CurrTerm(func=mdp.TerrainLevelsVel) ## diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.pyi index c277ec98e076..1384097609f8 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.pyi +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.pyi @@ -4,6 +4,7 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "ModifyRewardWeight", "position_command_error", "position_command_error_tanh", "orientation_command_error", @@ -13,4 +14,5 @@ __all__ = [ # override with reach-specific terms below. from isaaclab_experimental.envs.mdp import * # noqa: F401, F403 +from .curriculums import ModifyRewardWeight from .rewards import orientation_command_error, position_command_error, position_command_error_tanh diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/curriculums.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/curriculums.py new file mode 100644 index 000000000000..66bb0b0eabfb --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/curriculums.py @@ -0,0 +1,98 @@ +# 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 + +"""Warp-first curriculum terms for reach environments.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import warp as wp +from isaaclab_experimental.managers import CurriculumTermCfg, ManagerTermBase + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +@wp.kernel +def _mark_reset_requested(env_mask: wp.array(dtype=wp.bool), reset_requested: wp.array(dtype=wp.int32)): + env_id = wp.tid() + if env_mask[env_id]: + wp.atomic_max(reset_requested, 0, 1) + + +@wp.kernel +def _update_reward_weight( + reset_requested: wp.array(dtype=wp.int32), + common_step_counter: int, + num_steps: int, + target_weight: float, + term_weight: wp.array(dtype=wp.float32), + out: wp.array(dtype=wp.float32), +): + if reset_requested[0] != 0: + if common_step_counter > num_steps: + term_weight[0] = target_weight + out[0] = term_weight[0] + reset_requested[0] = 0 + + +class ModifyRewardWeight(ManagerTermBase): + """Update one reward weight after a reset reaches the configured step threshold.""" + + # The Python-owned common step counter is intentionally read in eager mode. + # A future captured path should first move this counter to persistent device storage. + _warp_capturable = False + + def __init__(self, cfg: CurriculumTermCfg, env: ManagerBasedRLEnv): + """Initialize persistent device state. + + Args: + cfg: Curriculum term configuration. + env: Environment containing the reward manager. + """ + super().__init__(cfg, env) + self._term_weight_wp = env.reward_manager.get_term_weight_wp(cfg.params["term_name"]) + self._reset_requested_wp = wp.zeros(1, dtype=wp.int32, device=self.device) + + def __call__( + self, + env: ManagerBasedRLEnv, + env_mask: wp.array(dtype=wp.bool), + out: wp.array(dtype=wp.float32), + term_name: str, + weight: float, + num_steps: int, + ) -> None: + """Update the reward weight when at least one environment resets. + + Args: + env: Environment containing the common step counter. + env_mask: Boolean mask selecting resetting environments. + out: Persistent one-element output containing the current weight. + term_name: Reward term name, resolved during initialization. + weight: Target reward weight. + num_steps: Step threshold after which the target weight is applied. + """ + del term_name + wp.launch( + _mark_reset_requested, + dim=self.num_envs, + inputs=[env_mask, self._reset_requested_wp], + device=self.device, + ) + wp.launch( + _update_reward_weight, + dim=1, + inputs=[ + self._reset_requested_wp, + env.common_step_counter, + num_steps, + weight, + self._term_weight_wp, + out, + ], + device=self.device, + ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py index 7435752528d6..51a828177a42 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py @@ -6,7 +6,6 @@ """Warp-first reward terms for the reach task. All functions follow the ``func(env, out, **params) -> None`` signature. -Command tensors are cached as zero-copy warp views on first call. """ from __future__ import annotations @@ -50,12 +49,7 @@ def _position_command_error_kernel( def position_command_error(env: ManagerBasedRLEnv, out, command_name: str, asset_cfg: SceneEntityCfg) -> None: """Penalize tracking of the position error using L2-norm.""" - fn = position_command_error asset: Articulation = env.scene[asset_cfg.name] - if not hasattr(fn, "_cmd_wp") or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name wp.launch( kernel=_position_command_error_kernel, dim=env.num_envs, @@ -63,7 +57,7 @@ def position_command_error(env: ManagerBasedRLEnv, out, command_name: str, asset asset.data.root_pos_w.warp, asset.data.root_quat_w.warp, asset.data.body_pos_w.warp, - fn._cmd_wp, + env.command_manager.get_command_wp(command_name), asset_cfg.body_ids[0], out, ], @@ -101,12 +95,7 @@ def position_command_error_tanh( env: ManagerBasedRLEnv, out, std: float, command_name: str, asset_cfg: SceneEntityCfg ) -> None: """Reward tracking of the position using the tanh kernel.""" - fn = position_command_error_tanh asset: Articulation = env.scene[asset_cfg.name] - if not hasattr(fn, "_cmd_wp") or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name wp.launch( kernel=_position_command_error_tanh_kernel, dim=env.num_envs, @@ -114,7 +103,7 @@ def position_command_error_tanh( asset.data.root_pos_w.warp, asset.data.root_quat_w.warp, asset.data.body_pos_w.warp, - fn._cmd_wp, + env.command_manager.get_command_wp(command_name), asset_cfg.body_ids[0], 1.0 / std, out, @@ -152,15 +141,16 @@ def _orientation_command_error_kernel( def orientation_command_error(env: ManagerBasedRLEnv, out, command_name: str, asset_cfg: SceneEntityCfg) -> None: """Penalize tracking orientation error using shortest path.""" - fn = orientation_command_error asset: Articulation = env.scene[asset_cfg.name] - if not hasattr(fn, "_cmd_wp") or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name wp.launch( kernel=_orientation_command_error_kernel, dim=env.num_envs, - inputs=[asset.data.root_quat_w.warp, asset.data.body_quat_w.warp, fn._cmd_wp, asset_cfg.body_ids[0], out], + inputs=[ + asset.data.root_quat_w.warp, + asset.data.body_quat_w.warp, + env.command_manager.get_command_wp(command_name), + asset_cfg.body_ids[0], + out, + ], device=env.device, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py index e83a27e44db5..ebb94cab0c5c 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py @@ -5,6 +5,7 @@ from dataclasses import MISSING +from isaaclab_experimental.managers import CurriculumTermCfg as CurrTerm from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm from isaaclab_experimental.managers import RewardTermCfg as RewTerm from isaaclab_experimental.managers import SceneEntityCfg @@ -14,7 +15,6 @@ from isaaclab.assets import ArticulationCfg, AssetBaseCfg from isaaclab.envs import ManagerBasedRLEnvCfg from isaaclab.managers import ActionTermCfg as ActionTerm -from isaaclab.managers import CurriculumTermCfg as CurrTerm from isaaclab.managers import EventTermCfg as EventTerm from isaaclab.managers import ObservationGroupCfg as ObsGroup from isaaclab.scene import InteractiveSceneCfg @@ -166,11 +166,13 @@ class CurriculumCfg: """Curriculum terms for the MDP.""" action_rate = CurrTerm( - func=mdp.modify_reward_weight, params={"term_name": "action_rate", "weight": -0.005, "num_steps": 4500} + func=mdp.ModifyRewardWeight, + params={"term_name": "action_rate", "weight": -0.005, "num_steps": 4500}, ) joint_vel = CurrTerm( - func=mdp.modify_reward_weight, params={"term_name": "joint_vel", "weight": -0.001, "num_steps": 4500} + func=mdp.ModifyRewardWeight, + params={"term_name": "joint_vel", "weight": -0.001, "num_steps": 4500}, ) From f8180c9ba50caaefcf83eb0c7feaef0a6aab829e Mon Sep 17 00:00:00 2001 From: jichuanh Date: Sun, 19 Jul 2026 03:35:53 -0700 Subject: [PATCH 37/58] Keep Newton actuators opt-in Preserve each task's configured actuator path instead of forcing the native Newton adapter in every Warp environment. The native path currently makes representative Go2 physics substeps roughly ten times slower, so it remains available for targeted validation without regressing the default frontend. --- .../isaaclab_experimental/envs/direct_rl_env_warp.py | 4 ---- .../isaaclab_experimental/envs/manager_based_env_warp.py | 4 ---- 2 files changed, 8 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py index dcd6fdf29a38..02eaff52ff4e 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py @@ -102,10 +102,6 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs RuntimeError: If a simulation context already exists. The environment must always create one since it configures the simulation context and controls the simulation. """ - # Keep the Warp frontend on the Warp-native actuator path. This covers - # implicit, PD, delayed, DC, and neural actuator configurations without - # crossing through the Torch-based Isaac Lab actuator loop each step. - cfg.sim.use_newton_actuators = True # check that the config is valid cfg.validate() # store inputs to class diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py index f6979f7052de..03295b324010 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py @@ -73,10 +73,6 @@ def __init__(self, cfg: ManagerBasedEnvCfg): RuntimeError: If a simulation context already exists. The environment must always create one since it configures the simulation context and controls the simulation. """ - # Keep the Warp frontend on the Warp-native actuator path. This covers - # implicit, PD, delayed, DC, and neural actuator configurations without - # crossing through the Torch-based Isaac Lab actuator loop each step. - cfg.sim.use_newton_actuators = True # check that the config is valid cfg.validate() # store inputs to class From ffa8a73b60339eb8016f5e7b1e2b092ab2273f6a Mon Sep 17 00:00:00 2001 From: jichuanh Date: Sun, 19 Jul 2026 03:41:12 -0700 Subject: [PATCH 38/58] Skip empty eager reset pipelines Keep reset masks canonical while using one explicit host predicate to avoid dispatching every reset manager for an empty selection. This recovers the eager hot path until recorded reset launches can remove the boundary. --- .../envs/manager_based_rl_env_warp.py | 6 ++++++ .../test/envs/test_manager_based_rl_env_warp.py | 16 ++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index 1c3632ef88db..259f9ac66bd3 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -580,6 +580,12 @@ def _reset_idx( def _reset_terminated_envs(self) -> None: """Reset terminated environments, compacting IDs only for host consumers.""" reset_mask = self.termination_manager.dones_wp + # The eager reset pipeline contains many small launches. Keep the mask as + # the canonical selection, but use one explicit host predicate to avoid + # dispatching the entire pipeline when it is empty. Record/replay can + # remove this boundary once reset stages have replay-safe launch caches. + if not self.reset_buf.any().item(): + return reset_env_ids = None if self._reset_requires_host_selection(): with Timer( diff --git a/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py index 00d4c09222be..7fbce5fe4899 100644 --- a/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py @@ -190,3 +190,19 @@ def test_empty_reset_skips_legacy_host_boundary_and_preserves_logs() -> None: env._reset_host_pre.assert_not_called() env._reset_host_post.assert_not_called() assert env.extras["log"] == {"previous": 1.0} + + +def test_empty_reset_skips_warp_pipeline_without_compacting_ids() -> None: + """An empty mask should skip eager reset stages without materializing IDs.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + reset_mask = wp.zeros(3, dtype=wp.bool, device="cpu") + env.termination_manager = SimpleNamespace(dones_wp=reset_mask) + env.reset_buf = torch.zeros(3, dtype=torch.bool) + env._reset_requires_host_selection = Mock(return_value=False) + env._reset_idx = Mock() + env.extras = {"log": {"previous": 1.0}} + + env._reset_terminated_envs() + + env._reset_idx.assert_not_called() + assert env.extras["log"] == {"previous": 1.0} From b3707f484c607390b2cd698b41c8e7ac34971eab Mon Sep 17 00:00:00 2001 From: jichuanh Date: Sun, 19 Jul 2026 03:45:07 -0700 Subject: [PATCH 39/58] Clarify reset optimization stages --- .../isaaclab_experimental/envs/manager_based_rl_env_warp.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index 259f9ac66bd3..e62112b5ff58 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -582,8 +582,8 @@ def _reset_terminated_envs(self) -> None: reset_mask = self.termination_manager.dones_wp # The eager reset pipeline contains many small launches. Keep the mask as # the canonical selection, but use one explicit host predicate to avoid - # dispatching the entire pipeline when it is empty. Record/replay can - # remove this boundary once reset stages have replay-safe launch caches. + # dispatching the entire pipeline when it is empty. Recorded launches + # can reduce that cost; capture can later make this predicate unnecessary. if not self.reset_buf.any().item(): return reset_env_ids = None From 89a25f1111b8c1bb69ca9ad461d6fe800180b480 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Mon, 20 Jul 2026 00:32:15 -0700 Subject: [PATCH 40/58] Refine Warp frontend reset paths Keep Warp-native reset and manager paths mask-based through their host boundaries, and remove unsupported PhysX forwarding. Move frontend-only helpers into the experimental package, make terrain origin buffers dual-access, and keep capture state scoped to its owning environment. --- docs/source/how-to/proxy_array.rst | 13 +- .../migration/migrating_to_isaaclab_3-0.rst | 5 + scripts/demos/procedural_terrain.py | 2 +- scripts/tutorials/06_deploy/anymal_c_env.py | 2 +- .../jichuanh-warp-frontend-cleanup.minor.rst | 17 +- .../envs/mdp/commands/pose_2d_command.py | 2 +- .../envs/mdp/commands/pose_command.py | 35 +- .../envs/mdp/commands/velocity_command.py | 62 +- source/isaaclab/isaaclab/envs/mdp/events.py | 2 +- .../isaaclab/scene/interactive_scene.py | 19 +- .../isaaclab/terrains/terrain_importer.py | 165 +++-- .../test/terrains/check_terrain_importer.py | 2 +- .../test/terrains/test_terrain_importer.py | 13 +- .../jichuanh-warp-frontend-cleanup.minor.rst | 12 +- .../envs/direct_rl_env_warp.py | 17 +- .../envs/manager_based_env_warp.py | 114 +--- .../envs/manager_based_rl_env_warp.py | 33 +- .../envs/mdp/commands/_debug_vis.py | 6 +- .../envs/mdp/commands/pose_command.py | 3 +- .../envs/mdp/commands/velocity_command.py | 29 +- .../isaaclab_experimental/envs/mdp/events.py | 644 ++++++++++-------- .../managers/curriculum_manager.py | 42 +- .../managers/manager_base.py | 9 +- .../utils/manager_call_switch.py | 180 +++-- .../utils/warp_graph_cache.py | 91 ++- .../test/envs/mdp/commands/test_commands.py | 54 +- .../test/envs/mdp/parity_helpers.py | 7 +- .../test/envs/mdp/test_events_warp_parity.py | 175 ++++- .../test/envs/test_interactive_scene_warp.py | 30 + .../envs/test_manager_based_rl_env_warp.py | 55 +- .../test/managers/test_curriculum_manager.py | 165 +++-- .../test/utils/test_warp_graph_cache.py | 21 + .../jichuanh-warp-frontend-cleanup.rst | 5 - .../assets/articulation/articulation.py | 12 +- .../assets/rigid_object/rigid_object.py | 4 +- .../rigid_object_collection.py | 6 +- .../test_physx_mask_reset_forwarding.py | 70 -- .../jichuanh-warp-frontend-cleanup.rst | 3 +- .../contrib/anymal_c_direct/anymal_c_env.py | 2 +- .../dexsuite/mdp/commands/pose_commands.py | 2 +- .../core/dexsuite/mdp/events.py | 2 +- .../core/velocity/mdp/curriculums.py | 2 +- .../jichuanh-warp-frontend-cleanup.minor.rst | 6 + .../inhand_manipulation_warp_env.py | 3 +- .../direct/locomotion/locomotion_env_warp.py | 2 +- .../manager_based/classic/ant/ant_env_cfg.py | 2 +- .../classic/humanoid/humanoid_env_cfg.py | 2 +- .../locomotion/velocity/mdp/curriculums.py | 4 +- .../locomotion/velocity/mdp/terminations.py | 31 +- .../locomotion/velocity/velocity_env_cfg.py | 8 +- .../manipulation/reach/reach_env_cfg.py | 2 +- .../test/test_terrain_out_of_bounds.py | 41 ++ 52 files changed, 1357 insertions(+), 878 deletions(-) rename source/{isaaclab/isaaclab => isaaclab_experimental/isaaclab_experimental}/envs/mdp/commands/_debug_vis.py (94%) delete mode 100644 source/isaaclab_physx/changelog.d/jichuanh-warp-frontend-cleanup.rst delete mode 100644 source/isaaclab_physx/test/assets/test_physx_mask_reset_forwarding.py create mode 100644 source/isaaclab_tasks_experimental/test/test_terrain_out_of_bounds.py diff --git a/docs/source/how-to/proxy_array.rst b/docs/source/how-to/proxy_array.rst index 7bf0a90374f3..48459e1455dd 100644 --- a/docs/source/how-to/proxy_array.rst +++ b/docs/source/how-to/proxy_array.rst @@ -5,9 +5,9 @@ Working with ProxyArray .. currentmodule:: isaaclab.utils.warp -Isaac Lab data classes return :class:`ProxyArray` — a lightweight, warp-first wrapper that -provides zero-copy access to simulation data as either a :class:`warp.array` or a -:class:`torch.Tensor`. +Isaac Lab data classes and :class:`~isaaclab.terrains.TerrainImporter` buffers return +:class:`ProxyArray` — a lightweight, warp-first wrapper that provides zero-copy access as either a +:class:`warp.array` or a :class:`torch.Tensor`. .. note:: @@ -20,7 +20,9 @@ Quick Start ~~~~~~~~~~~ Every property on asset and sensor data classes (e.g., ``robot.data.joint_pos``, -``sensor.data.net_forces_w``) returns a ``ProxyArray``: +``sensor.data.net_forces_w``) returns a ``ProxyArray``. The ``terrain_origins``, +``env_origins``, ``terrain_levels``, and ``terrain_types`` buffers on ``TerrainImporter`` follow +the same interface: .. code-block:: python @@ -36,6 +38,9 @@ Every property on asset and sensor data classes (e.g., ``robot.data.joint_pos``, # Pass directly to warp kernels — no unwrapping needed wp.launch(my_kernel, inputs=[robot.data.joint_pos], ...) # works via __cuda_array_interface__ + # Terrain buffers use the same explicit accessors + terrain_levels = env.scene.terrain.terrain_levels.torch + The ``.torch`` and ``.warp`` Accessors ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/docs/source/migration/migrating_to_isaaclab_3-0.rst b/docs/source/migration/migrating_to_isaaclab_3-0.rst index 0f9580f8fe3a..566781b413d9 100644 --- a/docs/source/migration/migrating_to_isaaclab_3-0.rst +++ b/docs/source/migration/migrating_to_isaaclab_3-0.rst @@ -1269,6 +1269,11 @@ change applies to all asset classes (:class:`~isaaclab.assets.Articulation`, (:class:`~isaaclab_physx.sensors.ContactSensor`, :class:`~isaaclab_physx.sensors.Imu`, :class:`~isaaclab_physx.sensors.Pva`, :class:`~isaaclab_physx.sensors.FrameTransformer`). +The terrain-origin fields on :class:`~isaaclab.terrains.TerrainImporter` follow the same +interface. :attr:`~isaaclab.scene.InteractiveScene.env_origins` remains a ``torch.Tensor`` +compatibility boundary; use :attr:`~isaaclab.scene.InteractiveScene.env_origins_wp` for its +Warp view. + To use a data property as a ``torch.Tensor``, append ``.torch``: .. code-block:: python diff --git a/scripts/demos/procedural_terrain.py b/scripts/demos/procedural_terrain.py index 6e26a8b26536..6dd0c2f57bf6 100644 --- a/scripts/demos/procedural_terrain.py +++ b/scripts/demos/procedural_terrain.py @@ -145,7 +145,7 @@ def design_scene() -> tuple[dict, torch.Tensor]: # return the scene information scene_entities = {"terrain": terrain_importer} - return scene_entities, terrain_importer.env_origins + return scene_entities, terrain_importer.env_origins.torch def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, AssetBase], origins: torch.Tensor): diff --git a/scripts/tutorials/06_deploy/anymal_c_env.py b/scripts/tutorials/06_deploy/anymal_c_env.py index 2f0003fcfe23..25194594c4ad 100644 --- a/scripts/tutorials/06_deploy/anymal_c_env.py +++ b/scripts/tutorials/06_deploy/anymal_c_env.py @@ -195,7 +195,7 @@ def _reset_idx(self, env_ids: torch.Tensor | None): joint_vel = self._robot.data.default_joint_vel.torch[env_ids] default_root_pose = self._robot.data.default_root_pose.torch[env_ids] default_root_vel = self._robot.data.default_root_vel.torch[env_ids] - default_root_pose[:, :3] += self._terrain.env_origins[env_ids] + default_root_pose[:, :3] += self._terrain.env_origins.torch[env_ids] self._robot.write_root_pose_to_sim_index(root_pose=default_root_pose, env_ids=env_ids) self._robot.write_root_velocity_to_sim_index(root_velocity=default_root_vel, env_ids=env_ids) self._robot.write_joint_position_to_sim_index(position=joint_pos, env_ids=env_ids) diff --git a/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst index e4bb7eec7da8..dbf6d1927d0b 100644 --- a/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst +++ b/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -1,9 +1,20 @@ Added ^^^^^ -* Added :attr:`~isaaclab.terrains.TerrainImporter.terrain_levels_wp` and - :meth:`~isaaclab.terrains.TerrainImporter.update_env_origins_mask` for - mask-based Warp terrain curriculum updates. +* Added :attr:`~isaaclab.scene.InteractiveScene.env_origins_wp` for zero-copy Warp access to + scene origins. +* Added :meth:`~isaaclab.terrains.TerrainImporter.update_env_origins_mask` for mask-based Warp + terrain curriculum updates. + +Changed +^^^^^^^ + +* Changed :attr:`~isaaclab.terrains.TerrainImporter.terrain_origins`, + :attr:`~isaaclab.terrains.TerrainImporter.env_origins`, + :attr:`~isaaclab.terrains.TerrainImporter.terrain_levels`, and + :attr:`~isaaclab.terrains.TerrainImporter.terrain_types` to return + :class:`~isaaclab.utils.warp.ProxyArray`. Use ``.torch`` for Torch operations + and ``.warp`` for Warp kernels. Fixed ^^^^^ diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py index ffd7bdad5e49..9d3a3803925a 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py @@ -204,7 +204,7 @@ def _resample_command(self, env_ids: Sequence[int]): # sample new position targets from the terrain ids = torch.randint(0, self.valid_targets.shape[2], size=(len(env_ids),), device=self.device) self.pos_command_w[env_ids] = self.valid_targets[ - self.terrain.terrain_levels[env_ids], self.terrain.terrain_types[env_ids], ids + self.terrain.terrain_levels.torch[env_ids], self.terrain.terrain_types.torch[env_ids], ids ] # offset the position command by the current root height self.pos_command_w[env_ids, 2] += self.robot.data.default_root_pose.torch[env_ids, 2] diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py index 876b5fb9ee30..d98a6b5e7f81 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py @@ -14,18 +14,17 @@ from isaaclab.assets import Articulation from isaaclab.managers import CommandTerm +from isaaclab.markers import VisualizationMarkers from isaaclab.utils.leapp import POSE7_ELEMENT_NAMES from isaaclab.utils.math import combine_frame_transforms, compute_pose_error, quat_from_euler_xyz, quat_unique -from ._debug_vis import _PoseCommandDebugVis - if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv from .commands_cfg import UniformPoseCommandCfg -class UniformPoseCommand(_PoseCommandDebugVis, CommandTerm): +class UniformPoseCommand(CommandTerm): """Command generator for generating pose commands uniformly. The command generator generates poses by sampling positions uniformly within specified @@ -64,7 +63,7 @@ def __init__(self, cfg: UniformPoseCommandCfg, env: ManagerBasedEnv): # create buffers # -- commands: (x, y, z, qx, qy, qz, qw) in root frame self.pose_command_b = torch.zeros(self.num_envs, 7, device=self.device) - self.pose_command_b[:, 6] = 1.0 + self.pose_command_b[:, 3] = 1.0 self.pose_command_w = torch.zeros_like(self.pose_command_b) # -- metrics self.metrics["position_error"] = torch.zeros(self.num_envs, device=self.device) @@ -152,3 +151,31 @@ def _resample_command(self, env_ids: Sequence[int]): def _update_command(self): pass + + def _set_debug_vis_impl(self, debug_vis: bool): + # create markers if necessary for the first time + if debug_vis: + if not hasattr(self, "goal_pose_visualizer"): + # -- goal pose + self.goal_pose_visualizer = VisualizationMarkers(self.cfg.goal_pose_visualizer_cfg) + # -- current body pose + self.current_pose_visualizer = VisualizationMarkers(self.cfg.current_pose_visualizer_cfg) + # set their visibility to true + self.goal_pose_visualizer.set_visibility(True) + self.current_pose_visualizer.set_visibility(True) + else: + if hasattr(self, "goal_pose_visualizer"): + self.goal_pose_visualizer.set_visibility(False) + self.current_pose_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event): + # check if robot is initialized + # note: this is needed in-case the robot is de-initialized. we can't access the data + if not self.robot.is_initialized: + return + # update the markers + # -- goal pose + self.goal_pose_visualizer.visualize(self.pose_command_w[:, :3], self.pose_command_w[:, 3:]) + # -- current body pose + body_link_pose_w = self.robot.data.body_link_pose_w.torch[:, self.body_idx] + self.current_pose_visualizer.visualize(body_link_pose_w[:, :3], body_link_pose_w[:, 3:7]) diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py index 31fad8296bc4..930f663f65d3 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py @@ -16,8 +16,7 @@ import isaaclab.utils.math as math_utils from isaaclab.assets import Articulation from isaaclab.managers import CommandTerm - -from ._debug_vis import _VelocityCommandDebugVis +from isaaclab.markers import VisualizationMarkers if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -28,7 +27,7 @@ logger = logging.getLogger(__name__) -class UniformVelocityCommand(_VelocityCommandDebugVis, CommandTerm): +class UniformVelocityCommand(CommandTerm): r"""Command generator that generates a velocity command in SE(2) from uniform distribution. The command comprises of a linear velocity in x and y direction and an angular velocity around @@ -198,6 +197,63 @@ def _update_command(self): standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten() self.vel_command_b[standing_env_ids, :] = 0.0 + def _set_debug_vis_impl(self, debug_vis: bool): + # set visibility of markers + # note: parent only deals with callbacks. not their visibility + if debug_vis: + # create markers if necessary for the first time + if not hasattr(self, "goal_vel_visualizer"): + # -- goal + self.goal_vel_visualizer = VisualizationMarkers(self.cfg.goal_vel_visualizer_cfg) + # -- current + self.current_vel_visualizer = VisualizationMarkers(self.cfg.current_vel_visualizer_cfg) + # set their visibility to true + self.goal_vel_visualizer.set_visibility(True) + self.current_vel_visualizer.set_visibility(True) + else: + if hasattr(self, "goal_vel_visualizer"): + self.goal_vel_visualizer.set_visibility(False) + self.current_vel_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event): + # check if robot is initialized + # note: this is needed in-case the robot is de-initialized. we can't access the data + if not self.robot.is_initialized: + return + # get marker location + # -- base state + base_pos_w = self.robot.data.root_pos_w.torch.clone() + base_pos_w[:, 2] += 0.5 + # -- resolve the scales and quaternions + vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2]) + vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow( + self.robot.data.root_lin_vel_b.torch[:, :2] + ) + # display markers + self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) + self.current_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale) + + """ + Internal helpers. + """ + + def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Converts the XY base velocity command to arrow direction rotation.""" + # obtain default scale of the marker + default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale + # arrow-scale + arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1) + arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0 + # arrow-direction + heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0]) + zeros = torch.zeros_like(heading_angle) + arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) + # convert everything back from base to world frame + base_quat_w = self.robot.data.root_quat_w.torch + arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) + + return arrow_scale, arrow_quat + class NormalVelocityCommand(UniformVelocityCommand): """Command generator that generates a velocity command in SE(2) from a normal distribution. diff --git a/source/isaaclab/isaaclab/envs/mdp/events.py b/source/isaaclab/isaaclab/envs/mdp/events.py index 936a79adebd8..b0c7d9d9e988 100644 --- a/source/isaaclab/isaaclab/envs/mdp/events.py +++ b/source/isaaclab/isaaclab/envs/mdp/events.py @@ -2000,7 +2000,7 @@ def reset_root_state_from_terrain( # sample random valid poses ids = torch.randint(0, valid_positions.shape[2], size=(len(env_ids),), device=env.device) - positions = valid_positions[terrain.terrain_levels[env_ids], terrain.terrain_types[env_ids], ids] + positions = valid_positions[terrain.terrain_levels.torch[env_ids], terrain.terrain_types.torch[env_ids], ids] positions += asset.data.default_root_pose.torch[env_ids, :3] # sample random orientations diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index 15b693197a46..e6d8405ada29 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -36,6 +36,7 @@ from isaaclab.sim import SimulationContext from isaaclab.sim.utils.stage import get_current_stage, get_current_stage_id from isaaclab.sim.views import FrameView +from isaaclab.utils.warp import ProxyArray # Note: This is a temporary import for the VisuoTactileSensorCfg class. # It will be removed once the VisuoTactileSensor class is added to the core Isaac Lab framework. @@ -188,6 +189,11 @@ def __init__(self, cfg: InteractiveSceneCfg): if self._is_scene_setup_from_cfg(): self._add_entities_from_cfg() + if self._terrain is None: + # Scene origins are immutable after construction, so one ProxyArray safely serves both frontends. + positions = self.sim.get_clone_plan().positions + self._env_origins = ProxyArray(wp.from_torch(positions.contiguous(), dtype=wp.vec3f)) + self._aggregate_scene_data_requirements(requested_viz_types) # Collision filtering is PhysX-only (matches both physx and ovphysx). @@ -364,12 +370,19 @@ def num_envs(self) -> int: @property def env_origins(self) -> torch.Tensor: - """Per-env world origins, shape ``(num_envs, 3)``. From the terrain when registered, + """Per-environment world origins [m], shape ``(num_envs, 3)``. From the terrain when registered, else from the published :class:`~isaaclab.cloner.ClonePlan`. """ if self._terrain is not None: - return self._terrain.env_origins - return self.sim.get_clone_plan().positions + return self._terrain.env_origins.torch + return self._env_origins.torch + + @property + def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): + """Zero-copy Warp view of environment world origins [m], shape ``(num_envs,)`` of ``wp.vec3f``.""" + if self._terrain is not None: + return self._terrain.env_origins.warp + return self._env_origins.warp @property def terrain(self) -> TerrainImporter | None: diff --git a/source/isaaclab/isaaclab/terrains/terrain_importer.py b/source/isaaclab/isaaclab/terrains/terrain_importer.py index 2320996bd774..d6428d3f18d8 100644 --- a/source/isaaclab/isaaclab/terrains/terrain_importer.py +++ b/source/isaaclab/isaaclab/terrains/terrain_importer.py @@ -16,6 +16,7 @@ import isaaclab.sim as sim_utils from isaaclab.markers import VisualizationMarkers from isaaclab.markers.config import FRAME_MARKER_CFG +from isaaclab.utils.warp import ProxyArray from .utils import create_prim_from_mesh @@ -38,10 +39,12 @@ def _update_env_origins_mask( env_origins: wp.array(dtype=wp.vec3f), max_terrain_level: int, ): + """Update terrain levels and origins selected by an environment mask.""" env_id = wp.tid() if env_mask[env_id]: state = rng_state[env_id] random_level = wp.randi(state, 0, max_terrain_level) + # Keep shared terrain indices int64 so their zero-copy Torch views support advanced indexing. level = terrain_levels[env_id] + wp.int64(move_up[env_id]) - wp.int64(move_down[env_id]) if level >= wp.int64(max_terrain_level): level = wp.int64(random_level) @@ -72,15 +75,32 @@ class TerrainImporter: terrain_prim_paths: list[str] """A list containing the USD prim paths to the imported terrains.""" - terrain_origins: torch.Tensor | None - """The origins of the sub-terrains in the added terrain mesh. Shape is (num_rows, num_cols, 3). + terrain_origins: ProxyArray | None + """The origins [m] of the sub-terrains, shape ``(num_rows, num_cols)`` of ``wp.vec3f``. + + The :attr:`~isaaclab.utils.warp.ProxyArray.torch` view resolves to shape ``(num_rows, num_cols, 3)``. If terrain origins is not None, the environment origins are computed based on the terrain origins. Otherwise, the environment origins are computed based on the grid spacing. """ - env_origins: torch.Tensor - """The origins of the environments. Shape is (num_envs, 3).""" + env_origins: ProxyArray + """The environment origins [m], shape ``(num_envs,)`` of ``wp.vec3f``. + + The :attr:`~isaaclab.utils.warp.ProxyArray.torch` view resolves to shape ``(num_envs, 3)``. + """ + + terrain_levels: ProxyArray | None + """The terrain level assigned to each environment, shape ``(num_envs,)`` of ``wp.int64``. + + This is available only when sub-terrain origins are configured. + """ + + terrain_types: ProxyArray | None + """The terrain type assigned to each environment, shape ``(num_envs,)`` of ``wp.int64``. + + This is available only when sub-terrain origins are configured. + """ def __init__(self, cfg: TerrainImporterCfg): """Initialize the terrain importer. @@ -103,11 +123,8 @@ def __init__(self, cfg: TerrainImporterCfg): # create buffers for the terrains self.terrain_prim_paths = list() self.terrain_origins = None - self.env_origins = None # assigned later when `configure_env_origins` is called - self._terrain_levels_wp = None - self._terrain_types_wp = None - self._terrain_origins_wp = None - self._env_origins_wp = None + self.terrain_levels = None + self.terrain_types = None # private variables self._terrain_flat_patches = dict() @@ -175,12 +192,6 @@ def terrain_names(self) -> list[str]: """A list of names of the imported terrains.""" return [f"'{path.split('/')[-1]}'" for path in self.terrain_prim_paths] - @property - def terrain_levels_wp(self) -> wp.array(dtype=wp.int64): - """Pointer-stable Warp view of the per-environment terrain levels.""" - self._initialize_warp_origin_views() - return self._terrain_levels_wp - """ Operations - Visibility. """ @@ -205,11 +216,9 @@ def set_debug_vis(self, debug_vis: bool) -> bool: cfg=FRAME_MARKER_CFG.replace(prim_path="/Visuals/TerrainOrigin") ) if self.terrain_origins is not None: - self.origin_visualizer.visualize(self.terrain_origins.reshape(-1, 3)) - elif self.env_origins is not None: - self.origin_visualizer.visualize(self.env_origins.reshape(-1, 3)) + self.origin_visualizer.visualize(self.terrain_origins.torch.reshape(-1, 3)) else: - raise RuntimeError("Terrain origins are not configured.") + self.origin_visualizer.visualize(self.env_origins.torch.reshape(-1, 3)) # set visibility self.origin_visualizer.set_visibility(True) else: @@ -325,46 +334,59 @@ def import_usd(self, name: str, usd_path: str): Operations - Origins. """ - def configure_env_origins(self, origins: np.ndarray | torch.Tensor | None = None): + def configure_env_origins(self, origins: np.ndarray | torch.Tensor | None = None) -> None: """Configure the origins of the environments based on the added terrain. Args: - origins: The origins of the sub-terrains. Shape is (num_rows, num_cols, 3). + origins: The origins [m] of the sub-terrains. Shape is (num_rows, num_cols, 3). + + Note: + This setup operation replaces the :class:`~isaaclab.utils.warp.ProxyArray` instances. Call it before + managers or kernels cache their ``.warp`` views. """ # decide whether to compute origins in a grid or based on curriculum if origins is not None: - # convert to numpy - if isinstance(origins, np.ndarray): - origins = torch.from_numpy(origins) - # store the origins - self.terrain_origins = origins.to(self.device, dtype=torch.float) + origins_torch = torch.as_tensor(origins, dtype=torch.float32, device=self.device).contiguous() + # ``wp.from_torch`` is zero-copy and keeps the source tensor alive through the Warp array. + self.terrain_origins = ProxyArray(wp.from_torch(origins_torch, dtype=wp.vec3f)) # compute environment origins self.env_origins = self._compute_env_origins_curriculum(self.cfg.num_envs, self.terrain_origins) else: self.terrain_origins = None + self.terrain_levels = None + self.terrain_types = None # check if env spacing is valid if self.cfg.env_spacing is None: raise ValueError("Environment spacing must be specified for configuring grid-like origins.") # compute environment origins self.env_origins = self._compute_env_origins_grid(self.cfg.num_envs, self.cfg.env_spacing) - self._invalidate_warp_origin_views() - def update_env_origins(self, env_ids: torch.Tensor, move_up: torch.Tensor, move_down: torch.Tensor): - """Update the environment origins based on the terrain levels.""" + def update_env_origins(self, env_ids: torch.Tensor, move_up: torch.Tensor, move_down: torch.Tensor) -> None: + """Update environment origins through the Torch ID-based interface. + + Args: + env_ids: Environment indices to update. + move_up: Flags that increment the selected terrain levels. + move_down: Flags that decrement the selected terrain levels. + """ # check if grid-like spawning - if self.terrain_origins is None: + if self.terrain_origins is None or self.terrain_levels is None or self.terrain_types is None: return + terrain_origins = self.terrain_origins.torch + terrain_levels = self.terrain_levels.torch + terrain_types = self.terrain_types.torch + env_origins = self.env_origins.torch # update terrain level for the envs - self.terrain_levels[env_ids] += 1 * move_up - 1 * move_down + terrain_levels[env_ids] += 1 * move_up - 1 * move_down # robots that solve the last level are sent to a random one # the minimum level is zero - self.terrain_levels[env_ids] = torch.where( - self.terrain_levels[env_ids] >= self.max_terrain_level, - torch.randint_like(self.terrain_levels[env_ids], self.max_terrain_level), - torch.clip(self.terrain_levels[env_ids], 0), + terrain_levels[env_ids] = torch.where( + terrain_levels[env_ids] >= self.max_terrain_level, + torch.randint_like(terrain_levels[env_ids], self.max_terrain_level), + torch.clip(terrain_levels[env_ids], 0), ) # update the env origins - self.env_origins[env_ids] = self.terrain_origins[self.terrain_levels[env_ids], self.terrain_types[env_ids]] + env_origins[env_ids] = terrain_origins[terrain_levels[env_ids], terrain_types[env_ids]] def update_env_origins_mask( self, @@ -373,27 +395,22 @@ def update_env_origins_mask( move_down: wp.array(dtype=wp.bool), rng_state: wp.array(dtype=wp.uint32), ) -> None: - """Update terrain levels and origins for a boolean environment mask. + """Update terrain levels and origins through the Warp mask-based interface. - This method preserves the level update semantics of :meth:`update_env_origins`: moving past the - maximum level wraps to a random level, moving below zero clamps to zero, and unselected environments - remain unchanged. + Moving past the maximum level wraps to a random level, moving below zero clamps to zero, and + unselected environments remain unchanged. Args: - env_mask: Boolean Warp mask selecting environments to update. Shape is - ``(num_envs,)`` on :attr:`device`. - move_up: Boolean Warp flags that increment selected terrain levels. - Shape is ``(num_envs,)`` on :attr:`device`. - move_down: Boolean Warp flags that decrement selected terrain levels. - Shape is ``(num_envs,)`` on :attr:`device`. - rng_state: Per-environment Warp random-number-generator state. Shape is - ``(num_envs,)`` on :attr:`device`. + env_mask: Boolean mask selecting environments, shape ``(num_envs,)``. + move_up: Flags that increment terrain levels, shape ``(num_envs,)``. + move_down: Flags that decrement terrain levels, shape ``(num_envs,)``. + rng_state: Per-environment random-number-generator state, shape ``(num_envs,)``. Raises: - ValueError: If an input array does not match the configured environment count. - TypeError: If an input array has the wrong Warp data type. + TypeError: If an input is not a Warp array with the required data type. + ValueError: If an input has the wrong shape or device. """ - if self.terrain_origins is None: + if self.terrain_origins is None or self.terrain_levels is None or self.terrain_types is None: return num_envs = self.terrain_levels.shape[0] arrays = { @@ -410,7 +427,6 @@ def update_env_origins_mask( if array.device != wp.get_device(self.device): raise ValueError(f"{name} must be on device {self.device}; received {array.device}.") - self._initialize_warp_origin_views() wp.launch( kernel=_update_env_origins_mask, dim=num_envs, @@ -419,10 +435,10 @@ def update_env_origins_mask( move_up, move_down, rng_state, - self._terrain_levels_wp, - self._terrain_types_wp, - self._terrain_origins_wp, - self._env_origins_wp, + self.terrain_levels.warp, + self.terrain_types.warp, + self.terrain_origins.warp, + self.env_origins.warp, self.max_terrain_level, ], device=self.device, @@ -432,25 +448,7 @@ def update_env_origins_mask( Internal helpers. """ - def _invalidate_warp_origin_views(self) -> None: - """Invalidate cached Warp views after origin storage changes.""" - self._terrain_levels_wp = None - self._terrain_types_wp = None - self._terrain_origins_wp = None - self._env_origins_wp = None - - def _initialize_warp_origin_views(self) -> None: - """Create persistent Warp views of terrain-origin Torch storage.""" - if self._terrain_levels_wp is not None: - return - if self.terrain_origins is None or self.env_origins is None: - raise RuntimeError("Terrain origins are not configured for curriculum updates.") - self._terrain_levels_wp = wp.from_torch(self.terrain_levels, dtype=wp.int64) - self._terrain_types_wp = wp.from_torch(self.terrain_types, dtype=wp.int64) - self._terrain_origins_wp = wp.from_torch(self.terrain_origins, dtype=wp.vec3f) - self._env_origins_wp = wp.from_torch(self.env_origins, dtype=wp.vec3f) - - def _compute_env_origins_curriculum(self, num_envs: int, origins: torch.Tensor) -> torch.Tensor: + def _compute_env_origins_curriculum(self, num_envs: int, origins: ProxyArray) -> ProxyArray: """Compute the origins of the environments defined by the sub-terrains origins.""" # extract number of rows and cols num_rows, num_cols = origins.shape[:2] @@ -462,21 +460,22 @@ def _compute_env_origins_curriculum(self, num_envs: int, origins: torch.Tensor) # store maximum terrain level possible self.max_terrain_level = num_rows # define all terrain levels and types available - self.terrain_levels = torch.randint(0, max_init_level + 1, (num_envs,), device=self.device) - self.terrain_types = torch.div( + terrain_levels = torch.randint(0, max_init_level + 1, (num_envs,), device=self.device) + terrain_types = torch.div( torch.arange(num_envs, device=self.device), (num_envs / num_cols), rounding_mode="floor" ).to(torch.long) - # create tensor based on number of environments - env_origins = torch.zeros(num_envs, 3, device=self.device) - env_origins[:] = origins[self.terrain_levels, self.terrain_types] - return env_origins + env_origins = origins.torch[terrain_levels, terrain_types].contiguous() + + self.terrain_levels = ProxyArray(wp.from_torch(terrain_levels, dtype=wp.int64)) + self.terrain_types = ProxyArray(wp.from_torch(terrain_types, dtype=wp.int64)) + return ProxyArray(wp.from_torch(env_origins, dtype=wp.vec3f)) - def _compute_env_origins_grid(self, num_envs: int, env_spacing: float) -> torch.Tensor: + def _compute_env_origins_grid(self, num_envs: int, env_spacing: float) -> ProxyArray: """Compute the origins of the environments in a grid based on configured spacing.""" from isaaclab.cloner import grid_transforms env_origins, _ = grid_transforms(num_envs, env_spacing, device=self.device) - return env_origins + return ProxyArray(wp.from_torch(env_origins.contiguous(), dtype=wp.vec3f)) """ Deprecated. diff --git a/source/isaaclab/test/terrains/check_terrain_importer.py b/source/isaaclab/test/terrains/check_terrain_importer.py index 901cfc420be3..7861f5c2b272 100644 --- a/source/isaaclab/test/terrains/check_terrain_importer.py +++ b/source/isaaclab/test/terrains/check_terrain_importer.py @@ -156,7 +156,7 @@ def main(): # Set ball positions over terrain origins using FrameView (before simulation starts) xform_view = sim_utils.FrameView("/World/envs/env_.*/ball") # cache initial state of the balls - ball_initial_positions = terrain_importer.env_origins.clone() + ball_initial_positions = terrain_importer.env_origins.torch.clone() ball_initial_positions[:, 2] += 5.0 # set initial poses (writes to USD before simulation) with xform_view.xform_world_space_writer() as w: diff --git a/source/isaaclab/test/terrains/test_terrain_importer.py b/source/isaaclab/test/terrains/test_terrain_importer.py index 553f01f79233..83192a472e1b 100644 --- a/source/isaaclab/test/terrains/test_terrain_importer.py +++ b/source/isaaclab/test/terrains/test_terrain_importer.py @@ -29,6 +29,7 @@ from isaaclab.terrains import TerrainImporter, TerrainImporterCfg from isaaclab.terrains.config.rough import ROUGH_TERRAINS_CFG from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.warp import ProxyArray pytestmark = pytest.mark.integration @@ -50,7 +51,15 @@ def test_terrain_importer_env_origins(device, env_spacing, num_envs): ) terrain_importer = TerrainImporter(terrain_importer_cfg) # obtain env origins using terrain importer - terrain_importer_origins = terrain_importer.env_origins + terrain_importer_origins = terrain_importer.env_origins.torch + + # check that the canonical Warp storage and Torch view share memory + assert isinstance(terrain_importer.env_origins, ProxyArray) + assert terrain_importer.env_origins.warp.dtype == wp.vec3f + assert terrain_importer.env_origins.warp.shape == (num_envs,) + assert terrain_importer_origins.shape == (num_envs, 3) + assert terrain_importer.env_origins.warp is terrain_importer.env_origins.warp + assert terrain_importer_origins.data_ptr() == terrain_importer.env_origins.warp.ptr # obtain env origins using Lab's grid_transforms lab_grid_origins, _ = lab_cloner.grid_transforms(num_envs, spacing=env_spacing, device=sim.device) @@ -314,7 +323,7 @@ def _populate_scene(sim: SimulationContext, num_balls: int = 2048, geom_sphere: # Create a view over all the balls using Isaac Lab's FrameView ball_view = sim_utils.FrameView("/World/envs/env_.*/ball") # cache initial state of the balls - ball_initial_positions = terrain_importer.env_origins.clone() + ball_initial_positions = terrain_importer.env_origins.torch.clone() ball_initial_positions[:, 2] += 5.0 # set initial poses # note: setting here writes to USD :) diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst index 234876b28416..2a6540bb8518 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -6,14 +6,16 @@ Added * Added :class:`~isaaclab_experimental.managers.CurriculumManager` and a boolean-mask reset path for Warp manager-based environments, retaining compact environment IDs only for legacy host consumers. +* Added Warp-native terrain curricula backed by the canonical Warp-owned + terrain and scene origins. Changed ^^^^^^^ -* Changed Warp manager and direct-environment stages to execute eagerly by - default while stateful capture semantics are validated. Set - ``MANAGER_CALL_CONFIG='{"default": 2}'`` for manager capture or - ``ISAACLAB_WARP_DIRECT_CAPTURE=1`` for the legacy direct capture path. +* Changed Warp manager stages to execute eagerly by default while stateful + capture semantics are validated. Set + ``MANAGER_CALL_CONFIG='{"default": 2}'`` to opt eligible manager stages into + CUDA graph capture. Direct-environment stages remained eager. Deprecated ^^^^^^^^^^ @@ -26,3 +28,5 @@ Fixed * Fixed Warp event state leaking across environments and configurations, and corrected center-of-mass sampling and termination metrics for partial resets. +* Fixed captured velocity-command terms to read current root state on every + replay instead of stale lazy-derived buffers. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py index 02eaff52ff4e..f6f327b9ada2 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py @@ -54,9 +54,6 @@ DEBUG_TIMERS = os.environ.get("DEBUG_TIMERS", "0") == "1" """Enable all fine-grained inner timers (adds wp.synchronize per sub-phase). Set DEBUG_TIMERS=1 env var to enable.""" -WARP_DIRECT_CAPTURE = os.environ.get("ISAACLAB_WARP_DIRECT_CAPTURE", "0") == "1" -"""Enable direct-environment stage capture. Eager execution is the correctness-first default.""" - class DirectRLEnvWarp(DirectRLEnv): """The superclass for the direct workflow to design environments. @@ -228,10 +225,10 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.torch_reset_time_outs: torch.Tensor = None self.torch_episode_length_buf: torch.Tensor = None - # Direct stages are Warp eager by default. The owner-held cache keeps - # execution dispatch out of call sites and can enable capture later once - # stateful warm-up semantics are defined. - self._graph_cache = WarpGraphCache(enabled=WARP_DIRECT_CAPTURE) + # Direct stages stay eager until their complete backend boundaries are + # verified capture-safe. The owner-held executor keeps that policy out of + # individual stage call sites. + self._warp_graph_cache = WarpGraphCache(enabled=False) # setup the action and observation spaces for Gym self._configure_gym_env_spaces() @@ -414,7 +411,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: # set actions into buffers # simulate with Timer(name="apply_action", msg="Action processing step took:", enable=DEBUG_TIMERS): - self._graph_cache.capture_or_replay("action", self.step_warp_action) + self._warp_graph_cache.call("action", self.step_warp_action) # Keep scene writes outside the task graph until scene, sensor, and # actuator capturability have been validated as one backend boundary. @@ -434,12 +431,12 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self.common_step_counter += 1 # total step (common for all envs) with Timer(name="end_pre_graph", msg="End pre-graph took:", enable=DEBUG_TIMERS): - self._graph_cache.capture_or_replay("end_pre", self._step_warp_end_pre) + self._warp_graph_cache.call("end_pre", self._step_warp_end_pre) # Keep the post-reset scene write at the explicit backend boundary. with Timer(name="write_data_to_sim_post", msg="Write data to sim (post-reset) took:", enable=DEBUG_TIMERS): self.scene.write_data_to_sim() with Timer(name="end_post_graph", msg="End post-graph took:", enable=DEBUG_TIMERS): - self._graph_cache.capture_or_replay("end_post", self._step_warp_end_post) + self._warp_graph_cache.call("end_post", self._step_warp_end_post) # Visualization hook — runs after CUDA graph scope. Override in subclass # to update markers or other non-graphable visual elements. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py index 03295b324010..f1f390e8f163 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py @@ -16,11 +16,9 @@ # import builtins import contextlib -import importlib import logging import warnings from collections.abc import Sequence -from copy import deepcopy from typing import Any import torch @@ -38,8 +36,9 @@ from isaaclab.utils.version import has_kit from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp as InteractiveScene -from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode, ManagerCallSwitch +from isaaclab_experimental.utils.manager_call_switch import ManagerCallSwitch from isaaclab_experimental.utils.warp import resolve_1d_mask +from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache # import logger logger = logging.getLogger(__name__) @@ -79,11 +78,18 @@ def __init__(self, cfg: ManagerBasedEnvCfg): self.cfg = cfg # initialize internal variables self._is_closed = False - # temporary debug runtime config for manager source/call switching. - cfg_source: dict | str | None = getattr(self.cfg, "manager_call_config", None) + # Keep the execution policy on the environment. ManagerCallSwitch is a + # compatibility router that can be removed without replacing this owner. + self._warp_graph_cache = WarpGraphCache(enabled=True) + # Temporary runtime config for stable/Warp manager routing. + cfg_source: dict[str, int] | str | None = getattr(self.cfg, "manager_call_config", None) max_modes: dict[str, int] | None = getattr(self.cfg, "manager_call_max_mode", None) - self._manager_call_switch = ManagerCallSwitch(cfg_source, max_modes=max_modes) - self._apply_manager_term_cfg_profile() + self._manager_call_switch = ManagerCallSwitch( + cfg_source, + max_modes=max_modes, + graph_cache=self._warp_graph_cache, + ) + self._manager_call_switch.apply_term_cfg_profile(self.cfg) # set the seed for the environment if self.cfg.seed is not None: @@ -273,15 +279,9 @@ def device(self): return self.sim.device @property - def env_origins_wp(self) -> wp.array: - """Scene env origins as a warp ``vec3f`` array. Cached on first access.""" - if not hasattr(self, "_env_origins_wp"): - origins = self.scene.env_origins - if isinstance(origins, wp.array): - self._env_origins_wp = origins - else: - self._env_origins_wp = wp.from_torch(origins, dtype=wp.vec3f) - return self._env_origins_wp + def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): + """Warp-owned environment origins [m] provided by the scene.""" + return self.scene.env_origins_wp def resolve_env_mask( self, @@ -444,7 +444,7 @@ def reset( device=self.device, ) - self._reset_idx(env_mask=reset_mask, env_ids=host_env_ids) + self._reset_idx(env_mask=reset_mask) if host_env_ids is not None: self.extras["log"].update(self._reset_host_post(host_env_ids)) @@ -502,7 +502,7 @@ def reset_to( env_mask = self.resolve_env_mask(env_ids=env_ids) self._reset_host_pre(env_ids) - self._reset_idx(env_mask=env_mask, env_ids=env_ids) + self._reset_idx(env_mask=env_mask) self.extras["log"].update(self._reset_host_post(env_ids)) # set the state @@ -629,94 +629,16 @@ def close(self): Helper functions. """ - def _resolve_stable_cfg_counterpart(self) -> ManagerBasedEnvCfg | None: - """Resolve the legacy stable task config counterpart.""" - cfg_cls = self.cfg.__class__ - cfg_module_name = cfg_cls.__module__ - if "isaaclab_tasks_experimental" not in cfg_module_name: - return None - - stable_module_name = cfg_module_name.replace("isaaclab_tasks_experimental", "isaaclab_tasks", 1) - try: - stable_module = importlib.import_module(stable_module_name) - except Exception as exc: - logger.warning( - "Failed to import stable task cfg module '%s' for manager_call_config stable mode: %s", - stable_module_name, - exc, - ) - return None - - stable_cfg_cls = getattr(stable_module, cfg_cls.__name__, None) - if stable_cfg_cls is None: - logger.warning("Stable task cfg class '%s' not found in module '%s'.", cfg_cls.__name__, stable_module_name) - return None - - try: - return stable_cfg_cls() - except Exception as exc: - logger.warning( - "Failed to instantiate stable task cfg '%s.%s': %s", - stable_module_name, - cfg_cls.__name__, - exc, - ) - return None - - def _apply_manager_term_cfg_profile(self) -> None: - """Preserve the deprecated mixed stable-manager configuration path.""" - manager_to_cfg_attr = { - "ActionManager": "actions", - "ObservationManager": "observations", - "EventManager": "events", - "RecorderManager": "recorders", - "CommandManager": "commands", - "TerminationManager": "terminations", - "RewardManager": "rewards", - "CurriculumManager": "curriculum", - } - stable_managers = [ - name - for name in manager_to_cfg_attr - if self._manager_call_switch.get_mode_for_manager(name) == ManagerCallMode.STABLE - ] - if not stable_managers: - return - - warnings.warn( - "Selecting STABLE managers inside ManagerBasedEnvWarp is deprecated. Use ManagerBasedEnv for Torch " - "managers or WARP_NOT_CAPTURED for the Warp frontend.", - DeprecationWarning, - stacklevel=2, - ) - stable_cfg = self._resolve_stable_cfg_counterpart() - if stable_cfg is None: - logger.warning( - "Stable managers requested (%s), but no stable cfg counterpart could be resolved. Keeping " - "experimental term configs.", - ", ".join(stable_managers), - ) - return - - for manager_name, cfg_attr in manager_to_cfg_attr.items(): - if self._manager_call_switch.get_mode_for_manager(manager_name) != ManagerCallMode.STABLE: - continue - if hasattr(self.cfg, cfg_attr) and hasattr(stable_cfg, cfg_attr): - setattr(self.cfg, cfg_attr, deepcopy(getattr(stable_cfg, cfg_attr))) - def _reset_idx( self, *, env_mask: wp.array(dtype=wp.bool), - env_ids: Sequence[int] | torch.Tensor | None = None, ) -> None: """Reset Warp-owned state for selected environments. Args: env_mask: Boolean Warp mask selecting environments to reset. - env_ids: Optional compact IDs for explicit host-only terms. """ - del env_ids # reset the internal buffers of the scene elements self.scene.reset(env_mask=env_mask) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index e62112b5ff58..696bf001b42b 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -15,7 +15,6 @@ import math import os -from collections.abc import Sequence from typing import Any, ClassVar import gymnasium as gym @@ -203,7 +202,7 @@ def invalidate_wp_graphs(self) -> None: Call this if the captured launch topology changes (e.g. different term list, shapes, etc.). """ - self._manager_call_switch.invalidate_graphs() + self._warp_graph_cache.invalidate() def step_warp_termination_compute(self) -> None: """Captured stage: compute terminations (env-step frequency).""" @@ -466,14 +465,11 @@ def _reset_idx( self, *, env_mask: wp.array(dtype=wp.bool), - env_ids: Sequence[int] | torch.Tensor | None = None, ) -> None: """Reset Warp-owned RL state for selected environments. Args: env_mask: Boolean Warp mask selecting environments to reset. - env_ids: Optional compact IDs for legacy curriculum terms that - explicitly require a host selection. """ if env_mask is not self.reset_mask_wp: wp.copy(self.reset_mask_wp, env_mask) @@ -483,7 +479,6 @@ def _reset_idx( "CurriculumManager_compute", self.curriculum_manager.compute, env_mask=env_mask, - env_ids=env_ids, _timer=DEBUG_TIMER_RESET, ) @@ -531,7 +526,6 @@ def _reset_idx( "CurriculumManager_reset", self.curriculum_manager.reset, env_mask=env_mask, - env_ids=env_ids, _timer=DEBUG_TIMER_RESET, ) @@ -607,29 +601,12 @@ def _reset_terminated_envs(self) -> None: enable=DEBUG_TIMER_STEP, time_unit="us", ): - self._reset_idx(env_mask=reset_mask, env_ids=reset_env_ids) + self._reset_idx(env_mask=reset_mask) if reset_env_ids is not None and reset_env_ids.numel() > 0: self.extras["log"].update(self._reset_host_post(reset_env_ids)) - if self.has_rtx_sensors and self.cfg.num_rerenders_on_reset > 0: - for _ in range(self.cfg.num_rerenders_on_reset): - self.sim.render() if self._has_recorders: self.recorder_manager.record_post_reset(reset_env_ids) - - def _reset_host_pre(self, env_ids: Sequence[int] | torch.Tensor) -> None: - """Run host-only pre-reset work for selected environments.""" - super()._reset_host_pre(env_ids) - - def _reset_host_post(self, env_ids: Sequence[int] | torch.Tensor) -> dict[str, Any]: - """Run host-only manager resets and return their logging values.""" - return super()._reset_host_post(env_ids) - - def _reset_requires_host_selection(self) -> bool: - """Return whether enabled reset features require a host-visible selection.""" - return bool( - self.curriculum_manager.requires_host_boundary - or self._has_recorders - or self.scene.surface_grippers - or (self.has_rtx_sensors and self.cfg.num_rerenders_on_reset > 0) - ) + if self.has_rtx_sensors and self.cfg.num_rerenders_on_reset > 0: + for _ in range(self.cfg.num_rerenders_on_reset): + self.sim.render() diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/_debug_vis.py similarity index 94% rename from source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py rename to source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/_debug_vis.py index 72fb7503ef72..e58e4a42e8ba 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/_debug_vis.py @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Shared debug-visualization helpers for command terms.""" +"""Private debug-visualization helpers for Warp-native command terms.""" from __future__ import annotations @@ -14,7 +14,7 @@ class _VelocityCommandDebugVis: - """Debug visualization shared by stable and Warp velocity commands.""" + """Debug visualization for Warp-native velocity commands.""" def _set_debug_vis_impl(self, debug_vis: bool): if debug_vis: @@ -52,7 +52,7 @@ def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torc class _PoseCommandDebugVis: - """Debug visualization shared by stable and Warp pose commands.""" + """Debug visualization for Warp-native pose commands.""" def _set_debug_vis_impl(self, debug_vis: bool): if debug_vis: diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py index b88d3be3d546..deaa5a5f569d 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py @@ -13,11 +13,12 @@ import warp as wp from isaaclab.assets import Articulation -from isaaclab.envs.mdp.commands._debug_vis import _PoseCommandDebugVis from isaaclab.utils.leapp import POSE7_ELEMENT_NAMES from isaaclab_experimental.managers import CommandTerm +from ._debug_vis import _PoseCommandDebugVis + if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py index 57dae69d59cf..24fbcefcdf16 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py @@ -13,13 +13,15 @@ import torch import warp as wp +from isaaclab_newton.kernels.state_kernels import body_ang_vel_from_root, body_lin_vel_from_root from isaaclab.assets import Articulation -from isaaclab.envs.mdp.commands._debug_vis import _VelocityCommandDebugVis from isaaclab_experimental.managers import CommandTerm from isaaclab_experimental.utils.warp import wrap_to_pi +from ._debug_vis import _VelocityCommandDebugVis + if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -31,17 +33,19 @@ @wp.kernel def _accumulate_velocity_metrics( command: wp.array(dtype=wp.float32, ndim=2), - root_lin_vel_b: wp.array(dtype=wp.vec3f), - root_ang_vel_b: wp.array(dtype=wp.vec3f), + root_pose_w: wp.array(dtype=wp.transformf), + root_vel_w: wp.array(dtype=wp.spatial_vectorf), error_xy_sum: wp.array(dtype=wp.float32), error_yaw_sum: wp.array(dtype=wp.float32), step_count: wp.array(dtype=wp.float32), ): env_id = wp.tid() - dx = command[env_id, 0] - root_lin_vel_b[env_id][0] - dy = command[env_id, 1] - root_lin_vel_b[env_id][1] + root_lin_vel_b = body_lin_vel_from_root(root_pose_w[env_id], root_vel_w[env_id]) + root_ang_vel_b = body_ang_vel_from_root(root_pose_w[env_id], root_vel_w[env_id]) + dx = command[env_id, 0] - root_lin_vel_b[0] + dy = command[env_id, 1] - root_lin_vel_b[1] error_xy_sum[env_id] += wp.sqrt(dx * dx + dy * dy) - error_yaw_sum[env_id] += wp.abs(command[env_id, 2] - root_ang_vel_b[env_id][2]) + error_yaw_sum[env_id] += wp.abs(command[env_id, 2] - root_ang_vel_b[2]) step_count[env_id] += 1.0 @@ -113,7 +117,7 @@ def _update_velocity_command( heading_target: wp.array(dtype=wp.float32), is_heading_env: wp.array(dtype=wp.bool), is_standing_env: wp.array(dtype=wp.bool), - heading_w: wp.array(dtype=wp.float32), + root_pose_w: wp.array(dtype=wp.transformf), heading_command: bool, heading_control_stiffness: float, ang_vel_z_min: float, @@ -121,7 +125,10 @@ def _update_velocity_command( ): env_id = wp.tid() if heading_command and is_heading_env[env_id]: - heading_error = wrap_to_pi(heading_target[env_id] - heading_w[env_id]) + root_quat_w = wp.transform_get_rotation(root_pose_w[env_id]) + forward_w = wp.quat_rotate(root_quat_w, wp.vec3f(1.0, 0.0, 0.0)) + heading_w = wp.atan2(forward_w[1], forward_w[0]) + heading_error = wrap_to_pi(heading_target[env_id] - heading_w) command[env_id, 2] = wp.clamp( heading_control_stiffness * heading_error, ang_vel_z_min, @@ -244,8 +251,8 @@ def _update_metrics(self): dim=self.num_envs, inputs=[ self._vel_command_b_wp, - self.robot.data.root_lin_vel_b.warp, - self.robot.data.root_ang_vel_b.warp, + self.robot.data.root_pose_w.warp, + self.robot.data.root_vel_w.warp, self._error_xy_sum_wp, self._error_yaw_sum_wp, self._step_count_wp, @@ -289,7 +296,7 @@ def _update_command(self): self._heading_target_wp, self._is_heading_env_wp, self._is_standing_env_wp, - self.robot.data.heading_w.warp, + self.robot.data.root_pose_w.warp, self.cfg.heading_command, self.cfg.heading_control_stiffness, self.cfg.ranges.ang_vel_z[0], diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py index 0a4c3c274b9e..edff91465a9b 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py @@ -5,7 +5,7 @@ """Warp-first overrides for common event terms. -These functions are intended to be used with the experimental Warp-first +These terms are intended to be used with the experimental Warp-first :class:`isaaclab_experimental.managers.EventManager` (mask-based interval/reset). Why this exists: @@ -17,6 +17,9 @@ These Warp-first implementations avoid that by writing directly into the sim-bound Warp state buffers (`asset.data.joint_pos` / `asset.data.joint_vel`) for the selected envs/joints. +Stateful terms use :class:`~isaaclab_experimental.managers.ManagerTermBase` so +persistent buffers and parsed constants are created once during manager setup. + Notes: - These terms assume the Newton/Warp backend (Warp arrays are available for joint state and defaults). - For best performance, pass :class:`isaaclab_experimental.managers.SceneEntityCfg` so `joint_ids_wp` is cached. @@ -24,145 +27,39 @@ from __future__ import annotations -from typing import TYPE_CHECKING, ClassVar +import warnings +from typing import TYPE_CHECKING import warp as wp -from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_experimental.utils.warp import WarpCapturable +from isaaclab_experimental.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg +from isaaclab_experimental.utils.warp import WarpCapturable, warp_capturable + +__all__ = [ + "ApplyExternalForceTorque", + "PushBySettingVelocity", + "RandomizeRigidBodyCom", + "ResetRootStateUniform", + "apply_external_force_torque", + "push_by_setting_velocity", + "randomize_rigid_body_com", + "reset_joints_by_offset", + "reset_joints_by_scale", + "reset_root_state_uniform", +] if TYPE_CHECKING: from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedEnv -def _range_cache_key( - range_values: dict[str, tuple[float, float]] | tuple[float, float], -) -> tuple[object, ...]: - """Return an immutable snapshot of an event range for state-cache lookup.""" - if isinstance(range_values, dict): - return tuple((name, tuple(bounds)) for name, bounds in sorted(range_values.items())) - return tuple(range_values) - - -class _EventStateCache: - """Environment-owned storage for persistent Warp event buffers and constants. - - The cache lifetime matches the environment lifetime. Keys include the event - callable, asset, term configuration, and immutable range values so stale parsed - constants are not reused after a range changes. - """ - - _ENV_ATTRIBUTE: ClassVar[str] = "_warp_event_state_cache" - - @classmethod - def for_env(cls, env: ManagerBasedEnv) -> dict[tuple[object, ...], object]: - """Return the event-state cache owned by an environment.""" - cache = getattr(env, cls._ENV_ATTRIBUTE, None) - if cache is None: - cache = {} - setattr(env, cls._ENV_ATTRIBUTE, cache) - return cache - - -class _RandomizeRigidBodyComState: - """Persistent arguments for center-of-mass randomization.""" - - def __init__( - self, - asset: RigidObject | Articulation, - asset_cfg: SceneEntityCfg, - com_range: dict[str, tuple[float, float]], - ): - self.asset = asset - self.asset_cfg = asset_cfg - self.com_range = com_range - self.default_com = wp.clone(asset.data.body_com_pos_b.warp) - ranges = [com_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z")] - self.com_lo = wp.vec3f(ranges[0][0], ranges[1][0], ranges[2][0]) - self.com_hi = wp.vec3f(ranges[0][1], ranges[1][1], ranges[2][1]) - - -class _ApplyExternalForceTorqueState: - """Persistent output buffers for external wrench randomization.""" - - def __init__( - self, - env: ManagerBasedEnv, - asset: RigidObject | Articulation, - asset_cfg: SceneEntityCfg, - force_range: tuple[float, float], - torque_range: tuple[float, float], - ): - self.asset = asset - self.asset_cfg = asset_cfg - self.force_range = force_range - self.torque_range = torque_range - self.forces = wp.zeros((env.num_envs, asset.num_bodies), dtype=wp.vec3f, device=env.device) - self.torques = wp.zeros((env.num_envs, asset.num_bodies), dtype=wp.vec3f, device=env.device) - if asset_cfg.body_ids is None or asset_cfg.body_ids == slice(None): - body_ids = list(range(asset.num_bodies)) - elif isinstance(asset_cfg.body_ids, int): - body_ids = [asset_cfg.body_ids] - else: - body_ids = list(asset_cfg.body_ids) - body_mask = [False] * asset.num_bodies - for body_id in body_ids: - body_mask[body_id] = True - self.body_ids = wp.array(body_ids, dtype=wp.int32, device=env.device) - self.body_mask = wp.array(body_mask, dtype=wp.bool, device=env.device) - - -class _PushBySettingVelocityState: - """Persistent output buffer and ranges for root-velocity pushes.""" - - def __init__( - self, - env: ManagerBasedEnv, - asset: RigidObject | Articulation, - asset_cfg: SceneEntityCfg, - velocity_range: dict[str, tuple[float, float]], - ): - self.asset = asset - self.asset_cfg = asset_cfg - self.velocity_range = velocity_range - self.velocity = wp.zeros((env.num_envs,), dtype=wp.spatial_vectorf, device=env.device) - ranges = [velocity_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z", "roll", "pitch", "yaw")] - self.lin_lo = wp.vec3f(ranges[0][0], ranges[1][0], ranges[2][0]) - self.lin_hi = wp.vec3f(ranges[0][1], ranges[1][1], ranges[2][1]) - self.ang_lo = wp.vec3f(ranges[3][0], ranges[4][0], ranges[5][0]) - self.ang_hi = wp.vec3f(ranges[3][1], ranges[4][1], ranges[5][1]) - - -class _ResetRootStateUniformState: - """Persistent output buffers and ranges for uniform root-state reset.""" - - def __init__( - self, - env: ManagerBasedEnv, - asset: RigidObject | Articulation, - asset_cfg: SceneEntityCfg, - pose_range: dict[str, tuple[float, float]], - velocity_range: dict[str, tuple[float, float]], - ): - self.asset = asset - self.asset_cfg = asset_cfg - self.pose_range = pose_range - self.velocity_range = velocity_range - self.pose = wp.zeros((env.num_envs,), dtype=wp.transformf, device=env.device) - self.velocity = wp.zeros((env.num_envs,), dtype=wp.spatial_vectorf, device=env.device) - - pose_ranges = [pose_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z", "roll", "pitch", "yaw")] - self.pos_lo = wp.vec3f(pose_ranges[0][0], pose_ranges[1][0], pose_ranges[2][0]) - self.pos_hi = wp.vec3f(pose_ranges[0][1], pose_ranges[1][1], pose_ranges[2][1]) - self.rot_lo = wp.vec3f(pose_ranges[3][0], pose_ranges[4][0], pose_ranges[5][0]) - self.rot_hi = wp.vec3f(pose_ranges[3][1], pose_ranges[4][1], pose_ranges[5][1]) - - velocity_ranges = [velocity_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z", "roll", "pitch", "yaw")] - self.vel_lin_lo = wp.vec3f(velocity_ranges[0][0], velocity_ranges[1][0], velocity_ranges[2][0]) - self.vel_lin_hi = wp.vec3f(velocity_ranges[0][1], velocity_ranges[1][1], velocity_ranges[2][1]) - self.vel_ang_lo = wp.vec3f(velocity_ranges[3][0], velocity_ranges[4][0], velocity_ranges[5][0]) - self.vel_ang_hi = wp.vec3f(velocity_ranges[3][1], velocity_ranges[4][1], velocity_ranges[5][1]) +def _resolve_body_ids(asset_cfg: SceneEntityCfg, num_bodies: int) -> list[int]: + """Resolve a configured body selection without allocating device storage.""" + if asset_cfg.body_ids is None or asset_cfg.body_ids == slice(None): + return list(range(num_bodies)) + if isinstance(asset_cfg.body_ids, int): + return [asset_cfg.body_ids] + return list(asset_cfg.body_ids) # --------------------------------------------------------------------------- @@ -196,44 +93,85 @@ def _randomize_com_kernel( rng_state[env_id] = state +@warp_capturable(False) +class RandomizeRigidBodyCom(ManagerTermBase): + """Randomize rigid-body centers of mass from a persistent default baseline. + + This term is not CUDA-graph capturable because notifying the solver of changed + inertial properties calls :meth:`SimulationManager.add_model_change`. + """ + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv) -> None: + """Initialize persistent center-of-mass randomization state. + + Args: + cfg: Event term configuration. + env: Environment containing the randomized asset. + """ + super().__init__(cfg, env) + asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self._asset: RigidObject | Articulation = env.scene[asset_cfg.name] + self._default_com = wp.clone(self._asset.data.body_com_pos_b.warp) + body_ids = _resolve_body_ids(asset_cfg, self._asset.num_bodies) + self._body_ids = wp.array(body_ids, dtype=wp.int32, device=env.device) + + com_range = cfg.params["com_range"] + ranges = [com_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z")] + self._com_lo = wp.vec3f(ranges[0][0], ranges[1][0], ranges[2][0]) + self._com_hi = wp.vec3f(ranges[0][1], ranges[1][1], ranges[2][1]) + + @WarpCapturable(False, reason="set_coms_mask calls SimulationManager.add_model_change") + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), + com_range: dict[str, tuple[float, float]], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + ) -> None: + """Randomize selected center-of-mass offsets [m]. + + Args: + env: Environment containing the randomized asset. + env_mask: Boolean Warp mask selecting environments. + com_range: Per-axis offset ranges [m]. Parsed during initialization. + asset_cfg: Scene entity selection. Resolved during initialization. + """ + wp.launch( + kernel=_randomize_com_kernel, + dim=env.num_envs, + inputs=[ + env_mask, + env.rng_state_wp, + self._default_com, + self._asset.data.body_com_pos_b.warp, + self._body_ids, + self._com_lo, + self._com_hi, + ], + device=env.device, + ) + + self._asset.set_coms_mask(coms=self._asset.data.body_com_pos_b.warp, env_mask=env_mask) + + @WarpCapturable(False, reason="set_coms_mask calls SimulationManager.add_model_change") def randomize_rigid_body_com( - env, - env_mask: wp.array, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), com_range: dict[str, tuple[float, float]], asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -): - """Randomize the center of mass (CoM) of rigid bodies by adding random offsets. - - Warp-first override of :func:`isaaclab.envs.mdp.events.randomize_rigid_body_com`. - Writes directly into the sim-bound ``body_com_pos_b`` buffer, then notifies the solver - via :meth:`set_coms_mask` so it recomputes inertial properties. - """ - asset: Articulation = env.scene[asset_cfg.name] - cache = _EventStateCache.for_env(env) - cache_key = (randomize_rigid_body_com, id(asset), id(asset_cfg), _range_cache_key(com_range)) - state = cache.get(cache_key) - if state is None: - state = _RandomizeRigidBodyComState(asset, asset_cfg, com_range) - cache[cache_key] = state - - wp.launch( - kernel=_randomize_com_kernel, - dim=env.num_envs, - inputs=[ - env_mask, - env.rng_state_wp, - state.default_com, - asset.data.body_com_pos_b.warp, - state.asset_cfg.body_ids_wp, - state.com_lo, - state.com_hi, - ], - device=env.device, +) -> None: + """Deprecated compatibility adapter for :class:`RandomizeRigidBodyCom`.""" + warnings.warn( + "'randomize_rigid_body_com' is deprecated; use 'RandomizeRigidBodyCom' in event configurations.", + DeprecationWarning, + stacklevel=2, ) - - # Notify the solver that inertial properties changed (COM position affects inertia). - asset.set_coms_mask(coms=asset.data.body_com_pos_b.warp, env_mask=env_mask) + params = {"com_range": com_range, "asset_cfg": asset_cfg} + cfg = EventTermCfg(func=RandomizeRigidBodyCom, mode="reset", params={}) + cfg.params = params + term = RandomizeRigidBodyCom(cfg, env) + term(env, env_mask, **params) # --------------------------------------------------------------------------- @@ -273,54 +211,97 @@ def _apply_external_force_torque_kernel( rng_state[env_id] = state +class ApplyExternalForceTorque(ManagerTermBase): + """Apply random external forces and torques using persistent wrench buffers.""" + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv) -> None: + """Initialize persistent external-wrench state. + + Args: + cfg: Event term configuration. + env: Environment containing the randomized asset. + """ + super().__init__(cfg, env) + asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self._asset: RigidObject | Articulation = env.scene[asset_cfg.name] + self._forces = wp.zeros((env.num_envs, self._asset.num_bodies), dtype=wp.vec3f, device=env.device) + self._torques = wp.zeros((env.num_envs, self._asset.num_bodies), dtype=wp.vec3f, device=env.device) + + body_ids = _resolve_body_ids(asset_cfg, self._asset.num_bodies) + body_mask = [False] * self._asset.num_bodies + for body_id in body_ids: + body_mask[body_id] = True + self._body_ids = wp.array(body_ids, dtype=wp.int32, device=env.device) + self._body_mask = wp.array(body_mask, dtype=wp.bool, device=env.device) + + force_range = cfg.params["force_range"] + torque_range = cfg.params["torque_range"] + self._force_lo = float(force_range[0]) + self._force_hi = float(force_range[1]) + self._torque_lo = float(torque_range[0]) + self._torque_hi = float(torque_range[1]) + + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), + force_range: tuple[float, float], + torque_range: tuple[float, float], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + ) -> None: + """Apply sampled forces [N] and torques [N·m] to selected bodies. + + Args: + env: Environment containing the randomized asset. + env_mask: Boolean Warp mask selecting environments. + force_range: Component-wise force range [N]. Parsed during initialization. + torque_range: Component-wise torque range [N·m]. Parsed during initialization. + asset_cfg: Scene entity selection. Resolved during initialization. + """ + wp.launch( + kernel=_apply_external_force_torque_kernel, + dim=env.num_envs, + inputs=[ + env_mask, + env.rng_state_wp, + self._body_ids, + self._forces, + self._torques, + self._force_lo, + self._force_hi, + self._torque_lo, + self._torque_hi, + ], + device=env.device, + ) + + self._asset.permanent_wrench_composer.set_forces_and_torques_mask( + forces=self._forces, + torques=self._torques, + body_mask=self._body_mask, + env_mask=env_mask, + ) + + +@WarpCapturable(False, reason="deprecated adapter constructs term state; use ApplyExternalForceTorque") def apply_external_force_torque( - env, - env_mask: wp.array, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), force_range: tuple[float, float], torque_range: tuple[float, float], asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -): - """Randomize external forces and torques applied to the asset's bodies. - - Warp-first override of :func:`isaaclab.envs.mdp.events.apply_external_force_torque`. - """ - asset: Articulation = env.scene[asset_cfg.name] - cache = _EventStateCache.for_env(env) - cache_key = ( - apply_external_force_torque, - id(asset), - id(asset_cfg), - _range_cache_key(force_range), - _range_cache_key(torque_range), - ) - state = cache.get(cache_key) - if state is None: - state = _ApplyExternalForceTorqueState(env, asset, asset_cfg, force_range, torque_range) - cache[cache_key] = state - - wp.launch( - kernel=_apply_external_force_torque_kernel, - dim=env.num_envs, - inputs=[ - env_mask, - env.rng_state_wp, - state.body_ids, - state.forces, - state.torques, - state.force_range[0], - state.force_range[1], - state.torque_range[0], - state.torque_range[1], - ], - device=env.device, - ) - - asset.permanent_wrench_composer.set_forces_and_torques_mask( - forces=state.forces, - torques=state.torques, - body_mask=state.body_mask, - env_mask=env_mask, +) -> None: + """Deprecated compatibility adapter for :class:`ApplyExternalForceTorque`.""" + warnings.warn( + "'apply_external_force_torque' is deprecated; use 'ApplyExternalForceTorque' in event configurations.", + DeprecationWarning, + stacklevel=2, ) + params = {"force_range": force_range, "torque_range": torque_range, "asset_cfg": asset_cfg} + cfg = EventTermCfg(func=ApplyExternalForceTorque, mode="reset", params={}) + cfg.params = params + term = ApplyExternalForceTorque(cfg, env) + term(env, env_mask, **params) # --------------------------------------------------------------------------- @@ -358,41 +339,80 @@ def _push_by_setting_velocity_kernel( rng_state[env_id] = state +class PushBySettingVelocity(ManagerTermBase): + """Push an asset by sampling into a persistent root-velocity buffer.""" + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv) -> None: + """Initialize persistent root-velocity push state. + + Args: + cfg: Event term configuration. + env: Environment containing the pushed asset. + """ + super().__init__(cfg, env) + asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self._asset: RigidObject | Articulation = env.scene[asset_cfg.name] + self._velocity = wp.zeros(env.num_envs, dtype=wp.spatial_vectorf, device=env.device) + + velocity_range = cfg.params["velocity_range"] + ranges = [velocity_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z", "roll", "pitch", "yaw")] + self._lin_lo = wp.vec3f(ranges[0][0], ranges[1][0], ranges[2][0]) + self._lin_hi = wp.vec3f(ranges[0][1], ranges[1][1], ranges[2][1]) + self._ang_lo = wp.vec3f(ranges[3][0], ranges[4][0], ranges[5][0]) + self._ang_hi = wp.vec3f(ranges[3][1], ranges[4][1], ranges[5][1]) + + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), + velocity_range: dict[str, tuple[float, float]], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + ) -> None: + """Add sampled linear [m/s] and angular [rad/s] root velocities. + + Args: + env: Environment containing the pushed asset. + env_mask: Boolean Warp mask selecting environments. + velocity_range: Per-axis velocity ranges [m/s or rad/s]. Parsed during initialization. + asset_cfg: Scene entity selection. Resolved during initialization. + """ + wp.launch( + kernel=_push_by_setting_velocity_kernel, + dim=env.num_envs, + inputs=[ + env_mask, + env.rng_state_wp, + self._asset.data.root_vel_w.warp, + self._velocity, + self._lin_lo, + self._lin_hi, + self._ang_lo, + self._ang_hi, + ], + device=env.device, + ) + + self._asset.write_root_velocity_to_sim_mask(root_velocity=self._velocity, env_mask=env_mask) + + +@WarpCapturable(False, reason="deprecated adapter constructs term state; use PushBySettingVelocity") def push_by_setting_velocity( - env, - env_mask: wp.array, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), velocity_range: dict[str, tuple[float, float]], asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -): - """Push the asset by setting the root velocity to a random value within the given ranges. - - Warp-first override of :func:`isaaclab.envs.mdp.events.push_by_setting_velocity`. - """ - asset: Articulation = env.scene[asset_cfg.name] - cache = _EventStateCache.for_env(env) - cache_key = (push_by_setting_velocity, id(asset), id(asset_cfg), _range_cache_key(velocity_range)) - state = cache.get(cache_key) - if state is None: - state = _PushBySettingVelocityState(env, asset, asset_cfg, velocity_range) - cache[cache_key] = state - - wp.launch( - kernel=_push_by_setting_velocity_kernel, - dim=env.num_envs, - inputs=[ - env_mask, - env.rng_state_wp, - asset.data.root_vel_w.warp, - state.velocity, - state.lin_lo, - state.lin_hi, - state.ang_lo, - state.ang_hi, - ], - device=env.device, +) -> None: + """Deprecated compatibility adapter for :class:`PushBySettingVelocity`.""" + warnings.warn( + "'push_by_setting_velocity' is deprecated; use 'PushBySettingVelocity' in event configurations.", + DeprecationWarning, + stacklevel=2, ) - - asset.write_root_velocity_to_sim_mask(root_velocity=state.velocity, env_mask=env_mask) + params = {"velocity_range": velocity_range, "asset_cfg": asset_cfg} + cfg = EventTermCfg(func=PushBySettingVelocity, mode="interval", params={}) + cfg.params = params + term = PushBySettingVelocity(cfg, env) + term(env, env_mask, **params) # --------------------------------------------------------------------------- @@ -464,56 +484,102 @@ def _reset_root_state_uniform_kernel( rng_state[env_id] = state +class ResetRootStateUniform(ManagerTermBase): + """Reset root pose and velocity using persistent Warp output buffers.""" + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv) -> None: + """Initialize persistent root-state reset state. + + Args: + cfg: Event term configuration. + env: Environment containing the reset asset. + """ + super().__init__(cfg, env) + asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self._asset: RigidObject | Articulation = env.scene[asset_cfg.name] + self._default_root_pose = self._asset.data.default_root_pose.warp + self._default_root_velocity = self._asset.data.default_root_vel.warp + self._env_origins = env.env_origins_wp + self._pose = wp.zeros(env.num_envs, dtype=wp.transformf, device=env.device) + self._velocity = wp.zeros(env.num_envs, dtype=wp.spatial_vectorf, device=env.device) + + pose_range = cfg.params["pose_range"] + pose_ranges = [pose_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z", "roll", "pitch", "yaw")] + self._pos_lo = wp.vec3f(pose_ranges[0][0], pose_ranges[1][0], pose_ranges[2][0]) + self._pos_hi = wp.vec3f(pose_ranges[0][1], pose_ranges[1][1], pose_ranges[2][1]) + self._rot_lo = wp.vec3f(pose_ranges[3][0], pose_ranges[4][0], pose_ranges[5][0]) + self._rot_hi = wp.vec3f(pose_ranges[3][1], pose_ranges[4][1], pose_ranges[5][1]) + + velocity_range = cfg.params["velocity_range"] + velocity_ranges = [velocity_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z", "roll", "pitch", "yaw")] + self._vel_lin_lo = wp.vec3f(velocity_ranges[0][0], velocity_ranges[1][0], velocity_ranges[2][0]) + self._vel_lin_hi = wp.vec3f(velocity_ranges[0][1], velocity_ranges[1][1], velocity_ranges[2][1]) + self._vel_ang_lo = wp.vec3f(velocity_ranges[3][0], velocity_ranges[4][0], velocity_ranges[5][0]) + self._vel_ang_hi = wp.vec3f(velocity_ranges[3][1], velocity_ranges[4][1], velocity_ranges[5][1]) + + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), + pose_range: dict[str, tuple[float, float]], + velocity_range: dict[str, tuple[float, float]], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + ) -> None: + """Reset root pose [m, rad] and velocity [m/s, rad/s]. + + Args: + env: Environment containing the reset asset. + env_mask: Boolean Warp mask selecting environments. + pose_range: Position and Euler-angle ranges [m or rad]. Parsed during initialization. + velocity_range: Linear and angular velocity ranges [m/s or rad/s]. Parsed during initialization. + asset_cfg: Scene entity selection. Resolved during initialization. + """ + wp.launch( + kernel=_reset_root_state_uniform_kernel, + dim=env.num_envs, + inputs=[ + env_mask, + env.rng_state_wp, + self._default_root_pose, + self._default_root_velocity, + self._env_origins, + self._pose, + self._velocity, + self._pos_lo, + self._pos_hi, + self._rot_lo, + self._rot_hi, + self._vel_lin_lo, + self._vel_lin_hi, + self._vel_ang_lo, + self._vel_ang_hi, + ], + device=env.device, + ) + + self._asset.write_root_pose_to_sim_mask(root_pose=self._pose, env_mask=env_mask) + self._asset.write_root_velocity_to_sim_mask(root_velocity=self._velocity, env_mask=env_mask) + + +@WarpCapturable(False, reason="deprecated adapter constructs term state; use ResetRootStateUniform") def reset_root_state_uniform( - env, - env_mask: wp.array, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), pose_range: dict[str, tuple[float, float]], velocity_range: dict[str, tuple[float, float]], asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -): - """Reset the asset root state to a random position and velocity uniformly within the given ranges. - - Warp-first override of :func:`isaaclab.envs.mdp.events.reset_root_state_uniform`. - """ - asset: Articulation = env.scene[asset_cfg.name] - cache = _EventStateCache.for_env(env) - cache_key = ( - reset_root_state_uniform, - id(asset), - id(asset_cfg), - _range_cache_key(pose_range), - _range_cache_key(velocity_range), +) -> None: + """Deprecated compatibility adapter for :class:`ResetRootStateUniform`.""" + warnings.warn( + "'reset_root_state_uniform' is deprecated; use 'ResetRootStateUniform' in event configurations.", + DeprecationWarning, + stacklevel=2, ) - state = cache.get(cache_key) - if state is None: - state = _ResetRootStateUniformState(env, asset, asset_cfg, pose_range, velocity_range) - cache[cache_key] = state - - wp.launch( - kernel=_reset_root_state_uniform_kernel, - dim=env.num_envs, - inputs=[ - env_mask, - env.rng_state_wp, - asset.data.default_root_pose.warp, - asset.data.default_root_vel.warp, - env.env_origins_wp, - state.pose, - state.velocity, - state.pos_lo, - state.pos_hi, - state.rot_lo, - state.rot_hi, - state.vel_lin_lo, - state.vel_lin_hi, - state.vel_ang_lo, - state.vel_ang_hi, - ], - device=env.device, - ) - - asset.write_root_pose_to_sim_mask(root_pose=state.pose, env_mask=env_mask) - asset.write_root_velocity_to_sim_mask(root_velocity=state.velocity, env_mask=env_mask) + params = {"pose_range": pose_range, "velocity_range": velocity_range, "asset_cfg": asset_cfg} + cfg = EventTermCfg(func=ResetRootStateUniform, mode="reset", params={}) + cfg.params = params + term = ResetRootStateUniform(cfg, env) + term(env, env_mask, **params) # --------------------------------------------------------------------------- @@ -567,12 +633,12 @@ def _reset_joints_by_offset_kernel( def reset_joints_by_offset( - env, - env_mask: wp.array, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), position_range: tuple[float, float], velocity_range: tuple[float, float], asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -): +) -> None: """Warp-first reset of joint state by random offsets around defaults. This overrides the stable `isaaclab.envs.mdp.events.reset_joints_by_offset` when importing @@ -663,12 +729,12 @@ def _reset_joints_by_scale_kernel( def reset_joints_by_scale( - env, - env_mask: wp.array, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), position_range: tuple[float, float], velocity_range: tuple[float, float], asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -): +) -> None: """Warp-first reset of joint state by scaling defaults with random factors.""" asset: Articulation = env.scene[asset_cfg.name] diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py index 611246aeb55d..5f94d157c440 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py @@ -91,30 +91,28 @@ def reset_extras(self) -> dict[str, torch.Tensor]: def reset( self, env_mask: wp.array(dtype=wp.bool), - *, - env_ids: Sequence[int] | torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Reset selected class terms and return persistent curriculum logging outputs. Args: env_mask: Boolean Warp mask selecting environments to reset. - env_ids: Compact environment IDs for legacy class terms that explicitly require them. Returns: Persistent scalar curriculum states keyed by their logging paths. """ env_mask = self._resolve_reset_mask(None, env_mask) + compact_env_ids = None for term_cfg, mode in zip(self._term_cfgs, self._term_modes): if isinstance(term_cfg.func, ManagerTermBase): term_cfg.func.reset(env_mask=env_mask) elif isinstance(term_cfg.func, StableManagerTermBase): if mode == "legacy_ids": - if env_ids is None: - raise RuntimeError("Legacy curriculum reset requires compact env_ids.") - legacy_env_ids = env_ids + if compact_env_ids is None: + compact_env_ids = self._compact_legacy_env_ids(env_mask) + term_env_ids = compact_env_ids else: - legacy_env_ids = slice(None) - term_cfg.func.reset(env_ids=legacy_env_ids) + term_env_ids = slice(None) + term_cfg.func.reset(env_ids=term_env_ids) return self._reset_extras @property @@ -130,31 +128,26 @@ def requires_host_boundary(self) -> bool: def compute( self, env_mask: wp.array(dtype=wp.bool), - *, - env_ids: Sequence[int] | torch.Tensor | None = None, ) -> None: """Update curriculum terms for selected environments. Args: env_mask: Boolean Warp mask selecting environments to update. - env_ids: Compact environment IDs for legacy terms that explicitly require them. - - Raises: - RuntimeError: If a legacy ID term is active and ``env_ids`` is not provided. """ env_mask = self._resolve_reset_mask(None, env_mask) self._term_states_wp.zero_() + compact_env_ids = None for term_idx, (term_cfg, mode) in enumerate(zip(self._term_cfgs, self._term_modes)): if mode == "mask": term_cfg.func(self._env, env_mask, term_cfg.out, **term_cfg.params) continue if mode == "legacy_ids": - if env_ids is None: - raise RuntimeError(f"Curriculum term '{self._term_names[term_idx]}' requires compact env_ids.") - legacy_env_ids = env_ids + if compact_env_ids is None: + compact_env_ids = self._compact_legacy_env_ids(env_mask) + term_env_ids = compact_env_ids else: - legacy_env_ids = slice(None) - state = term_cfg.func(self._env, legacy_env_ids, **term_cfg.params) + term_env_ids = slice(None) + state = term_cfg.func(self._env, term_env_ids, **term_cfg.params) if state is not None: if isinstance(state, torch.Tensor) and state.numel() != 1: raise TypeError( @@ -238,8 +231,13 @@ def _resolve_legacy_term_cfg(self, term_name: str, term_cfg: StableCurriculumTer f"The legacy curriculum term '{term_name}' expects parameters {sorted(accepted)}," f" but received {sorted(provided)}." ) - switch = getattr(self._env, "_manager_call_switch", None) - if switch is not None: - switch.register_manager_capturability(type(self).__name__, False) + graph_cache = getattr(self._env, "_warp_graph_cache", None) + if graph_cache is not None: + graph_cache.register_capturability(type(self).__name__, False) if self._env.sim.is_playing(): self._process_term_cfg_at_play(term_name, term_cfg) + + @staticmethod + def _compact_legacy_env_ids(env_mask: wp.array(dtype=wp.bool)) -> torch.Tensor: + """Materialize compact Torch IDs at the legacy curriculum boundary.""" + return wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py index 31bfceb62516..291d807ddadc 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py @@ -446,11 +446,12 @@ def _resolve_common_term_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg, f" and optional parameters: {args_with_defaults}, but received: {term_params}." ) - # register non-capturable terms with the call switch for mode=2 fallback + # Register capture safety with the environment-owned executor. The + # manager call switch only consumes this policy while it still exists. if not is_warp_capturable(term_cfg.func): - switch = getattr(self._env, "_manager_call_switch", None) - if switch is not None: - switch.register_manager_capturability(type(self).__name__, False) + graph_cache = getattr(self._env, "_warp_graph_cache", None) + if graph_cache is not None: + graph_cache.register_capturability(type(self).__name__, False) # process attributes at runtime # these properties are only resolvable once the simulation starts playing diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py b/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py index 5296a96a6bb7..4999646e402b 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py @@ -9,15 +9,23 @@ import importlib import json +import logging import os +import warnings from collections.abc import Callable +from copy import deepcopy from enum import IntEnum -from typing import Any +from typing import TYPE_CHECKING, Any from isaaclab.utils.timer import Timer from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnvCfg + +logger = logging.getLogger(__name__) + class ManagerCallMode(IntEnum): """Execution mode for manager stage calls. @@ -34,11 +42,12 @@ class ManagerCallMode(IntEnum): class ManagerCallSwitch: - """Per-manager call switch for stable/warp/captured execution. + """Compatibility router for stable and Warp manager implementations. - Routes each manager stage call through the configured execution path: - stable Python, Warp (eager), or Warp (captured CUDA graph). Optionally - wraps each call in a :class:`Timer` context for profiling. + This temporary layer selects stable or Warp manager classes and forwards + Warp stages to an environment-owned :class:`WarpGraphCache`. Execution and + capture policy therefore remain reusable after mixed stable-manager routing + is removed. Calls may optionally use a :class:`Timer` for profiling. """ # Warp eager is the correctness-first default. Capture remains an explicit @@ -70,12 +79,14 @@ class ManagerCallSwitch: def __init__( self, - cfg_source: dict | str | None = None, + cfg_source: dict[str, int] | str | None = None, *, max_modes: dict[str, int] | None = None, + graph_cache: WarpGraphCache | None = None, ): - # The graph cache is reached only for an explicit WARP_CAPTURED mode. - self._graph_cache = WarpGraphCache(enabled=True) + # The environment normally owns this durable executor. Creating one here + # keeps the compatibility router independently usable in tests and tools. + self._graph_cache = graph_cache if graph_cache is not None else WarpGraphCache(enabled=True) # Merge caller-supplied max_modes with the class-level MAX_MODE_OVERRIDES. self._max_modes = dict(self.MAX_MODE_OVERRIDES) if max_modes is not None: @@ -100,6 +111,53 @@ def invalidate_graphs(self) -> None: """Invalidate cached capture graphs and their cached return values.""" self._graph_cache.invalidate() + def apply_term_cfg_profile(self, cfg: ManagerBasedEnvCfg) -> None: + """Apply the deprecated mixed stable-manager configuration profile. + + This compatibility hook is intentionally isolated from the Warp + environment. It can be removed together with stable manager routing + once all task configurations use the Warp frontend natively. + + Args: + cfg: Experimental environment configuration to update in place. + """ + manager_to_cfg_attr = { + "ActionManager": "actions", + "ObservationManager": "observations", + "EventManager": "events", + "RecorderManager": "recorders", + "CommandManager": "commands", + "TerminationManager": "terminations", + "RewardManager": "rewards", + "CurriculumManager": "curriculum", + } + stable_managers = [ + name for name in manager_to_cfg_attr if self.get_mode_for_manager(name) == ManagerCallMode.STABLE + ] + if not stable_managers: + return + + warnings.warn( + "Selecting STABLE managers inside ManagerBasedEnvWarp is deprecated. Use ManagerBasedEnv for Torch " + "managers or WARP_NOT_CAPTURED for the Warp frontend.", + DeprecationWarning, + stacklevel=2, + ) + stable_cfg = self._resolve_stable_cfg_counterpart(cfg) + if stable_cfg is None: + logger.warning( + "Stable managers requested (%s), but no stable cfg counterpart could be resolved. Keeping " + "experimental term configs.", + ", ".join(stable_managers), + ) + return + + for manager_name, cfg_attr in manager_to_cfg_attr.items(): + if self.get_mode_for_manager(manager_name) != ManagerCallMode.STABLE: + continue + if hasattr(cfg, cfg_attr) and hasattr(stable_cfg, cfg_attr): + setattr(cfg, cfg_attr, deepcopy(getattr(stable_cfg, cfg_attr))) + # ------------------------------------------------------------------ # Stage dispatch # ------------------------------------------------------------------ @@ -131,11 +189,16 @@ def call( The stage result, optionally transformed by :paramref:`_output`. """ with Timer(name=stage, msg=f"{stage} took:", enable=_timer, time_unit="us"): - mode = self.get_mode_for_manager(self._manager_name_from_stage(stage)) - if mode == ManagerCallMode.WARP_CAPTURED: - result = self._graph_cache.capture_or_replay(stage, fn, args=args, kwargs=kwargs) - else: - result = fn(*args, **kwargs) + manager_name = self._manager_name_from_stage(stage) + mode = self.get_mode_for_manager(manager_name) + result = self._graph_cache.call( + stage, + fn, + args=args, + kwargs=kwargs, + capture=mode == ManagerCallMode.WARP_CAPTURED, + group=manager_name, + ) return _output(result) if _output is not None else result def call_stage( @@ -161,7 +224,7 @@ def call_stage( stage: Stage identifier in the form ``"ManagerName_function_name"``. warp_call: Call spec for the warp path (eager or captured). stable_call: Call spec for the stable (torch) path. Defaults to ``None``. - timer: Whether to wrap execution in a :class:`Timer`. Defaults to ``True`` + timer: Whether to wrap execution in a :class:`Timer`. Defaults to ``False`` (controlled by the global :attr:`Timer.enable` class-level toggle). Pass a module-level flag like ``TIMER_ENABLED_STEP`` to make timing conditional on that flag. @@ -169,28 +232,21 @@ def call_stage( Returns: The (possibly transformed) return value of the stage. """ - with Timer(name=stage, msg=f"{stage} took:", enable=timer, time_unit="us"): - return self._dispatch(stage, stable_call, warp_call) - - def _dispatch( - self, - stage: str, - stable_call: dict[str, Any] | None, - warp_call: dict[str, Any], - ) -> Any: - """Select call path based on mode, execute, and apply output.""" mode = self.get_mode_for_manager(self._manager_name_from_stage(stage)) if mode == ManagerCallMode.STABLE: if stable_call is None: raise ValueError(f"Stage '{stage}' is configured as STABLE (mode=0) but no stable_call was provided.") - call, result = stable_call, self._run_call(stable_call) - elif mode == ManagerCallMode.WARP_CAPTURED: - call, result = warp_call, self._wp_capture_or_launch(stage, warp_call) + call = stable_call else: - call, result = warp_call, self._run_call(warp_call) - - output_fn = call.get("output") - return output_fn(result) if output_fn is not None else result + call = warp_call + return self.call( + stage, + call["fn"], + *call.get("args", ()), + _output=call.get("output"), + _timer=timer, + **call.get("kwargs", {}), + ) # ------------------------------------------------------------------ # Manager resolution @@ -212,6 +268,8 @@ def get_mode_for_manager(self, manager_name: str) -> ManagerCallMode: cap = self._max_modes.get(manager_name) if cap is not None: mode_value = min(mode_value, cap) + if mode_value == ManagerCallMode.WARP_CAPTURED and not self._graph_cache.is_capturable(manager_name): + mode_value = ManagerCallMode.WARP_NOT_CAPTURED return ManagerCallMode(mode_value) def resolve_manager_class(self, manager_name: str, mode_override: ManagerCallMode | int | None = None) -> type: @@ -229,34 +287,13 @@ def register_manager_capturability(self, manager_name: str, capturable: bool) -> Called by :class:`ManagerBase` during term preparation when a term is decorated with ``@warp_capturable(False)``. """ - if not capturable: - self._max_modes[manager_name] = min( - self._max_modes.get(manager_name, ManagerCallMode.WARP_CAPTURED), - ManagerCallMode.WARP_NOT_CAPTURED, - ) + self._graph_cache.register_capturability(manager_name, capturable) # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ - def _run_call(self, call: dict[str, Any]) -> Any: - """Execute a single call spec eagerly.""" - return call["fn"](*call.get("args", ()), **call.get("kwargs", {})) - - def _wp_capture_or_launch(self, stage: str, call: dict[str, Any]) -> Any: - """Capture Warp CUDA graph on first call, then replay. - - Delegates to :class:`WarpGraphCache` which handles warm-up, capture, - caching the return value, and replay. - """ - return self._graph_cache.capture_or_replay( - stage, - call["fn"], - args=call.get("args", ()), - kwargs=call.get("kwargs", {}), - ) - - def _load_cfg(self, cfg_source: dict | str | None) -> dict[str, int]: + def _load_cfg(self, cfg_source: dict[str, int] | str | None) -> dict[str, int]: if cfg_source is None: cfg = dict(self.DEFAULT_CONFIG) elif isinstance(cfg_source, dict): @@ -298,3 +335,38 @@ def _load_cfg(self, cfg_source: dict | str | None) -> dict[str, int]: cfg[name] = max_mode return cfg + + @staticmethod + def _resolve_stable_cfg_counterpart(cfg: ManagerBasedEnvCfg) -> ManagerBasedEnvCfg | None: + """Resolve the legacy stable task configuration counterpart.""" + cfg_cls = cfg.__class__ + cfg_module_name = cfg_cls.__module__ + if "isaaclab_tasks_experimental" not in cfg_module_name: + return None + + stable_module_name = cfg_module_name.replace("isaaclab_tasks_experimental", "isaaclab_tasks", 1) + try: + stable_module = importlib.import_module(stable_module_name) + except Exception as exc: + logger.warning( + "Failed to import stable task cfg module '%s' for manager_call_config stable mode: %s", + stable_module_name, + exc, + ) + return None + + stable_cfg_cls = getattr(stable_module, cfg_cls.__name__, None) + if stable_cfg_cls is None: + logger.warning("Stable task cfg class '%s' not found in module '%s'.", cfg_cls.__name__, stable_module_name) + return None + + try: + return stable_cfg_cls() + except Exception as exc: + logger.warning( + "Failed to instantiate stable task cfg '%s.%s': %s", + stable_module_name, + cfg_cls.__name__, + exc, + ) + return None diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py index 047e0a3fd0e5..dd9f2d2f7229 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py @@ -12,7 +12,7 @@ class WarpGraphCache: - """Run Warp stages eagerly or cache CUDA graphs by stage name. + """Execute Warp stages eagerly or through cached CUDA graphs. On the very first call for a given stage, an **eager warm-up** run executes *before* graph capture. This lets one-time initialisation @@ -27,9 +27,9 @@ class WarpGraphCache: Usage:: cache = WarpGraphCache() - result = cache.capture_or_replay("my_stage", my_warp_function) + result = cache.call("my_stage", my_warp_function) # uncaptured work here ... - result2 = cache.capture_or_replay("my_stage_post", my_other_function) + result2 = cache.call("my_stage_post", my_other_function) """ def __init__(self, *, enabled: bool = True): @@ -41,12 +41,46 @@ def __init__(self, *, enabled: bool = True): self._enabled = enabled self._graphs: dict[str, Any] = {} self._results: dict[str, Any] = {} + self._capturable: dict[str, bool] = {} + + def call( + self, + stage: str, + fn: Callable[..., Any], + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + *, + capture: bool | None = None, + group: str | None = None, + ) -> Any: + """Execute a stage using its registered capture policy. + + Args: + stage: Unique stage identifier. + fn: Callable implementing the stage. + args: Positional arguments forwarded to :paramref:`fn`. + kwargs: Keyword arguments forwarded to :paramref:`fn`. + capture: Whether capture is requested for this call. When omitted, + follows the cache-wide :attr:`_enabled` setting. + group: Optional capturability group. This lets an owner register one + decision for several stages, such as all calls on one manager. + + Returns: + The eager stage result or cached capture result. + """ + if kwargs is None: + kwargs = {} + capture_requested = self._enabled if capture is None else capture + capture_key = stage if group is None else group + if not self._enabled or not capture_requested or not self.is_capturable(capture_key): + return fn(*args, **kwargs) + return self._capture_or_replay(stage, fn, args, kwargs) def capture_or_replay( self, stage: str, fn: Callable[..., Any], - args: tuple = (), + args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None, ) -> Any: """Capture *fn* into a CUDA graph on the first call, then replay. @@ -62,10 +96,42 @@ def capture_or_replay( Returns: The eager result when disabled, otherwise the cached result from capture. """ - if kwargs is None: - kwargs = {} - if not self._enabled: - return fn(*args, **kwargs) + return self.call(stage, fn, args, kwargs, capture=True) + + def register_capturability(self, key: str, capturable: bool) -> None: + """Register capture eligibility for a stage or owner group. + + Repeated registrations are conservative: once a key is non-capturable, + later registrations cannot silently re-enable it. + + Args: + key: Stage or group identifier. + capturable: Whether every operation registered under the key is safe + during CUDA graph capture. + """ + self._capturable[key] = self._capturable.get(key, True) and capturable + + def is_capturable(self, key: str) -> bool: + """Return whether a stage or group is eligible for capture.""" + return self._capturable.get(key, True) + + def invalidate(self, stage: str | None = None) -> None: + """Drop cached graph(s). If *stage* is ``None``, drop all.""" + if stage is None: + self._graphs.clear() + self._results.clear() + else: + self._graphs.pop(stage, None) + self._results.pop(stage, None) + + def _capture_or_replay( + self, + stage: str, + fn: Callable[..., Any], + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> Any: + """Capture a stage on first use and replay it afterward.""" graph = self._graphs.get(stage) if graph is not None: wp.capture_launch(graph) @@ -78,12 +144,3 @@ def capture_or_replay( self._graphs[stage] = capture.graph self._results[stage] = result return result - - def invalidate(self, stage: str | None = None) -> None: - """Drop cached graph(s). If *stage* is ``None``, drop all.""" - if stage is None: - self._graphs.clear() - self._results.clear() - else: - self._graphs.pop(stage, None) - self._results.pop(stage, None) diff --git a/source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py b/source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py index f6dca451673e..c07980f0b6c3 100644 --- a/source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py +++ b/source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py @@ -22,9 +22,7 @@ ) from isaaclab_experimental.managers import CommandManager, CommandTerm -from isaaclab.envs.mdp.commands import UniformPoseCommand as StableUniformPoseCommand from isaaclab.envs.mdp.commands import UniformPoseCommandCfg as StableUniformPoseCommandCfg -from isaaclab.envs.mdp.commands import UniformVelocityCommand as StableUniformVelocityCommand from isaaclab.envs.mdp.commands import UniformVelocityCommandCfg as StableUniformVelocityCommandCfg from isaaclab.utils.math import combine_frame_transforms, compute_pose_error, quat_from_euler_xyz, quat_unique @@ -43,7 +41,11 @@ class _Robot: def __init__(self, num_envs: int): identity_quat = np.zeros((num_envs, 4), dtype=np.float32) identity_quat[:, 3] = 1.0 + identity_pose = np.zeros((num_envs, 7), dtype=np.float32) + identity_pose[:, 6] = 1.0 self.data = SimpleNamespace( + root_pose_w=_ArrayProxy(identity_pose, wp.transformf), + root_vel_w=_ArrayProxy(np.zeros((num_envs, 6), dtype=np.float32), wp.spatial_vectorf), root_pos_w=_ArrayProxy(np.zeros((num_envs, 3), dtype=np.float32), wp.vec3f), root_quat_w=_ArrayProxy(identity_quat, wp.quatf), root_lin_vel_b=_ArrayProxy(np.zeros((num_envs, 3), dtype=np.float32), wp.vec3f), @@ -115,20 +117,19 @@ def _pose_cfg() -> UniformPoseCommandCfg: def test_configs_exports_and_command_accessors_preserve_shared_contracts(): - """Configs should inherit stable fields and terms should expose shared debug UI and persistent storage.""" + """Configs should inherit stable fields and terms should expose debug UI and persistent storage.""" assert issubclass(UniformVelocityCommandCfg, StableUniformVelocityCommandCfg) assert issubclass(UniformPoseCommandCfg, StableUniformPoseCommandCfg) assert experimental_mdp.UniformVelocityCommandCfg is UniformVelocityCommandCfg assert experimental_mdp.UniformPoseCommandCfg is UniformPoseCommandCfg assert str(_velocity_cfg().class_type).endswith(".commands.velocity_command:UniformVelocityCommand") assert str(_pose_cfg().class_type).endswith(".commands.pose_command:UniformPoseCommand") - assert UniformVelocityCommand._set_debug_vis_impl is StableUniformVelocityCommand._set_debug_vis_impl - assert UniformPoseCommand._set_debug_vis_impl is StableUniformPoseCommand._set_debug_vis_impl + assert UniformVelocityCommand._set_debug_vis_impl.__module__.endswith(".commands._debug_vis") + assert UniformPoseCommand._set_debug_vis_impl.__module__.endswith(".commands._debug_vis") env = _Env() velocity_term = UniformVelocityCommand(_velocity_cfg(), env) pose_term = UniformPoseCommand(_pose_cfg(), env) - stable_pose_term = StableUniformPoseCommand(_pose_cfg(), env) manager = CommandManager.__new__(CommandManager) manager._terms = {"velocity": velocity_term, "pose": pose_term} @@ -140,8 +141,6 @@ def test_configs_exports_and_command_accessors_preserve_shared_contracts(): assert manager.get_command("pose") is pose_term.command torch.testing.assert_close(pose_term.command[:, :6], torch.zeros(4, 6)) torch.testing.assert_close(pose_term.command[:, 6], torch.ones(4)) - torch.testing.assert_close(stable_pose_term.command[:, :6], torch.zeros(4, 6)) - torch.testing.assert_close(stable_pose_term.command[:, 6], torch.ones(4)) def test_debug_visualization_uses_current_simulation_registry(): @@ -174,10 +173,10 @@ def test_uniform_velocity_command_updates_and_resets_selected_rows(): command_wp = term.command_wp term.command[:] = torch.tensor([[1.0, 2.0, 0.3], [2.0, -1.0, -0.2], [0.0, 0.5, 0.1], [-1.0, -2.0, 0.0]]) - env.scene["robot"].data.root_lin_vel_b.torch[:] = torch.tensor( + env.scene["robot"].data.root_vel_w.torch[:, :3] = torch.tensor( [[0.0, 0.0, 0.0], [1.0, -1.0, 0.0], [0.0, 0.0, 0.0], [-2.0, -2.0, 0.0]] ) - env.scene["robot"].data.root_ang_vel_b.torch[:, 2] = torch.tensor([0.1, -0.1, 0.4, 0.0]) + env.scene["robot"].data.root_vel_w.torch[:, 5] = torch.tensor([0.1, -0.1, 0.4, 0.0]) term._update_metrics() wp.synchronize() torch.testing.assert_close( @@ -213,7 +212,6 @@ def test_uniform_velocity_command_updates_and_resets_selected_rows(): term.heading_target[:] = torch.tensor([np.pi / 2, 0.0, np.pi / 2, 0.0]) term.is_heading_env[:] = torch.tensor([True, False, True, False]) term.is_standing_env[:] = torch.tensor([False, False, True, True]) - env.scene["robot"].data.heading_w.torch.zero_() term._update_command() wp.synchronize() torch.testing.assert_close( @@ -222,6 +220,40 @@ def test_uniform_velocity_command_updates_and_resets_selected_rows(): ) +def test_uniform_velocity_command_replay_reads_current_root_state(): + """Captured command kernels should read current Tier-1 root state on every replay.""" + env = _Env() + cfg = _velocity_cfg() + cfg.ranges.ang_vel_z = (-2.0, 2.0) + term = UniformVelocityCommand(cfg, env) + robot_data = env.scene["robot"].data + + term.command.zero_() + term.command[:, 0] = 1.0 + term.heading_target.fill_(np.pi / 2) + term.is_heading_env.fill_(True) + term.is_standing_env.fill_(False) + + with wp.ScopedCapture(device=env.device) as metrics_capture: + term._update_metrics() + with wp.ScopedCapture(device=env.device) as command_capture: + term._update_command() + + # Identity pose with matching world-frame velocity has zero body-frame XY error. + robot_data.root_vel_w.torch[:, 0] = 1.0 + wp.capture_launch(metrics_capture.graph) + wp.synchronize() + torch.testing.assert_close(wp.to_torch(term._error_xy_sum_wp), torch.zeros(env.num_envs)) + + # A 90-degree root yaw exactly matches the heading target, so replay should command zero yaw rate. + half_sqrt = np.sqrt(0.5) + robot_data.root_pose_w.torch[:, 5] = half_sqrt + robot_data.root_pose_w.torch[:, 6] = half_sqrt + wp.capture_launch(command_capture.graph) + wp.synchronize() + torch.testing.assert_close(term.command[:, 2], torch.zeros(env.num_envs), atol=1.0e-6, rtol=1.0e-6) + + def test_uniform_pose_command_matches_stable_math_and_tracks_sticky_success(): """Pose commands should use xyzw math, update metrics, and reset sticky success by mask.""" env = _Env() diff --git a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py index a549592a9ada..7cb1e8b073be 100644 --- a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py +++ b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py @@ -383,7 +383,12 @@ class MockScene: def __init__(self, assets: dict, env_origins, sensors=None): self._assets = assets - self.env_origins = env_origins + if isinstance(env_origins, wp.array): + self.env_origins_wp = env_origins + self.env_origins = wp.to_torch(env_origins) + else: + self.env_origins = env_origins + self.env_origins_wp = wp.from_torch(env_origins, dtype=wp.vec3f) self.sensors = sensors or {} self.articulations = dict(assets) self.rigid_objects = {} diff --git a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py index ecdfb67549c4..402ea2bf44ec 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py @@ -17,7 +17,7 @@ pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") import isaaclab_experimental.envs.mdp.events as warp_evt -from isaaclab_experimental.managers import SceneEntityCfg +from isaaclab_experimental.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg from parity_helpers import ( DEVICE, NUM_ACTIONS, @@ -81,6 +81,7 @@ class _Env: env.action_manager = MockActionManagerWarp(action_wp[0], action_wp[1]) env.num_envs = NUM_ENVS env.device = DEVICE + env.env_origins_wp = scene.env_origins_wp env.episode_length_buf = episode_length_buf env.step_dt = 0.02 env.max_episode_length_s = 10.0 @@ -125,11 +126,24 @@ class _Env: env.scene = MockScene({"robot": asset}, origins) env.num_envs = NUM_ENVS env.device = DEVICE - env.env_origins_wp = origins + env.env_origins_wp = env.scene.env_origins_wp env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + seed, device=DEVICE) return env, data, asset +def _make_event_term( + term_type: type[ManagerTermBase], + env: object, + *, + mode: str = "reset", + **params: object, +) -> ManagerTermBase: + """Create a persistent Warp event term for a mock environment.""" + params.setdefault("asset_cfg", SceneEntityCfg("robot")) + cfg = EventTermCfg(func=term_type, mode=mode, params=params) + return term_type(cfg, env) + + # ============================================================================ # Event parity tests: deterministic (zero-width range) warp vs stable # ============================================================================ @@ -278,10 +292,16 @@ def test_push_by_setting_velocity(self, warp_env, art_data, all_joints_cfg): "pitch": (0.0, 0.0), "yaw": (0.0, 0.0), } + term = _make_event_term( + warp_evt.PushBySettingVelocity, + warp_env, + mode="interval", + velocity_range=zero_range, + ) - warp_evt.push_by_setting_velocity(warp_env, mask, velocity_range=zero_range) + term(warp_env, mask, **term.cfg.params) with wp.ScopedCapture() as cap: - warp_evt.push_by_setting_velocity(warp_env, mask, velocity_range=zero_range) + term(warp_env, mask, **term.cfg.params) # Mutate root_vel_w new_vel = np.tile([1.0, 2.0, 3.0, 0.1, 0.2, 0.3], (NUM_ENVS, 1)).astype(np.float32) @@ -304,11 +324,17 @@ def test_apply_external_force_torque(self, warp_env, art_data, all_joints_cfg): lambda **kwargs: captured.update(kwargs) ) zero_range = (0.0, 0.0) + term = _make_event_term( + warp_evt.ApplyExternalForceTorque, + warp_env, + force_range=zero_range, + torque_range=zero_range, + ) # Zero-range: forces and torques should be zero - warp_evt.apply_external_force_torque(warp_env, mask, force_range=zero_range, torque_range=zero_range) + term(warp_env, mask, **term.cfg.params) with wp.ScopedCapture() as cap: - warp_evt.apply_external_force_torque(warp_env, mask, force_range=zero_range, torque_range=zero_range) + term(warp_env, mask, **term.cfg.params) wp.capture_launch(cap.graph) wp.synchronize() @@ -350,6 +376,29 @@ def test_reset_joints_mask_selectivity(self, warp_env, art_data, all_joints_cfg) class TestEventStateOwnership: """Verify event scratch and parsed ranges are owned by one environment configuration.""" + def test_com_term_is_marked_non_capturable(self): + assert warp_evt.RandomizeRigidBodyCom._warp_capturable is False + + def test_lowercase_adapter_warns_and_forwards(self): + env, data, asset = _make_event_env(88) + copy_np_to_wp(data.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + captured = {} + asset.write_root_velocity_to_sim_mask = lambda **kwargs: captured.update(kwargs) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + asset_cfg = SceneEntityCfg("robot") + asset_cfg.body_ids_wp = wp.array([0], dtype=wp.int32, device=DEVICE) + + with pytest.warns(DeprecationWarning, match="PushBySettingVelocity"): + warp_evt.push_by_setting_velocity( + env, + env_mask, + velocity_range={"x": (1.0, 1.0)}, + asset_cfg=asset_cfg, + ) + wp.synchronize() + + assert_close(wp.to_torch(captured["root_velocity"])[:, 0], torch.ones(NUM_ENVS, device=DEVICE)) + def test_push_state_is_not_shared_between_environments(self): env_a, data_a, asset_a = _make_event_env(101) env_b, data_b, asset_b = _make_event_env(202) @@ -362,9 +411,11 @@ def test_push_state_is_not_shared_between_environments(self): env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) range_a = {"x": (1.0, 1.0)} range_b = {"x": (2.0, 2.0)} + term_a = _make_event_term(warp_evt.PushBySettingVelocity, env_a, mode="interval", velocity_range=range_a) + term_b = _make_event_term(warp_evt.PushBySettingVelocity, env_b, mode="interval", velocity_range=range_b) - warp_evt.push_by_setting_velocity(env_a, env_mask, velocity_range=range_a) - warp_evt.push_by_setting_velocity(env_b, env_mask, velocity_range=range_b) + term_a(env_a, env_mask, **term_a.cfg.params) + term_b(env_b, env_mask, **term_b.cfg.params) wp.synchronize() velocity_a = captured_a["root_velocity"] @@ -381,9 +432,11 @@ def test_push_state_is_not_shared_between_configurations(self): env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) range_a = {"x": (1.0, 1.0)} range_b = {"x": (2.0, 2.0)} + term_a = _make_event_term(warp_evt.PushBySettingVelocity, env, mode="interval", velocity_range=range_a) + term_b = _make_event_term(warp_evt.PushBySettingVelocity, env, mode="interval", velocity_range=range_b) - warp_evt.push_by_setting_velocity(env, env_mask, velocity_range=range_a) - warp_evt.push_by_setting_velocity(env, env_mask, velocity_range=range_b) + term_a(env, env_mask, **term_a.cfg.params) + term_b(env, env_mask, **term_b.cfg.params) wp.synchronize() velocity_a = captured[0]["root_velocity"] @@ -392,23 +445,29 @@ def test_push_state_is_not_shared_between_configurations(self): assert_close(wp.to_torch(velocity_a)[:, 0], torch.ones(NUM_ENVS, device=DEVICE)) assert_close(wp.to_torch(velocity_b)[:, 0], torch.full((NUM_ENVS,), 2.0, device=DEVICE)) - def test_push_state_updates_after_in_place_range_mutation(self): + def test_push_term_snapshots_range_during_initialization(self): env, data, asset = _make_event_env(277) copy_np_to_wp(data.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32)) captured = {} asset.write_root_velocity_to_sim_mask = lambda **kwargs: captured.update(kwargs) env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) velocity_range = {"x": (1.0, 1.0)} + term = _make_event_term( + warp_evt.PushBySettingVelocity, + env, + mode="interval", + velocity_range=velocity_range, + ) - warp_evt.push_by_setting_velocity(env, env_mask, velocity_range=velocity_range) + term(env, env_mask, **term.cfg.params) wp.synchronize() assert_close(wp.to_torch(captured["root_velocity"])[:, 0], torch.ones(NUM_ENVS, device=DEVICE)) velocity_range["x"] = (4.0, 4.0) - warp_evt.push_by_setting_velocity(env, env_mask, velocity_range=velocity_range) + term(env, env_mask, **term.cfg.params) wp.synchronize() - assert_close(wp.to_torch(captured["root_velocity"])[:, 0], torch.full((NUM_ENVS,), 4.0, device=DEVICE)) + assert_close(wp.to_torch(captured["root_velocity"])[:, 0], torch.ones(NUM_ENVS, device=DEVICE)) def test_push_state_respects_sparse_mask(self): env, data, asset = _make_event_env(303) @@ -417,8 +476,14 @@ def test_push_state_respects_sparse_mask(self): asset.write_root_velocity_to_sim_mask = lambda **kwargs: captured.update(kwargs) mask_np = np.arange(NUM_ENVS) % 3 == 0 env_mask = wp.array(mask_np, dtype=wp.bool, device=DEVICE) + term = _make_event_term( + warp_evt.PushBySettingVelocity, + env, + mode="interval", + velocity_range={"x": (3.0, 3.0)}, + ) - warp_evt.push_by_setting_velocity(env, env_mask, velocity_range={"x": (3.0, 3.0)}) + term(env, env_mask, **term.cfg.params) wp.synchronize() velocity = wp.to_torch(captured["root_velocity"]) @@ -437,9 +502,21 @@ def test_external_wrench_state_is_not_shared_between_environments(self): force_range_a = (1.0, 1.0) force_range_b = (2.0, 2.0) torque_range = (0.0, 0.0) + term_a = _make_event_term( + warp_evt.ApplyExternalForceTorque, + env_a, + force_range=force_range_a, + torque_range=torque_range, + ) + term_b = _make_event_term( + warp_evt.ApplyExternalForceTorque, + env_b, + force_range=force_range_b, + torque_range=torque_range, + ) - warp_evt.apply_external_force_torque(env_a, env_mask, force_range=force_range_a, torque_range=torque_range) - warp_evt.apply_external_force_torque(env_b, env_mask, force_range=force_range_b, torque_range=torque_range) + term_a(env_a, env_mask, **term_a.cfg.params) + term_b(env_b, env_mask, **term_b.cfg.params) wp.synchronize() forces_a = captured_a["forces"] @@ -454,14 +531,15 @@ def test_external_wrench_respects_body_selection(self): asset.permanent_wrench_composer.set_forces_and_torques_mask = lambda **kwargs: captured.update(kwargs) asset_cfg = SceneEntityCfg("robot", body_ids=[1]) env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) - - warp_evt.apply_external_force_torque( + term = _make_event_term( + warp_evt.ApplyExternalForceTorque, env, - env_mask, force_range=(2.0, 2.0), torque_range=(0.0, 0.0), asset_cfg=asset_cfg, ) + + term(env, env_mask, **term.cfg.params) wp.synchronize() expected_body_mask = torch.tensor([False, True, False], dtype=torch.bool, device=DEVICE) @@ -488,9 +566,21 @@ def test_root_reset_state_is_not_shared_between_environments(self): pose_range_a = {"x": (1.0, 1.0)} pose_range_b = {"x": (2.0, 2.0)} velocity_range = {} + term_a = _make_event_term( + warp_evt.ResetRootStateUniform, + env_a, + pose_range=pose_range_a, + velocity_range=velocity_range, + ) + term_b = _make_event_term( + warp_evt.ResetRootStateUniform, + env_b, + pose_range=pose_range_b, + velocity_range=velocity_range, + ) - warp_evt.reset_root_state_uniform(env_a, env_mask, pose_range=pose_range_a, velocity_range=velocity_range) - warp_evt.reset_root_state_uniform(env_b, env_mask, pose_range=pose_range_b, velocity_range=velocity_range) + term_a(env_a, env_mask, **term_a.cfg.params) + term_b(env_b, env_mask, **term_b.cfg.params) wp.synchronize() pose_a = captured_a["root_pose"] @@ -508,36 +598,51 @@ def test_com_ranges_are_not_shared_between_environments(self): asset_b.set_coms_mask = lambda **kwargs: None asset_cfg_a = SceneEntityCfg("robot") asset_cfg_b = SceneEntityCfg("robot") - asset_cfg_a.body_ids_wp = wp.array([0, 1], dtype=wp.int32, device=DEVICE) - asset_cfg_b.body_ids_wp = wp.array([0, 1], dtype=wp.int32, device=DEVICE) env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) com_range_a = {"x": (1.0, 1.0)} com_range_b = {"x": (2.0, 2.0)} + term_a = _make_event_term( + warp_evt.RandomizeRigidBodyCom, + env_a, + com_range=com_range_a, + asset_cfg=asset_cfg_a, + ) + term_b = _make_event_term( + warp_evt.RandomizeRigidBodyCom, + env_b, + com_range=com_range_b, + asset_cfg=asset_cfg_b, + ) - warp_evt.randomize_rigid_body_com(env_a, env_mask, com_range=com_range_a, asset_cfg=asset_cfg_a) - warp_evt.randomize_rigid_body_com(env_b, env_mask, com_range=com_range_b, asset_cfg=asset_cfg_b) + term_a(env_a, env_mask, **term_a.cfg.params) + term_b(env_b, env_mask, **term_b.cfg.params) wp.synchronize() - assert env_a._warp_event_state_cache is not env_b._warp_event_state_cache + assert term_a._default_com.ptr != term_b._default_com.ptr assert_close(data_a.body_com_pos_b.torch[..., 0], torch.ones((NUM_ENVS, 2), device=DEVICE)) assert_close(data_b.body_com_pos_b.torch[..., 0], torch.full((NUM_ENVS, 2), 2.0, device=DEVICE)) - def test_com_randomization_uses_cached_baseline(self): + def test_com_randomization_uses_persistent_baseline(self): env, data, asset = _make_event_env(1001, num_bodies=2) baseline = np.zeros((NUM_ENVS, 2, 3), dtype=np.float32) baseline[..., 0] = 0.25 copy_np_to_wp(data.body_com_pos_b, baseline) asset.set_coms_mask = lambda **kwargs: None asset_cfg = SceneEntityCfg("robot", body_ids=[0, 1]) - asset_cfg.body_ids_wp = wp.array([0, 1], dtype=wp.int32, device=DEVICE) env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) com_range = {"x": (1.0, 1.0)} + term = _make_event_term( + warp_evt.RandomizeRigidBodyCom, + env, + com_range=com_range, + asset_cfg=asset_cfg, + ) - warp_evt.randomize_rigid_body_com(env, env_mask, com_range=com_range, asset_cfg=asset_cfg) + term(env, env_mask, **term.cfg.params) wp.synchronize() assert_close(data.body_com_pos_b.torch[..., 0], torch.full((NUM_ENVS, 2), 1.25, device=DEVICE)) - warp_evt.randomize_rigid_body_com(env, env_mask, com_range=com_range, asset_cfg=asset_cfg) + term(env, env_mask, **term.cfg.params) wp.synchronize() assert_close(data.body_com_pos_b.torch[..., 0], torch.full((NUM_ENVS, 2), 1.25, device=DEVICE)) @@ -547,15 +652,15 @@ def test_com_randomization_broadcasts_one_offset_across_selected_bodies(self): copy_np_to_wp(data.body_com_pos_b, baseline) asset.set_coms_mask = lambda **kwargs: None asset_cfg = SceneEntityCfg("robot", body_ids=[0, 2]) - asset_cfg.body_ids_wp = wp.array([0, 2], dtype=wp.int32, device=DEVICE) env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) - - warp_evt.randomize_rigid_body_com( + term = _make_event_term( + warp_evt.RandomizeRigidBodyCom, env, - env_mask, com_range={"x": (-1.0, 1.0), "y": (-2.0, 2.0), "z": (-3.0, 3.0)}, asset_cfg=asset_cfg, ) + + term(env, env_mask, **term.cfg.params) wp.synchronize() assert_close(data.body_com_pos_b.torch[:, 0], data.body_com_pos_b.torch[:, 2]) diff --git a/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py b/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py index dbd56df36754..cc38406ad190 100644 --- a/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py +++ b/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py @@ -11,13 +11,43 @@ from unittest.mock import Mock import pytest +import torch import warp as wp from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp +from isaaclab.scene import InteractiveScene +from isaaclab.terrains import TerrainImporterCfg +from isaaclab.utils.warp import ProxyArray + class TestInteractiveSceneWarp: """Tests for :class:`InteractiveSceneWarp`.""" + def test_frontend_uses_canonical_terrain_storage_without_mutating_cfg( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Warp scenes should inherit canonical origin storage without replacing the configured class.""" + terrain_cfg = TerrainImporterCfg(prim_path="/World/Ground", terrain_type="plane") + cfg = SimpleNamespace(terrain=terrain_cfg) + class_type = terrain_cfg.class_type + origins_wp = wp.zeros(2, dtype=wp.vec3f, device="cpu") + terrain = SimpleNamespace(env_origins=ProxyArray(origins_wp)) + + def initialize_scene(scene: InteractiveScene, scene_cfg: SimpleNamespace) -> None: + scene.cfg = scene_cfg + scene._terrain = terrain + + monkeypatch.setattr(InteractiveScene, "__init__", initialize_scene) + + scene = InteractiveSceneWarp(cfg) + + assert terrain_cfg.class_type is class_type + assert isinstance(scene.env_origins, torch.Tensor) + assert scene.env_origins is terrain.env_origins.torch + assert scene.env_origins_wp is origins_wp + assert InteractiveSceneWarp.env_origins is InteractiveScene.env_origins + assert InteractiveSceneWarp.env_origins_wp is InteractiveScene.env_origins_wp + def test_mask_reset_stays_mask_based_for_supported_entities(self): """Mask-based reset should not pass environment IDs to Warp-capable entities.""" scene = InteractiveSceneWarp.__new__(InteractiveSceneWarp) diff --git a/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py index 7fbce5fe4899..49224c43708b 100644 --- a/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py @@ -15,27 +15,22 @@ import torch import warp as wp from isaaclab_experimental.envs.manager_based_rl_env_warp import ManagerBasedRLEnvWarp -from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode +from isaaclab_experimental.utils.manager_call_switch import ManagerCallSwitch def test_warp_environment_deprecates_and_preserves_stable_manager_profile() -> None: """The legacy stable-manager profile should warn while retaining its config bridge.""" - env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) - env.cfg = SimpleNamespace(rewards="warp_rewards", actions="warp_actions") - env._manager_call_switch = SimpleNamespace( - get_mode_for_manager=lambda name: ( - ManagerCallMode.STABLE if name == "RewardManager" else ManagerCallMode.WARP_NOT_CAPTURED - ), - ) - env._resolve_stable_cfg_counterpart = Mock( + cfg = SimpleNamespace(rewards="warp_rewards", actions="warp_actions") + switch = ManagerCallSwitch(cfg_source={"default": 1, "RewardManager": 0}) + switch._resolve_stable_cfg_counterpart = Mock( return_value=SimpleNamespace(rewards="stable_rewards", actions="stable_actions") ) with pytest.warns(DeprecationWarning, match="STABLE managers"): - env._apply_manager_term_cfg_profile() + switch.apply_term_cfg_profile(cfg) - assert env.cfg.rewards == "stable_rewards" - assert env.cfg.actions == "warp_actions" + assert cfg.rewards == "stable_rewards" + assert cfg.actions == "warp_actions" def test_reset_without_host_consumers_does_not_compact_ids() -> None: @@ -47,12 +42,14 @@ def test_reset_without_host_consumers_does_not_compact_ids() -> None: env.reset_buf.nonzero.side_effect = AssertionError("reset IDs should not be materialized") env._reset_requires_host_selection = Mock(return_value=False) env._reset_idx = Mock() + env.has_rtx_sensors = False + env.cfg = SimpleNamespace(num_rerenders_on_reset=0) env.extras = {"log": {}} env._reset_terminated_envs() env.reset_buf.nonzero.assert_not_called() - env._reset_idx.assert_called_once_with(env_mask=reset_mask, env_ids=None) + env._reset_idx.assert_called_once_with(env_mask=reset_mask) def test_public_reset_selection_stays_as_mask_without_host_consumers() -> None: @@ -101,14 +98,12 @@ def call(stage, fn, /, *args, **kwargs): env._episode_length_buf_wp = wp.ones(3, dtype=wp.int64, device="cpu") env.extras = {} reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") - reset_ids = torch.tensor([1]) - - env._reset_idx(env_mask=reset_mask, env_ids=reset_ids) + env._reset_idx(env_mask=reset_mask) assert [stage for stage, _ in stages[:2]] == ["CurriculumManager_compute", "Scene_reset"] assert stages[0][1]["env_mask"] is env.reset_mask_wp np.testing.assert_array_equal(stages[0][1]["env_mask"].numpy(), reset_mask.numpy()) - assert stages[0][1]["env_ids"] is reset_ids + assert "env_ids" not in stages[0][1] assert "CurriculumManager_reset" in [stage for stage, _ in stages] @@ -167,10 +162,34 @@ def test_reset_compacts_ids_for_enabled_host_consumer() -> None: reset_ids = env._reset_host_pre.call_args.args[0] torch.testing.assert_close(reset_ids, torch.tensor([1])) - env._reset_idx.assert_called_once_with(env_mask=reset_mask, env_ids=reset_ids) + env._reset_idx.assert_called_once_with(env_mask=reset_mask) assert env.extras["log"]["host"] == 1.0 +def test_rtx_rerender_does_not_require_compact_reset_ids() -> None: + """RTX reset rendering should use the nonempty predicate without materializing IDs.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") + env.termination_manager = SimpleNamespace(dones_wp=reset_mask) + env.reset_buf = torch.tensor([False, True, False]) + env._reset_requires_host_selection = Mock(return_value=False) + env._reset_idx = Mock() + env._reset_host_pre = Mock() + env._reset_host_post = Mock() + env._has_recorders = False + env.has_rtx_sensors = True + env.cfg = SimpleNamespace(num_rerenders_on_reset=2) + env.sim = SimpleNamespace(render=Mock()) + env.extras = {"log": {}} + + env._reset_terminated_envs() + + env._reset_idx.assert_called_once_with(env_mask=reset_mask) + env._reset_host_pre.assert_not_called() + env._reset_host_post.assert_not_called() + assert env.sim.render.call_count == 2 + + def test_empty_reset_skips_legacy_host_boundary_and_preserves_logs() -> None: """An empty host-visible selection should not execute reset side effects.""" env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) diff --git a/source/isaaclab_experimental/test/managers/test_curriculum_manager.py b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py index dcd9dbac6cf3..73c33a28fce7 100644 --- a/source/isaaclab_experimental/test/managers/test_curriculum_manager.py +++ b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py @@ -21,6 +21,7 @@ from isaaclab.managers import ManagerTermBase as StableManagerTermBase from isaaclab.terrains import TerrainImporter +from isaaclab.utils.warp import ProxyArray class _GlobalCurriculumTerm(StableManagerTermBase): @@ -31,9 +32,21 @@ def __call__(self, env, env_ids, value: float) -> float: return value -def _legacy_id_count(env, env_ids, scale: float) -> float: - """Legacy curriculum state that genuinely depends on compact environment IDs.""" - return len(env_ids) * scale +class _LegacyIdCurriculumTerm(StableManagerTermBase): + """Legacy term that records the compact IDs received by compute and reset.""" + + def __init__(self, cfg, env): + super().__init__(cfg, env) + self.compute_env_ids = None + self.reset_env_ids = None + + def __call__(self, env, env_ids, scale: float) -> float: + del env + self.compute_env_ids = env_ids.clone() + return len(env_ids) * scale + + def reset(self, env_ids=None) -> None: + self.reset_env_ids = env_ids.clone() def _structured_state(env, env_ids) -> dict[str, float]: @@ -111,7 +124,11 @@ def __init__(self, robot: _Robot, terrain: TerrainImporter): @property def env_origins(self) -> torch.Tensor: - return self.terrain.env_origins + return self.terrain.env_origins.torch + + @property + def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): + return self.terrain.env_origins.warp def __getitem__(self, name: str): return self._entities[name] @@ -128,31 +145,85 @@ def __init__(self, terrain: TerrainImporter, root_positions: np.ndarray, command self.device = "cpu" self.scene = _Scene(_Robot(root_positions), terrain) self.command_manager = _CommandManager(commands) - self.env_origins_wp = wp.from_torch(terrain.env_origins, dtype=wp.vec3f) self.rng_state_wp = wp.array(np.arange(self.num_envs, dtype=np.uint32) + 101, device=self.device) self.max_episode_length_s = 10.0 self.sim = SimpleNamespace(is_playing=lambda: True) + @property + def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): + """Return the scene's persistent Warp origin view.""" + return self.scene.env_origins_wp + def _make_terrain(levels: list[int]) -> TerrainImporter: num_envs = len(levels) terrain = TerrainImporter.__new__(TerrainImporter) terrain.device = "cpu" - terrain.cfg = SimpleNamespace(terrain_generator=SimpleNamespace(size=(8.0, 8.0))) - terrain.max_terrain_level = 3 - terrain.terrain_levels = torch.tensor(levels, dtype=torch.int64) - terrain.terrain_types = torch.tensor([index % 2 for index in range(num_envs)], dtype=torch.int64) - terrain.terrain_origins = torch.zeros(3, 2, 3, dtype=torch.float32) + terrain.cfg = SimpleNamespace( + num_envs=num_envs, + max_init_terrain_level=2, + terrain_generator=SimpleNamespace(size=(8.0, 8.0)), + ) + level_values = np.asarray(levels, dtype=np.int64) + type_values = np.asarray([index % 2 for index in range(num_envs)], dtype=np.int64) + origin_values = np.zeros((3, 2, 3), dtype=np.float32) for level in range(3): for terrain_type in range(2): - terrain.terrain_origins[level, terrain_type] = torch.tensor( - [100.0 * level + 10.0 * terrain_type, float(level), float(terrain_type)] - ) - terrain.env_origins = terrain.terrain_origins[terrain.terrain_levels, terrain.terrain_types].clone() - terrain._invalidate_warp_origin_views() + origin_values[level, terrain_type] = [ + 100.0 * level + 10.0 * terrain_type, + float(level), + float(terrain_type), + ] + terrain.configure_env_origins(origin_values) + terrain.terrain_levels.torch.copy_(torch.from_numpy(level_values)) + terrain.terrain_types.torch.copy_(torch.from_numpy(type_values)) + terrain.env_origins.torch.copy_( + terrain.terrain_origins.torch[terrain.terrain_levels.torch, terrain.terrain_types.torch] + ) return terrain +def test_terrain_importer_exposes_persistent_proxy_array_views(): + """Canonical terrain storage should expose persistent zero-copy Torch and Warp views.""" + terrain = _make_terrain([0, 1, 2, 1]) + + assert isinstance(terrain.terrain_origins, ProxyArray) + assert isinstance(terrain.env_origins, ProxyArray) + assert isinstance(terrain.terrain_levels, ProxyArray) + assert isinstance(terrain.terrain_types, ProxyArray) + assert terrain.terrain_origins.dtype == wp.vec3f + assert terrain.env_origins.dtype == wp.vec3f + assert terrain.terrain_levels.dtype == wp.int64 + assert terrain.terrain_types.dtype == wp.int64 + assert terrain.terrain_origins.torch.data_ptr() == terrain.terrain_origins.warp.ptr + assert terrain.env_origins.torch.data_ptr() == terrain.env_origins.warp.ptr + assert terrain.terrain_levels.torch.data_ptr() == terrain.terrain_levels.warp.ptr + assert terrain.terrain_types.torch.data_ptr() == terrain.terrain_types.warp.ptr + + terrain.terrain_levels.torch[0] = 2 + assert terrain.terrain_levels.warp.numpy()[0] == 2 + replacement_levels = wp.array([1, 0, 2, 1], dtype=wp.int64, device="cpu") + wp.copy(terrain.terrain_levels.warp, replacement_levels) + wp.synchronize() + torch.testing.assert_close(terrain.terrain_levels.torch, torch.tensor([1, 0, 2, 1])) + + +def test_terrain_importer_grid_origins_use_proxy_array_without_curriculum_buffers(): + """Grid origins should use the same proxy contract without curriculum buffers.""" + terrain = TerrainImporter.__new__(TerrainImporter) + terrain.device = "cpu" + terrain.cfg = SimpleNamespace(num_envs=4, env_spacing=2.0) + terrain.configure_env_origins() + + assert terrain.terrain_origins is None + assert isinstance(terrain.env_origins, ProxyArray) + assert terrain.env_origins.shape == (4,) + assert terrain.env_origins.torch.shape == (4, 3) + assert terrain.env_origins.torch.data_ptr() == terrain.env_origins.warp.ptr + assert terrain.terrain_levels is None + assert terrain.terrain_types is None + + def test_terrain_importer_mask_update_preserves_sparse_and_stable_semantics(): """Mask updates should handle up/down/clamp/wrap while stable IDs remain supported.""" terrain = _make_terrain([0, 1, 2, 2, 1, 0]) @@ -161,27 +232,36 @@ def test_terrain_importer_mask_update_preserves_sparse_and_stable_semantics(): move_down = wp.array([True, False, False, True, False, True], dtype=wp.bool, device="cpu") rng_state = wp.array(np.arange(6, dtype=np.uint32) + 71, device="cpu") rng_before = rng_state.numpy().copy() - origins_before = terrain.env_origins.clone() - levels_wp = terrain.terrain_levels_wp + origins_before = terrain.env_origins.torch.clone() + levels_wp = terrain.terrain_levels.warp + origins_wp = terrain.env_origins.warp + levels_ptr = levels_wp.ptr + origins_ptr = origins_wp.ptr + assert terrain.terrain_levels.torch.data_ptr() == levels_ptr + assert terrain.env_origins.torch.data_ptr() == origins_ptr terrain.update_env_origins_mask(env_mask, move_up, move_down, rng_state) wp.synchronize() - levels = terrain.terrain_levels.tolist() + levels = terrain.terrain_levels.torch.tolist() assert levels[0] == 0 assert levels[1] == 2 assert 0 <= levels[2] < terrain.max_terrain_level assert levels[3] == 1 assert levels[4:] == [1, 0] torch.testing.assert_close( - terrain.env_origins[:4], terrain.terrain_origins[terrain.terrain_levels[:4], terrain.terrain_types[:4]] + terrain.env_origins.torch[:4], + terrain.terrain_origins.torch[terrain.terrain_levels.torch[:4], terrain.terrain_types.torch[:4]], ) - torch.testing.assert_close(terrain.env_origins[4:], origins_before[4:]) + torch.testing.assert_close(terrain.env_origins.torch[4:], origins_before[4:]) assert np.all(rng_state.numpy()[:4] != rng_before[:4]) np.testing.assert_array_equal(rng_state.numpy()[4:], rng_before[4:]) - assert terrain.terrain_levels_wp is levels_wp + assert terrain.terrain_levels.warp is levels_wp + assert terrain.env_origins.warp is origins_wp + assert terrain.terrain_levels.warp.ptr == levels_ptr + assert terrain.env_origins.warp.ptr == origins_ptr - levels_before_stable_update = terrain.terrain_levels.clone() + levels_before_stable_update = terrain.terrain_levels.torch.clone() env_ids = torch.tensor([1, 3], dtype=torch.int64) terrain.update_env_origins( env_ids, @@ -190,15 +270,18 @@ def test_terrain_importer_mask_update_preserves_sparse_and_stable_semantics(): ) expected_levels = levels_before_stable_update.clone() expected_levels[env_ids] -= 1 - torch.testing.assert_close(terrain.terrain_levels, expected_levels) - assert terrain.terrain_levels_wp is levels_wp + torch.testing.assert_close(terrain.terrain_levels.torch, expected_levels) + assert terrain.terrain_levels.warp is levels_wp + assert terrain.env_origins.warp is origins_wp + assert terrain.terrain_levels.warp.ptr == levels_ptr + assert terrain.env_origins.warp.ptr == origins_ptr def test_curriculum_manager_updates_masked_levels_and_logs_without_host_compaction(monkeypatch: pytest.MonkeyPatch): """Manager compute/reset should remain mask-native and expose persistent scalar logging.""" assert terrain_levels_vel is not TerrainLevelsVel terrain = _make_terrain([0, 1, 2, 2, 1, 0]) - root_positions = terrain.env_origins.numpy().copy() + root_positions = terrain.env_origins.torch.numpy().copy() root_positions[:, 0] += np.array([5.0, 1.0, 3.0, 5.0, 9.0, 0.0], dtype=np.float32) commands = np.array( [[0.1, 0.0, 0.0], [1.0, 0.0, 0.0], [0.2, 0.0, 0.0], [2.0, 0.0, 0.0], [0.1, 0.0, 0.0], [0.1, 0.0, 0.0]], @@ -235,7 +318,7 @@ def _fail_host_compaction(*args, **kwargs): extras = manager.reset(env_mask) wp.synchronize() - levels = terrain.terrain_levels.tolist() + levels = terrain.terrain_levels.torch.tolist() assert levels[0] == 1 assert levels[1] == 0 assert levels[2] == 2 @@ -245,34 +328,34 @@ def _fail_host_compaction(*args, **kwargs): np.testing.assert_array_equal(term._move_down_wp.numpy(), [False, True, False, False, False, False]) assert extras is extras_ref assert extras["Curriculum/terrain_levels"] is manager.reset_extras["Curriculum/terrain_levels"] - torch.testing.assert_close(extras["Curriculum/terrain_levels"], terrain.terrain_levels.float().mean()) + torch.testing.assert_close(extras["Curriculum/terrain_levels"], terrain.terrain_levels.torch.float().mean()) torch.testing.assert_close(extras["Curriculum/global_state"], torch.tensor(7.0)) assert term._move_up_wp.ptr == move_up_ptr assert term._move_down_wp.ptr == move_down_ptr assert manager._term_states_wp.ptr == state_ptr -def test_curriculum_manager_requests_ids_only_for_legacy_id_terms(): - """Legacy terms should require host IDs only when their config says they consume them.""" - terrain = _make_terrain([0, 1, 2, 2]) - env = _Env( - terrain, - terrain.env_origins.numpy().copy(), - np.zeros((4, 3), dtype=np.float32), +def test_curriculum_manager_compacts_ids_inside_legacy_term_boundary(): + """Legacy terms should receive compact IDs without polluting the environment reset API.""" + env = SimpleNamespace( + num_envs=4, + device="cpu", + sim=SimpleNamespace(is_playing=lambda: True), ) manager = CurriculumManager( - {"id_count": CurriculumTermCfg(func=_legacy_id_count, params={"scale": 2.0})}, + {"id_count": CurriculumTermCfg(func=_LegacyIdCurriculumTerm, params={"scale": 2.0})}, env, ) env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + term = manager._term_cfgs[0].func assert manager.requires_host_ids assert manager.requires_host_boundary - with pytest.raises(RuntimeError, match="requires compact env_ids"): - manager.compute(env_mask) + manager.compute(env_mask) + extras = manager.reset(env_mask) - manager.compute(env_mask, env_ids=[0, 2]) - extras = manager.reset(env_mask, env_ids=[0, 2]) + torch.testing.assert_close(term.compute_env_ids, torch.tensor([0, 2])) + torch.testing.assert_close(term.reset_env_ids, torch.tensor([0, 2])) torch.testing.assert_close(extras["Curriculum/id_count"], torch.tensor(4.0)) @@ -281,7 +364,7 @@ def test_registered_reach_curricula_update_weights_without_a_host_boundary(): terrain = _make_terrain([0, 1, 2, 2]) env = _Env( terrain, - terrain.env_origins.numpy().copy(), + terrain.env_origins.torch.numpy().copy(), np.zeros((4, 3), dtype=np.float32), ) env.reward_manager = _RewardManager() @@ -311,7 +394,7 @@ def test_registered_reach_curricula_update_weights_without_a_host_boundary(): def test_curriculum_manager_rejects_structured_legacy_state() -> None: """Unsupported fallback logging should fail explicitly instead of narrowing silently.""" terrain = _make_terrain([0, 1]) - env = _Env(terrain, terrain.env_origins.numpy().copy(), np.zeros((2, 3), dtype=np.float32)) + env = _Env(terrain, terrain.env_origins.torch.numpy().copy(), np.zeros((2, 3), dtype=np.float32)) manager = CurriculumManager( {"structured": CurriculumTermCfg(func=_structured_state, requires_host_ids=False)}, env, diff --git a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py index abe94920bc33..e7136c19ab54 100644 --- a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py +++ b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py @@ -41,6 +41,27 @@ def counted_call(value): self.assertEqual(cache.capture_or_replay("stage", counted_call, args=(2,)), 2) self.assertEqual(call_count, 2) + def test_non_capturable_group_stays_eager(self): + """A group-level safety decision should keep every member stage eager.""" + call_count = 0 + + def counted_call(): + nonlocal call_count + call_count += 1 + + self.cache.register_capturability("RewardManager", False) + self.cache.call("RewardManager_compute", counted_call, capture=True, group="RewardManager") + self.cache.call("RewardManager_reset", counted_call, capture=True, group="RewardManager") + + self.assertEqual(call_count, 2) + + def test_capturability_registration_is_conservative(self): + """A later positive registration must not override a known unsafe operation.""" + self.cache.register_capturability("EventManager", False) + self.cache.register_capturability("EventManager", True) + + self.assertFalse(self.cache.is_capturable("EventManager")) + # ------------------------------------------------------------------ # Warm-up # ------------------------------------------------------------------ diff --git a/source/isaaclab_physx/changelog.d/jichuanh-warp-frontend-cleanup.rst b/source/isaaclab_physx/changelog.d/jichuanh-warp-frontend-cleanup.rst deleted file mode 100644 index c56453d2cc8e..000000000000 --- a/source/isaaclab_physx/changelog.d/jichuanh-warp-frontend-cleanup.rst +++ /dev/null @@ -1,5 +0,0 @@ -Fixed -^^^^^ - -* Fixed PhysX asset resets to forward boolean environment masks to Newton - actuator and external-wrench state. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 6e43c71b78f9..0bb66b68543e 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -236,21 +236,15 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None # use ellipses object to skip initial indices. if (env_ids is None) or (env_ids == slice(None)): env_ids = slice(None) - actuator_env_ids = env_ids - if self.actuators and env_mask is not None and not getattr(self, "_has_newton_actuators", False): - # Explicit boundary for the legacy Torch actuator fallback. The Warp - # frontend enables Newton-native actuators, whose state resets below - # consume the mask directly without materializing compact IDs. - actuator_env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) - # reset Lab actuators + # reset actuators (including Newton-native adapter which owns its states) for actuator in self.actuators.values(): - actuator.reset(actuator_env_ids) + actuator.reset(env_ids) # Reset Newton-actuator per-env states (delay queues, neural hidden state, etc.). # The adapter is per-articulation on PhysX and is not part of ``self.actuators``. # ``getattr`` guards subclasses (e.g. ``Multirotor``) that override # ``_process_actuators_cfg`` and never initialize these attributes. if getattr(self, "_has_newton_actuators", False) and getattr(self, "newton_actuator_adapter", None) is not None: - self.newton_actuator_adapter.reset(env_ids, env_mask=env_mask) + self.newton_actuator_adapter.reset(env_ids) # reset external wrenches. self._instantaneous_wrench_composer.reset(env_ids, env_mask) self._permanent_wrench_composer.reset(env_ids, env_mask) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py index 4c42628313fd..16c9801cb47c 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py @@ -141,8 +141,8 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None if (env_ids is None) or (env_ids == slice(None)): env_ids = slice(None) # reset external wrench - self._instantaneous_wrench_composer.reset(env_ids, env_mask) - self._permanent_wrench_composer.reset(env_ids, env_mask) + self._instantaneous_wrench_composer.reset(env_ids) + self._permanent_wrench_composer.reset(env_ids) def write_data_to_sim(self) -> None: """Write external wrench to the simulation. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py index e2b36424774d..d03bfb78012f 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py @@ -170,8 +170,6 @@ def reset( Args: env_ids: Environment indices. If None, then all indices are used. object_ids: Object indices. If None, then all indices are used. - env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). - object_mask: Object mask. Not used currently. """ # resolve all indices if env_ids is None: @@ -179,8 +177,8 @@ def reset( if object_ids is None: object_ids = self._ALL_BODY_INDICES # reset external wrench - self._instantaneous_wrench_composer.reset(env_ids, env_mask) - self._permanent_wrench_composer.reset(env_ids, env_mask) + self._instantaneous_wrench_composer.reset(env_ids) + self._permanent_wrench_composer.reset(env_ids) def write_data_to_sim(self) -> None: """Write external wrench to the simulation. diff --git a/source/isaaclab_physx/test/assets/test_physx_mask_reset_forwarding.py b/source/isaaclab_physx/test/assets/test_physx_mask_reset_forwarding.py deleted file mode 100644 index 63e616b11b76..000000000000 --- a/source/isaaclab_physx/test/assets/test_physx_mask_reset_forwarding.py +++ /dev/null @@ -1,70 +0,0 @@ -# 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 - -"""Tests for Warp reset-mask forwarding to PhysX asset wrench storage.""" - -from isaaclab.app import AppLauncher -from isaaclab.test.utils import resolve_test_sim_device - -simulation_app = AppLauncher(headless=True, device=resolve_test_sim_device()).app - -"""Everything else follows.""" - -from unittest.mock import Mock - -import warp as wp -from isaaclab_physx.assets.articulation.articulation import Articulation -from isaaclab_physx.assets.rigid_object.rigid_object import RigidObject -from isaaclab_physx.assets.rigid_object_collection.rigid_object_collection import RigidObjectCollection - - -def test_rigid_object_forwards_reset_mask_to_wrench_composers() -> None: - """Rigid-object resets should preserve mask selection at wrench storage.""" - asset = RigidObject.__new__(RigidObject) - asset._initialize_handle = None - asset._invalidate_initialize_handle = None - asset._prim_deletion_handle = None - asset._instantaneous_wrench_composer = Mock() - asset._permanent_wrench_composer = Mock() - env_mask = wp.array([False, True], dtype=wp.bool, device="cpu") - - asset.reset(env_mask=env_mask) - - asset._instantaneous_wrench_composer.reset.assert_called_once_with(slice(None), env_mask) - asset._permanent_wrench_composer.reset.assert_called_once_with(slice(None), env_mask) - - -def test_rigid_object_collection_forwards_reset_mask_to_wrench_composers() -> None: - """Collection resets should preserve mask selection at wrench storage.""" - asset = RigidObjectCollection.__new__(RigidObjectCollection) - asset._initialize_handle = None - asset._invalidate_initialize_handle = None - asset._prim_deletion_handle = None - asset._instantaneous_wrench_composer = Mock() - asset._permanent_wrench_composer = Mock() - env_ids = [0] - env_mask = wp.array([True, False], dtype=wp.bool, device="cpu") - - asset.reset(env_ids=env_ids, object_ids=slice(None), env_mask=env_mask) - - asset._instantaneous_wrench_composer.reset.assert_called_once_with(env_ids, env_mask) - asset._permanent_wrench_composer.reset.assert_called_once_with(env_ids, env_mask) - - -def test_articulation_without_lab_actuators_keeps_mask_on_device() -> None: - """Mask resets should not materialize IDs when no Torch actuator consumes them.""" - asset = Articulation.__new__(Articulation) - asset._initialize_handle = None - asset._invalidate_initialize_handle = None - asset._prim_deletion_handle = None - asset.actuators = {} - asset._has_newton_actuators = False - asset._instantaneous_wrench_composer = Mock() - asset._permanent_wrench_composer = Mock() - env_mask = object() - - asset.reset(env_mask=env_mask) - - asset._instantaneous_wrench_composer.reset.assert_called_once_with(slice(None), env_mask) diff --git a/source/isaaclab_tasks/changelog.d/jichuanh-warp-frontend-cleanup.rst b/source/isaaclab_tasks/changelog.d/jichuanh-warp-frontend-cleanup.rst index 7d2f89fbbf9d..5157824a2c21 100644 --- a/source/isaaclab_tasks/changelog.d/jichuanh-warp-frontend-cleanup.rst +++ b/source/isaaclab_tasks/changelog.d/jichuanh-warp-frontend-cleanup.rst @@ -1,4 +1,5 @@ Fixed ^^^^^ -* Fixed identity quaternion initialization for Dexsuite uniform pose commands. +* Fixed DexSuite slab-clearance setup to reuse pointer-stable scene origin + storage. diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/anymal_c_direct/anymal_c_env.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/anymal_c_direct/anymal_c_env.py index f932022481e0..11ff1287a23b 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/anymal_c_direct/anymal_c_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/anymal_c_direct/anymal_c_env.py @@ -205,7 +205,7 @@ def _reset_idx(self, env_ids: torch.Tensor | None): joint_vel = self._robot.data.default_joint_vel.torch[env_ids] default_root_pose = self._robot.data.default_root_pose.torch[env_ids] default_root_vel = self._robot.data.default_root_vel.torch[env_ids] - default_root_pose[:, :3] += self._terrain.env_origins[env_ids] + default_root_pose[:, :3] += self._terrain.env_origins.torch[env_ids] self._robot.write_root_pose_to_sim_index(root_pose=default_root_pose, env_ids=env_ids) self._robot.write_root_velocity_to_sim_index(root_velocity=default_root_vel, env_ids=env_ids) self._robot.write_joint_position_to_sim_index(position=joint_pos, env_ids=env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/commands/pose_commands.py b/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/commands/pose_commands.py index c7014ba7b304..b0165e452ee3 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/commands/pose_commands.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/commands/pose_commands.py @@ -73,7 +73,7 @@ def __init__(self, cfg: dex_cmd_cfgs.ObjectUniformPoseCommandCfg, env: ManagerBa # create buffers # -- commands: (x, y, z, qx, qy, qz, qw) in root frame self.pose_command_b = torch.zeros(self.num_envs, 7, device=self.device) - self.pose_command_b[:, 6] = 1.0 + self.pose_command_b[:, 3] = 1.0 self.pose_command_w = torch.zeros_like(self.pose_command_b) # -- metrics self.metrics["position_error"] = torch.zeros(self.num_envs, device=self.device) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/events.py index 9f10f18a2ec0..5c6691d5f6df 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/events.py @@ -580,7 +580,7 @@ def __init__(self, cfg: SlabClearanceCfg, env: ManagerBasedEnv): wp.array(np.array([x is not None for x, _, _ in slabs], dtype=np.int32), device=device), wp.array(np.array([y is not None for _, y, _ in slabs], dtype=np.int32), device=device), ] - self._env_origins = wp.from_torch(env.scene.env_origins.contiguous(), dtype=wp.vec3) + self._env_origins = env.scene.env_origins_wp def __call__(self, env: ManagerBasedEnv, env_ids: torch.Tensor) -> torch.Tensor: num = len(env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/curriculums.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/curriculums.py index 1438cbf03801..6b05db44348b 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/curriculums.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/curriculums.py @@ -53,4 +53,4 @@ def terrain_levels_vel( # update terrain levels terrain.update_env_origins(env_ids, move_up, move_down) # return the mean terrain level - return torch.mean(terrain.terrain_levels.float()) + return torch.mean(terrain.terrain_levels.torch.float()) diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst index cd2905af15b2..6cfc728745b0 100644 --- a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst +++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -3,3 +3,9 @@ Added * Added Warp-native terrain and reward-weight curricula and connected locomotion and reach hot paths to pointer-stable Warp command arrays. + +Fixed +^^^^^ + +* Fixed terrain-boundary termination state leaking between environment + instances with different terrain configurations. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py index d2f1be03dd17..b8bb6bcab6ac 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py @@ -582,8 +582,7 @@ def __init__(self, cfg: AllegroHandWarpEnvCfg, render_mode: str | None = None, * self.y_unit_vec = wp.vec3f(0.0, 1.0, 0.0) self.z_unit_vec = wp.vec3f(0.0, 0.0, 1.0) - # Per-env origins (Warp view for kernels; Torch env uses `self.scene.env_origins` directly). - self.env_origins = wp.from_torch(self.scene.env_origins, dtype=wp.vec3f) + self.env_origins = self.scene.env_origins_wp # --------------------------------------------------------------------- # Warp buffers diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py index 62400c4d48fe..fedf35f23237 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py @@ -385,7 +385,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.rpy = wp.zeros((self.num_envs), dtype=wp.vec3f, device=self.sim.device) self.angle_to_target = wp.zeros((self.num_envs), dtype=wp.float32, device=self.sim.device) self.dof_pos_scaled = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device) - self.env_origins = wp.from_torch(self.scene.env_origins, dtype=wp.vec3f) + self.env_origins = self.scene.env_origins_wp self.actions_mapped = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device) # Initial states and targets diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py index 66a9131011f3..a8b7964617a4 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py @@ -100,7 +100,7 @@ class EventCfg: """Configuration for events.""" reset_base = EventTerm( - func=mdp.reset_root_state_uniform, + func=mdp.ResetRootStateUniform, mode="reset", params={"pose_range": {}, "velocity_range": {}}, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py index 9de689ffb98e..49bdb600e4e2 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py @@ -109,7 +109,7 @@ class EventCfg: """Configuration for events.""" reset_base = EventTerm( - func=mdp.reset_root_state_uniform, + func=mdp.ResetRootStateUniform, mode="reset", params={"pose_range": {}, "velocity_range": {}}, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py index f55af9254efa..734b496ad18b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py @@ -77,7 +77,7 @@ def __init__(self, cfg: CurriculumTermCfg, env: ManagerBasedRLEnv): self._root_pos_w_wp = self._asset.data.root_pos_w.warp self._env_origins_wp = env.env_origins_wp self._command_wp = env.command_manager.get_command_wp("base_velocity") - self._terrain_levels_wp = self._terrain.terrain_levels_wp + self._terrain_levels_wp = self._terrain.terrain_levels.warp self._move_up_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) self._move_down_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) self._terrain_half_length = float(self._terrain.cfg.terrain_generator.size[0]) * 0.5 @@ -141,4 +141,4 @@ def terrain_levels_vel( move_down = distance < torch.norm(command[env_ids, :2], dim=1) * env.max_episode_length_s * 0.5 move_down *= ~move_up terrain.update_env_origins(env_ids, move_up, move_down) - return torch.mean(terrain.terrain_levels.float()) + return torch.mean(terrain.terrain_levels.torch.float()) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py index 4e0c32c58e64..19583eb4edeb 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py @@ -36,31 +36,24 @@ def terrain_out_of_bounds( env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), distance_buffer: float = 3.0 ) -> None: """Terminate when the actor moves too close to the edge of the terrain.""" - fn = terrain_out_of_bounds - if not getattr(fn, "_is_warmed_up", False): - terrain_type = env.scene.cfg.terrain.terrain_type - if terrain_type == "plane": - fn._is_plane = True - elif terrain_type == "generator": - fn._is_plane = False - terrain_gen_cfg = env.scene.terrain.cfg.terrain_generator - grid_width, grid_length = terrain_gen_cfg.size - n_rows, n_cols = terrain_gen_cfg.num_rows, terrain_gen_cfg.num_cols - border_width = terrain_gen_cfg.border_width - fn._half_width = 0.5 * (n_rows * grid_width + 2 * border_width) - fn._half_height = 0.5 * (n_cols * grid_length + 2 * border_width) - else: - raise ValueError("Received unsupported terrain type, must be either 'plane' or 'generator'.") - fn._is_warmed_up = True - - if fn._is_plane: + terrain_type = env.scene.cfg.terrain.terrain_type + if terrain_type == "plane": out.zero_() return + if terrain_type != "generator": + raise ValueError("Received unsupported terrain type, must be either 'plane' or 'generator'.") + + terrain_gen_cfg = env.scene.terrain.cfg.terrain_generator + grid_width, grid_length = terrain_gen_cfg.size + n_rows, n_cols = terrain_gen_cfg.num_rows, terrain_gen_cfg.num_cols + border_width = terrain_gen_cfg.border_width + half_width = 0.5 * (n_rows * grid_width + 2 * border_width) + half_height = 0.5 * (n_cols * grid_length + 2 * border_width) asset: Articulation = env.scene[asset_cfg.name] wp.launch( kernel=_terrain_out_of_bounds_kernel, dim=env.num_envs, - inputs=[asset.data.root_pos_w.warp, fn._half_width, fn._half_height, distance_buffer, out], + inputs=[asset.data.root_pos_w.warp, half_width, half_height, distance_buffer, out], device=env.device, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py index 2807bf2c5a31..8999c4a01a64 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py @@ -146,7 +146,7 @@ class EventCfg: # FIXME(warp-migration): COM randomization in exp manager-based locomotion currently causes # NaNs and is temporarily disabled. # base_com = EventTerm( - # func=mdp.randomize_rigid_body_com, + # func=mdp.RandomizeRigidBodyCom, # mode="startup", # params={ # "asset_cfg": SceneEntityCfg("robot", body_names="base"), @@ -157,7 +157,7 @@ class EventCfg: # reset base_external_force_torque = EventTerm( - func=mdp.apply_external_force_torque, + func=mdp.ApplyExternalForceTorque, mode="reset", params={ "asset_cfg": SceneEntityCfg("robot", body_names="base"), @@ -167,7 +167,7 @@ class EventCfg: ) reset_base = EventTerm( - func=mdp.reset_root_state_uniform, + func=mdp.ResetRootStateUniform, mode="reset", params={ "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, @@ -193,7 +193,7 @@ class EventCfg: # interval push_robot = EventTerm( - func=mdp.push_by_setting_velocity, + func=mdp.PushBySettingVelocity, mode="interval", interval_range_s=(10.0, 15.0), params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}}, diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py index ebb94cab0c5c..e7c2684b7f96 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py @@ -109,7 +109,7 @@ class EventCfg: """Configuration for events.""" reset_base = EventTerm( - func=mdp.reset_root_state_uniform, + func=mdp.ResetRootStateUniform, mode="reset", params={"pose_range": {}, "velocity_range": {}}, ) diff --git a/source/isaaclab_tasks_experimental/test/test_terrain_out_of_bounds.py b/source/isaaclab_tasks_experimental/test/test_terrain_out_of_bounds.py new file mode 100644 index 000000000000..c306608323ec --- /dev/null +++ b/source/isaaclab_tasks_experimental/test/test_terrain_out_of_bounds.py @@ -0,0 +1,41 @@ +# 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 + +"""Tests for Warp locomotion termination terms.""" + +from types import SimpleNamespace + +import numpy as np +import warp as wp +from isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp.terminations import terrain_out_of_bounds + + +class _Scene(dict): + """Minimal scene mapping with terrain configuration.""" + + +def _make_env(terrain_type: str, root_x: float = 0.0): + root_pos_w = wp.array(np.full((2, 3), (root_x, 0.0, 0.0), dtype=np.float32), dtype=wp.vec3f, device="cpu") + scene = _Scene(robot=SimpleNamespace(data=SimpleNamespace(root_pos_w=SimpleNamespace(warp=root_pos_w)))) + scene.cfg = SimpleNamespace(terrain=SimpleNamespace(terrain_type=terrain_type)) + scene.terrain = SimpleNamespace( + cfg=SimpleNamespace( + terrain_generator=SimpleNamespace(size=(2.0, 2.0), num_rows=2, num_cols=2, border_width=0.0) + ) + ) + return SimpleNamespace(scene=scene, num_envs=2, device="cpu") + + +def test_terrain_out_of_bounds_does_not_share_terrain_state_between_envs(): + """A plane environment must not poison bounds used by a later generated-terrain environment.""" + for attribute in ("_is_warmed_up", "_is_plane", "_half_width", "_half_height"): + if hasattr(terrain_out_of_bounds, attribute): + delattr(terrain_out_of_bounds, attribute) + + out = wp.zeros(2, dtype=wp.bool, device="cpu") + terrain_out_of_bounds(_make_env("plane"), out) + terrain_out_of_bounds(_make_env("generator", root_x=2.0), out, distance_buffer=0.5) + + np.testing.assert_array_equal(out.numpy(), np.ones(2, dtype=np.bool_)) From dc738c364c6ccf9f1f56c736c5d0ef825b91c560 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Mon, 20 Jul 2026 00:45:07 -0700 Subject: [PATCH 41/58] Pin terrain origin ProxyArray buffers Store terrain origin, level, and type wrappers as importer-owned members. Public properties now return the same persistent objects so ownership and pointer stability are explicit. --- .../isaaclab/terrains/terrain_importer.py | 105 +++++++++--------- 1 file changed, 55 insertions(+), 50 deletions(-) diff --git a/source/isaaclab/isaaclab/terrains/terrain_importer.py b/source/isaaclab/isaaclab/terrains/terrain_importer.py index d6428d3f18d8..c9d740b438fe 100644 --- a/source/isaaclab/isaaclab/terrains/terrain_importer.py +++ b/source/isaaclab/isaaclab/terrains/terrain_importer.py @@ -75,33 +75,6 @@ class TerrainImporter: terrain_prim_paths: list[str] """A list containing the USD prim paths to the imported terrains.""" - terrain_origins: ProxyArray | None - """The origins [m] of the sub-terrains, shape ``(num_rows, num_cols)`` of ``wp.vec3f``. - - The :attr:`~isaaclab.utils.warp.ProxyArray.torch` view resolves to shape ``(num_rows, num_cols, 3)``. - - If terrain origins is not None, the environment origins are computed based on the terrain origins. - Otherwise, the environment origins are computed based on the grid spacing. - """ - - env_origins: ProxyArray - """The environment origins [m], shape ``(num_envs,)`` of ``wp.vec3f``. - - The :attr:`~isaaclab.utils.warp.ProxyArray.torch` view resolves to shape ``(num_envs, 3)``. - """ - - terrain_levels: ProxyArray | None - """The terrain level assigned to each environment, shape ``(num_envs,)`` of ``wp.int64``. - - This is available only when sub-terrain origins are configured. - """ - - terrain_types: ProxyArray | None - """The terrain type assigned to each environment, shape ``(num_envs,)`` of ``wp.int64``. - - This is available only when sub-terrain origins are configured. - """ - def __init__(self, cfg: TerrainImporterCfg): """Initialize the terrain importer. @@ -122,9 +95,10 @@ def __init__(self, cfg: TerrainImporterCfg): # create buffers for the terrains self.terrain_prim_paths = list() - self.terrain_origins = None - self.terrain_levels = None - self.terrain_types = None + self._terrain_origins: ProxyArray | None = None + self._env_origins: ProxyArray + self._terrain_levels: ProxyArray | None = None + self._terrain_types: ProxyArray | None = None # private variables self._terrain_flat_patches = dict() @@ -192,6 +166,39 @@ def terrain_names(self) -> list[str]: """A list of names of the imported terrains.""" return [f"'{path.split('/')[-1]}'" for path in self.terrain_prim_paths] + @property + def terrain_origins(self) -> ProxyArray | None: + """Sub-terrain origins [m], shape ``(num_rows, num_cols)`` of ``wp.vec3f``. + + The :attr:`~isaaclab.utils.warp.ProxyArray.torch` view has shape ``(num_rows, num_cols, 3)``. + This is ``None`` when environment origins are configured from grid spacing. + """ + return self._terrain_origins + + @property + def env_origins(self) -> ProxyArray: + """Environment origins [m], shape ``(num_envs,)`` of ``wp.vec3f``. + + The :attr:`~isaaclab.utils.warp.ProxyArray.torch` view has shape ``(num_envs, 3)``. + """ + return self._env_origins + + @property + def terrain_levels(self) -> ProxyArray | None: + """Terrain level per environment, shape ``(num_envs,)`` of ``wp.int64``. + + This is ``None`` when environment origins are configured from grid spacing. + """ + return self._terrain_levels + + @property + def terrain_types(self) -> ProxyArray | None: + """Terrain type per environment, shape ``(num_envs,)`` of ``wp.int64``. + + This is ``None`` when environment origins are configured from grid spacing. + """ + return self._terrain_types + """ Operations - Visibility. """ @@ -346,20 +353,12 @@ def configure_env_origins(self, origins: np.ndarray | torch.Tensor | None = None """ # decide whether to compute origins in a grid or based on curriculum if origins is not None: - origins_torch = torch.as_tensor(origins, dtype=torch.float32, device=self.device).contiguous() - # ``wp.from_torch`` is zero-copy and keeps the source tensor alive through the Warp array. - self.terrain_origins = ProxyArray(wp.from_torch(origins_torch, dtype=wp.vec3f)) - # compute environment origins - self.env_origins = self._compute_env_origins_curriculum(self.cfg.num_envs, self.terrain_origins) + self._configure_env_origins_curriculum(self.cfg.num_envs, origins) else: - self.terrain_origins = None - self.terrain_levels = None - self.terrain_types = None # check if env spacing is valid if self.cfg.env_spacing is None: raise ValueError("Environment spacing must be specified for configuring grid-like origins.") - # compute environment origins - self.env_origins = self._compute_env_origins_grid(self.cfg.num_envs, self.cfg.env_spacing) + self._configure_env_origins_grid(self.cfg.num_envs, self.cfg.env_spacing) def update_env_origins(self, env_ids: torch.Tensor, move_up: torch.Tensor, move_down: torch.Tensor) -> None: """Update environment origins through the Torch ID-based interface. @@ -448,10 +447,13 @@ def update_env_origins_mask( Internal helpers. """ - def _compute_env_origins_curriculum(self, num_envs: int, origins: ProxyArray) -> ProxyArray: - """Compute the origins of the environments defined by the sub-terrains origins.""" + def _configure_env_origins_curriculum(self, num_envs: int, origins: np.ndarray | torch.Tensor) -> None: + """Allocate and initialize origin buffers from sub-terrain origins.""" + origins_torch = torch.as_tensor(origins, dtype=torch.float32, device=self.device).contiguous() + # ``wp.from_torch`` is zero-copy and keeps the source tensor alive through the Warp array. + self._terrain_origins = ProxyArray(wp.from_torch(origins_torch, dtype=wp.vec3f)) # extract number of rows and cols - num_rows, num_cols = origins.shape[:2] + num_rows, num_cols = self._terrain_origins.shape[:2] # maximum initial level possible for the terrains if self.cfg.max_init_terrain_level is None: max_init_level = num_rows - 1 @@ -464,18 +466,21 @@ def _compute_env_origins_curriculum(self, num_envs: int, origins: ProxyArray) -> terrain_types = torch.div( torch.arange(num_envs, device=self.device), (num_envs / num_cols), rounding_mode="floor" ).to(torch.long) - env_origins = origins.torch[terrain_levels, terrain_types].contiguous() + env_origins = self._terrain_origins.torch[terrain_levels, terrain_types].contiguous() - self.terrain_levels = ProxyArray(wp.from_torch(terrain_levels, dtype=wp.int64)) - self.terrain_types = ProxyArray(wp.from_torch(terrain_types, dtype=wp.int64)) - return ProxyArray(wp.from_torch(env_origins, dtype=wp.vec3f)) + self._terrain_levels = ProxyArray(wp.from_torch(terrain_levels, dtype=wp.int64)) + self._terrain_types = ProxyArray(wp.from_torch(terrain_types, dtype=wp.int64)) + self._env_origins = ProxyArray(wp.from_torch(env_origins, dtype=wp.vec3f)) - def _compute_env_origins_grid(self, num_envs: int, env_spacing: float) -> ProxyArray: - """Compute the origins of the environments in a grid based on configured spacing.""" + def _configure_env_origins_grid(self, num_envs: int, env_spacing: float) -> None: + """Allocate and initialize origin buffers from grid spacing.""" from isaaclab.cloner import grid_transforms env_origins, _ = grid_transforms(num_envs, env_spacing, device=self.device) - return ProxyArray(wp.from_torch(env_origins.contiguous(), dtype=wp.vec3f)) + self._terrain_origins = None + self._terrain_levels = None + self._terrain_types = None + self._env_origins = ProxyArray(wp.from_torch(env_origins.contiguous(), dtype=wp.vec3f)) """ Deprecated. From 6952c082656486ceead8f480fa8a97851a646737 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Mon, 20 Jul 2026 01:20:34 -0700 Subject: [PATCH 42/58] Simplify terrain Warp origin views Keep the existing Torch terrain API while caching zero-copy Warp views during origin configuration. Route shared origin access through the Warp scene without changing the stable scene API. --- docs/source/how-to/proxy_array.rst | 13 +- .../migration/migrating_to_isaaclab_3-0.rst | 5 - scripts/demos/procedural_terrain.py | 2 +- scripts/tutorials/06_deploy/anymal_c_env.py | 2 +- .../jichuanh-warp-frontend-cleanup.minor.rst | 18 +-- .../envs/mdp/commands/pose_2d_command.py | 2 +- source/isaaclab/isaaclab/envs/mdp/events.py | 2 +- .../isaaclab/scene/interactive_scene.py | 19 +-- .../isaaclab/terrains/terrain_importer.py | 152 +++++++++--------- .../test/terrains/check_terrain_importer.py | 2 +- .../test/terrains/test_terrain_importer.py | 17 +- .../jichuanh-warp-frontend-cleanup.minor.rst | 2 +- .../envs/interactive_scene_warp.py | 19 ++- .../test/envs/test_interactive_scene_warp.py | 10 +- .../test/managers/test_curriculum_manager.py | 112 ++++++------- .../contrib/anymal_c_direct/anymal_c_env.py | 2 +- .../core/dexsuite/mdp/events.py | 2 +- .../core/velocity/mdp/curriculums.py | 2 +- .../locomotion/velocity/mdp/curriculums.py | 4 +- 19 files changed, 181 insertions(+), 206 deletions(-) diff --git a/docs/source/how-to/proxy_array.rst b/docs/source/how-to/proxy_array.rst index 48459e1455dd..7bf0a90374f3 100644 --- a/docs/source/how-to/proxy_array.rst +++ b/docs/source/how-to/proxy_array.rst @@ -5,9 +5,9 @@ Working with ProxyArray .. currentmodule:: isaaclab.utils.warp -Isaac Lab data classes and :class:`~isaaclab.terrains.TerrainImporter` buffers return -:class:`ProxyArray` — a lightweight, warp-first wrapper that provides zero-copy access as either a -:class:`warp.array` or a :class:`torch.Tensor`. +Isaac Lab data classes return :class:`ProxyArray` — a lightweight, warp-first wrapper that +provides zero-copy access to simulation data as either a :class:`warp.array` or a +:class:`torch.Tensor`. .. note:: @@ -20,9 +20,7 @@ Quick Start ~~~~~~~~~~~ Every property on asset and sensor data classes (e.g., ``robot.data.joint_pos``, -``sensor.data.net_forces_w``) returns a ``ProxyArray``. The ``terrain_origins``, -``env_origins``, ``terrain_levels``, and ``terrain_types`` buffers on ``TerrainImporter`` follow -the same interface: +``sensor.data.net_forces_w``) returns a ``ProxyArray``: .. code-block:: python @@ -38,9 +36,6 @@ the same interface: # Pass directly to warp kernels — no unwrapping needed wp.launch(my_kernel, inputs=[robot.data.joint_pos], ...) # works via __cuda_array_interface__ - # Terrain buffers use the same explicit accessors - terrain_levels = env.scene.terrain.terrain_levels.torch - The ``.torch`` and ``.warp`` Accessors ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/docs/source/migration/migrating_to_isaaclab_3-0.rst b/docs/source/migration/migrating_to_isaaclab_3-0.rst index 566781b413d9..0f9580f8fe3a 100644 --- a/docs/source/migration/migrating_to_isaaclab_3-0.rst +++ b/docs/source/migration/migrating_to_isaaclab_3-0.rst @@ -1269,11 +1269,6 @@ change applies to all asset classes (:class:`~isaaclab.assets.Articulation`, (:class:`~isaaclab_physx.sensors.ContactSensor`, :class:`~isaaclab_physx.sensors.Imu`, :class:`~isaaclab_physx.sensors.Pva`, :class:`~isaaclab_physx.sensors.FrameTransformer`). -The terrain-origin fields on :class:`~isaaclab.terrains.TerrainImporter` follow the same -interface. :attr:`~isaaclab.scene.InteractiveScene.env_origins` remains a ``torch.Tensor`` -compatibility boundary; use :attr:`~isaaclab.scene.InteractiveScene.env_origins_wp` for its -Warp view. - To use a data property as a ``torch.Tensor``, append ``.torch``: .. code-block:: python diff --git a/scripts/demos/procedural_terrain.py b/scripts/demos/procedural_terrain.py index 6dd0c2f57bf6..6e26a8b26536 100644 --- a/scripts/demos/procedural_terrain.py +++ b/scripts/demos/procedural_terrain.py @@ -145,7 +145,7 @@ def design_scene() -> tuple[dict, torch.Tensor]: # return the scene information scene_entities = {"terrain": terrain_importer} - return scene_entities, terrain_importer.env_origins.torch + return scene_entities, terrain_importer.env_origins def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, AssetBase], origins: torch.Tensor): diff --git a/scripts/tutorials/06_deploy/anymal_c_env.py b/scripts/tutorials/06_deploy/anymal_c_env.py index 25194594c4ad..2f0003fcfe23 100644 --- a/scripts/tutorials/06_deploy/anymal_c_env.py +++ b/scripts/tutorials/06_deploy/anymal_c_env.py @@ -195,7 +195,7 @@ def _reset_idx(self, env_ids: torch.Tensor | None): joint_vel = self._robot.data.default_joint_vel.torch[env_ids] default_root_pose = self._robot.data.default_root_pose.torch[env_ids] default_root_vel = self._robot.data.default_root_vel.torch[env_ids] - default_root_pose[:, :3] += self._terrain.env_origins.torch[env_ids] + default_root_pose[:, :3] += self._terrain.env_origins[env_ids] self._robot.write_root_pose_to_sim_index(root_pose=default_root_pose, env_ids=env_ids) self._robot.write_root_velocity_to_sim_index(root_velocity=default_root_vel, env_ids=env_ids) self._robot.write_joint_position_to_sim_index(position=joint_pos, env_ids=env_ids) diff --git a/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst index dbf6d1927d0b..cf5e9d3e6841 100644 --- a/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst +++ b/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -1,20 +1,10 @@ Added ^^^^^ -* Added :attr:`~isaaclab.scene.InteractiveScene.env_origins_wp` for zero-copy Warp access to - scene origins. -* Added :meth:`~isaaclab.terrains.TerrainImporter.update_env_origins_mask` for mask-based Warp - terrain curriculum updates. - -Changed -^^^^^^^ - -* Changed :attr:`~isaaclab.terrains.TerrainImporter.terrain_origins`, - :attr:`~isaaclab.terrains.TerrainImporter.env_origins`, - :attr:`~isaaclab.terrains.TerrainImporter.terrain_levels`, and - :attr:`~isaaclab.terrains.TerrainImporter.terrain_types` to return - :class:`~isaaclab.utils.warp.ProxyArray`. Use ``.torch`` for Torch operations - and ``.warp`` for Warp kernels. +* Added :attr:`~isaaclab.terrains.TerrainImporter.env_origins_wp`, + :attr:`~isaaclab.terrains.TerrainImporter.terrain_levels_wp`, and + :meth:`~isaaclab.terrains.TerrainImporter.update_env_origins_mask` for zero-copy Warp terrain + access and mask-based curriculum updates. Fixed ^^^^^ diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py index 9d3a3803925a..ffd7bdad5e49 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py @@ -204,7 +204,7 @@ def _resample_command(self, env_ids: Sequence[int]): # sample new position targets from the terrain ids = torch.randint(0, self.valid_targets.shape[2], size=(len(env_ids),), device=self.device) self.pos_command_w[env_ids] = self.valid_targets[ - self.terrain.terrain_levels.torch[env_ids], self.terrain.terrain_types.torch[env_ids], ids + self.terrain.terrain_levels[env_ids], self.terrain.terrain_types[env_ids], ids ] # offset the position command by the current root height self.pos_command_w[env_ids, 2] += self.robot.data.default_root_pose.torch[env_ids, 2] diff --git a/source/isaaclab/isaaclab/envs/mdp/events.py b/source/isaaclab/isaaclab/envs/mdp/events.py index b0c7d9d9e988..936a79adebd8 100644 --- a/source/isaaclab/isaaclab/envs/mdp/events.py +++ b/source/isaaclab/isaaclab/envs/mdp/events.py @@ -2000,7 +2000,7 @@ def reset_root_state_from_terrain( # sample random valid poses ids = torch.randint(0, valid_positions.shape[2], size=(len(env_ids),), device=env.device) - positions = valid_positions[terrain.terrain_levels.torch[env_ids], terrain.terrain_types.torch[env_ids], ids] + positions = valid_positions[terrain.terrain_levels[env_ids], terrain.terrain_types[env_ids], ids] positions += asset.data.default_root_pose.torch[env_ids, :3] # sample random orientations diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index e6d8405ada29..15b693197a46 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -36,7 +36,6 @@ from isaaclab.sim import SimulationContext from isaaclab.sim.utils.stage import get_current_stage, get_current_stage_id from isaaclab.sim.views import FrameView -from isaaclab.utils.warp import ProxyArray # Note: This is a temporary import for the VisuoTactileSensorCfg class. # It will be removed once the VisuoTactileSensor class is added to the core Isaac Lab framework. @@ -189,11 +188,6 @@ def __init__(self, cfg: InteractiveSceneCfg): if self._is_scene_setup_from_cfg(): self._add_entities_from_cfg() - if self._terrain is None: - # Scene origins are immutable after construction, so one ProxyArray safely serves both frontends. - positions = self.sim.get_clone_plan().positions - self._env_origins = ProxyArray(wp.from_torch(positions.contiguous(), dtype=wp.vec3f)) - self._aggregate_scene_data_requirements(requested_viz_types) # Collision filtering is PhysX-only (matches both physx and ovphysx). @@ -370,19 +364,12 @@ def num_envs(self) -> int: @property def env_origins(self) -> torch.Tensor: - """Per-environment world origins [m], shape ``(num_envs, 3)``. From the terrain when registered, + """Per-env world origins, shape ``(num_envs, 3)``. From the terrain when registered, else from the published :class:`~isaaclab.cloner.ClonePlan`. """ if self._terrain is not None: - return self._terrain.env_origins.torch - return self._env_origins.torch - - @property - def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): - """Zero-copy Warp view of environment world origins [m], shape ``(num_envs,)`` of ``wp.vec3f``.""" - if self._terrain is not None: - return self._terrain.env_origins.warp - return self._env_origins.warp + return self._terrain.env_origins + return self.sim.get_clone_plan().positions @property def terrain(self) -> TerrainImporter | None: diff --git a/source/isaaclab/isaaclab/terrains/terrain_importer.py b/source/isaaclab/isaaclab/terrains/terrain_importer.py index c9d740b438fe..d9523497016f 100644 --- a/source/isaaclab/isaaclab/terrains/terrain_importer.py +++ b/source/isaaclab/isaaclab/terrains/terrain_importer.py @@ -16,7 +16,6 @@ import isaaclab.sim as sim_utils from isaaclab.markers import VisualizationMarkers from isaaclab.markers.config import FRAME_MARKER_CFG -from isaaclab.utils.warp import ProxyArray from .utils import create_prim_from_mesh @@ -75,6 +74,16 @@ class TerrainImporter: terrain_prim_paths: list[str] """A list containing the USD prim paths to the imported terrains.""" + terrain_origins: torch.Tensor | None + """Sub-terrain origins [m], shape ``(num_rows, num_cols, 3)``. + + If terrain origins is not None, the environment origins are computed based on the terrain origins. + Otherwise, the environment origins are computed based on the grid spacing. + """ + + env_origins: torch.Tensor + """Environment origins [m], shape ``(num_envs, 3)``.""" + def __init__(self, cfg: TerrainImporterCfg): """Initialize the terrain importer. @@ -95,10 +104,12 @@ def __init__(self, cfg: TerrainImporterCfg): # create buffers for the terrains self.terrain_prim_paths = list() - self._terrain_origins: ProxyArray | None = None - self._env_origins: ProxyArray - self._terrain_levels: ProxyArray | None = None - self._terrain_types: ProxyArray | None = None + self.terrain_origins = None + self.env_origins = None # assigned later when `configure_env_origins` is called + self._terrain_levels_wp = None + self._terrain_types_wp = None + self._terrain_origins_wp = None + self._env_origins_wp = None # private variables self._terrain_flat_patches = dict() @@ -167,37 +178,16 @@ def terrain_names(self) -> list[str]: return [f"'{path.split('/')[-1]}'" for path in self.terrain_prim_paths] @property - def terrain_origins(self) -> ProxyArray | None: - """Sub-terrain origins [m], shape ``(num_rows, num_cols)`` of ``wp.vec3f``. - - The :attr:`~isaaclab.utils.warp.ProxyArray.torch` view has shape ``(num_rows, num_cols, 3)``. - This is ``None`` when environment origins are configured from grid spacing. - """ - return self._terrain_origins - - @property - def env_origins(self) -> ProxyArray: - """Environment origins [m], shape ``(num_envs,)`` of ``wp.vec3f``. - - The :attr:`~isaaclab.utils.warp.ProxyArray.torch` view has shape ``(num_envs, 3)``. - """ - return self._env_origins - - @property - def terrain_levels(self) -> ProxyArray | None: - """Terrain level per environment, shape ``(num_envs,)`` of ``wp.int64``. - - This is ``None`` when environment origins are configured from grid spacing. - """ - return self._terrain_levels + def terrain_levels_wp(self) -> wp.array(dtype=wp.int64): + """Cached zero-copy Warp view of terrain levels, shape ``(num_envs,)``.""" + if self._terrain_levels_wp is None: + raise RuntimeError("Terrain levels are unavailable for grid-like environment origins.") + return self._terrain_levels_wp @property - def terrain_types(self) -> ProxyArray | None: - """Terrain type per environment, shape ``(num_envs,)`` of ``wp.int64``. - - This is ``None`` when environment origins are configured from grid spacing. - """ - return self._terrain_types + def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): + """Cached zero-copy Warp view of environment origins [m], shape ``(num_envs,)``.""" + return self._env_origins_wp """ Operations - Visibility. @@ -223,9 +213,11 @@ def set_debug_vis(self, debug_vis: bool) -> bool: cfg=FRAME_MARKER_CFG.replace(prim_path="/Visuals/TerrainOrigin") ) if self.terrain_origins is not None: - self.origin_visualizer.visualize(self.terrain_origins.torch.reshape(-1, 3)) + self.origin_visualizer.visualize(self.terrain_origins.reshape(-1, 3)) + elif self.env_origins is not None: + self.origin_visualizer.visualize(self.env_origins.reshape(-1, 3)) else: - self.origin_visualizer.visualize(self.env_origins.torch.reshape(-1, 3)) + raise RuntimeError("Terrain origins are not configured.") # set visibility self.origin_visualizer.set_visibility(True) else: @@ -346,19 +338,24 @@ def configure_env_origins(self, origins: np.ndarray | torch.Tensor | None = None Args: origins: The origins [m] of the sub-terrains. Shape is (num_rows, num_cols, 3). - - Note: - This setup operation replaces the :class:`~isaaclab.utils.warp.ProxyArray` instances. Call it before - managers or kernels cache their ``.warp`` views. """ # decide whether to compute origins in a grid or based on curriculum if origins is not None: - self._configure_env_origins_curriculum(self.cfg.num_envs, origins) + # convert to torch + if isinstance(origins, np.ndarray): + origins = torch.from_numpy(origins) + # store the origins + self.terrain_origins = origins.to(self.device, dtype=torch.float).contiguous() + # compute environment origins + self.env_origins = self._compute_env_origins_curriculum(self.cfg.num_envs, self.terrain_origins) else: + self.terrain_origins = None # check if env spacing is valid if self.cfg.env_spacing is None: raise ValueError("Environment spacing must be specified for configuring grid-like origins.") - self._configure_env_origins_grid(self.cfg.num_envs, self.cfg.env_spacing) + # compute environment origins + self.env_origins = self._compute_env_origins_grid(self.cfg.num_envs, self.cfg.env_spacing) + self._configure_warp_origin_views() def update_env_origins(self, env_ids: torch.Tensor, move_up: torch.Tensor, move_down: torch.Tensor) -> None: """Update environment origins through the Torch ID-based interface. @@ -369,23 +366,19 @@ def update_env_origins(self, env_ids: torch.Tensor, move_up: torch.Tensor, move_ move_down: Flags that decrement the selected terrain levels. """ # check if grid-like spawning - if self.terrain_origins is None or self.terrain_levels is None or self.terrain_types is None: + if self.terrain_origins is None: return - terrain_origins = self.terrain_origins.torch - terrain_levels = self.terrain_levels.torch - terrain_types = self.terrain_types.torch - env_origins = self.env_origins.torch # update terrain level for the envs - terrain_levels[env_ids] += 1 * move_up - 1 * move_down + self.terrain_levels[env_ids] += 1 * move_up - 1 * move_down # robots that solve the last level are sent to a random one # the minimum level is zero - terrain_levels[env_ids] = torch.where( - terrain_levels[env_ids] >= self.max_terrain_level, - torch.randint_like(terrain_levels[env_ids], self.max_terrain_level), - torch.clip(terrain_levels[env_ids], 0), + self.terrain_levels[env_ids] = torch.where( + self.terrain_levels[env_ids] >= self.max_terrain_level, + torch.randint_like(self.terrain_levels[env_ids], self.max_terrain_level), + torch.clip(self.terrain_levels[env_ids], 0), ) # update the env origins - env_origins[env_ids] = terrain_origins[terrain_levels[env_ids], terrain_types[env_ids]] + self.env_origins[env_ids] = self.terrain_origins[self.terrain_levels[env_ids], self.terrain_types[env_ids]] def update_env_origins_mask( self, @@ -409,7 +402,7 @@ def update_env_origins_mask( TypeError: If an input is not a Warp array with the required data type. ValueError: If an input has the wrong shape or device. """ - if self.terrain_origins is None or self.terrain_levels is None or self.terrain_types is None: + if self.terrain_origins is None: return num_envs = self.terrain_levels.shape[0] arrays = { @@ -434,10 +427,10 @@ def update_env_origins_mask( move_up, move_down, rng_state, - self.terrain_levels.warp, - self.terrain_types.warp, - self.terrain_origins.warp, - self.env_origins.warp, + self._terrain_levels_wp, + self._terrain_types_wp, + self._terrain_origins_wp, + self._env_origins_wp, self.max_terrain_level, ], device=self.device, @@ -447,13 +440,22 @@ def update_env_origins_mask( Internal helpers. """ - def _configure_env_origins_curriculum(self, num_envs: int, origins: np.ndarray | torch.Tensor) -> None: - """Allocate and initialize origin buffers from sub-terrain origins.""" - origins_torch = torch.as_tensor(origins, dtype=torch.float32, device=self.device).contiguous() - # ``wp.from_torch`` is zero-copy and keeps the source tensor alive through the Warp array. - self._terrain_origins = ProxyArray(wp.from_torch(origins_torch, dtype=wp.vec3f)) + def _configure_warp_origin_views(self) -> None: + """Create persistent zero-copy Warp views after configuring the Torch buffers.""" + self._env_origins_wp = wp.from_torch(self.env_origins, dtype=wp.vec3f) + if self.terrain_origins is None: + self._terrain_levels_wp = None + self._terrain_types_wp = None + self._terrain_origins_wp = None + return + self._terrain_levels_wp = wp.from_torch(self.terrain_levels, dtype=wp.int64) + self._terrain_types_wp = wp.from_torch(self.terrain_types, dtype=wp.int64) + self._terrain_origins_wp = wp.from_torch(self.terrain_origins, dtype=wp.vec3f) + + def _compute_env_origins_curriculum(self, num_envs: int, origins: torch.Tensor) -> torch.Tensor: + """Compute the origins of the environments defined by the sub-terrains origins.""" # extract number of rows and cols - num_rows, num_cols = self._terrain_origins.shape[:2] + num_rows, num_cols = origins.shape[:2] # maximum initial level possible for the terrains if self.cfg.max_init_terrain_level is None: max_init_level = num_rows - 1 @@ -462,25 +464,21 @@ def _configure_env_origins_curriculum(self, num_envs: int, origins: np.ndarray | # store maximum terrain level possible self.max_terrain_level = num_rows # define all terrain levels and types available - terrain_levels = torch.randint(0, max_init_level + 1, (num_envs,), device=self.device) - terrain_types = torch.div( + self.terrain_levels = torch.randint(0, max_init_level + 1, (num_envs,), device=self.device) + self.terrain_types = torch.div( torch.arange(num_envs, device=self.device), (num_envs / num_cols), rounding_mode="floor" ).to(torch.long) - env_origins = self._terrain_origins.torch[terrain_levels, terrain_types].contiguous() - - self._terrain_levels = ProxyArray(wp.from_torch(terrain_levels, dtype=wp.int64)) - self._terrain_types = ProxyArray(wp.from_torch(terrain_types, dtype=wp.int64)) - self._env_origins = ProxyArray(wp.from_torch(env_origins, dtype=wp.vec3f)) + # create tensor based on number of environments + env_origins = torch.zeros(num_envs, 3, device=self.device) + env_origins[:] = origins[self.terrain_levels, self.terrain_types] + return env_origins - def _configure_env_origins_grid(self, num_envs: int, env_spacing: float) -> None: - """Allocate and initialize origin buffers from grid spacing.""" + def _compute_env_origins_grid(self, num_envs: int, env_spacing: float) -> torch.Tensor: + """Compute the origins of the environments in a grid based on configured spacing.""" from isaaclab.cloner import grid_transforms env_origins, _ = grid_transforms(num_envs, env_spacing, device=self.device) - self._terrain_origins = None - self._terrain_levels = None - self._terrain_types = None - self._env_origins = ProxyArray(wp.from_torch(env_origins.contiguous(), dtype=wp.vec3f)) + return env_origins """ Deprecated. diff --git a/source/isaaclab/test/terrains/check_terrain_importer.py b/source/isaaclab/test/terrains/check_terrain_importer.py index 7861f5c2b272..901cfc420be3 100644 --- a/source/isaaclab/test/terrains/check_terrain_importer.py +++ b/source/isaaclab/test/terrains/check_terrain_importer.py @@ -156,7 +156,7 @@ def main(): # Set ball positions over terrain origins using FrameView (before simulation starts) xform_view = sim_utils.FrameView("/World/envs/env_.*/ball") # cache initial state of the balls - ball_initial_positions = terrain_importer.env_origins.torch.clone() + ball_initial_positions = terrain_importer.env_origins.clone() ball_initial_positions[:, 2] += 5.0 # set initial poses (writes to USD before simulation) with xform_view.xform_world_space_writer() as w: diff --git a/source/isaaclab/test/terrains/test_terrain_importer.py b/source/isaaclab/test/terrains/test_terrain_importer.py index 83192a472e1b..4ced50804d2e 100644 --- a/source/isaaclab/test/terrains/test_terrain_importer.py +++ b/source/isaaclab/test/terrains/test_terrain_importer.py @@ -29,7 +29,6 @@ from isaaclab.terrains import TerrainImporter, TerrainImporterCfg from isaaclab.terrains.config.rough import ROUGH_TERRAINS_CFG from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR -from isaaclab.utils.warp import ProxyArray pytestmark = pytest.mark.integration @@ -51,15 +50,15 @@ def test_terrain_importer_env_origins(device, env_spacing, num_envs): ) terrain_importer = TerrainImporter(terrain_importer_cfg) # obtain env origins using terrain importer - terrain_importer_origins = terrain_importer.env_origins.torch + terrain_importer_origins = terrain_importer.env_origins - # check that the canonical Warp storage and Torch view share memory - assert isinstance(terrain_importer.env_origins, ProxyArray) - assert terrain_importer.env_origins.warp.dtype == wp.vec3f - assert terrain_importer.env_origins.warp.shape == (num_envs,) + # check that the cached Warp view shares the canonical Torch storage + assert isinstance(terrain_importer.env_origins, torch.Tensor) + assert terrain_importer.env_origins_wp.dtype == wp.vec3f + assert terrain_importer.env_origins_wp.shape == (num_envs,) assert terrain_importer_origins.shape == (num_envs, 3) - assert terrain_importer.env_origins.warp is terrain_importer.env_origins.warp - assert terrain_importer_origins.data_ptr() == terrain_importer.env_origins.warp.ptr + assert terrain_importer.env_origins_wp is terrain_importer.env_origins_wp + assert terrain_importer_origins.data_ptr() == terrain_importer.env_origins_wp.ptr # obtain env origins using Lab's grid_transforms lab_grid_origins, _ = lab_cloner.grid_transforms(num_envs, spacing=env_spacing, device=sim.device) @@ -323,7 +322,7 @@ def _populate_scene(sim: SimulationContext, num_balls: int = 2048, geom_sphere: # Create a view over all the balls using Isaac Lab's FrameView ball_view = sim_utils.FrameView("/World/envs/env_.*/ball") # cache initial state of the balls - ball_initial_positions = terrain_importer.env_origins.torch.clone() + ball_initial_positions = terrain_importer.env_origins.clone() ball_initial_positions[:, 2] += 5.0 # set initial poses # note: setting here writes to USD :) diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst index 2a6540bb8518..97452d26225d 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -6,7 +6,7 @@ Added * Added :class:`~isaaclab_experimental.managers.CurriculumManager` and a boolean-mask reset path for Warp manager-based environments, retaining compact environment IDs only for legacy host consumers. -* Added Warp-native terrain curricula backed by the canonical Warp-owned +* Added Warp-native terrain curricula backed by cached zero-copy Warp views of terrain and scene origins. Changed diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py index 6722599f9032..6905988499e3 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py @@ -11,7 +11,7 @@ import warp as wp -from isaaclab.scene import InteractiveScene +from isaaclab.scene import InteractiveScene, InteractiveSceneCfg class InteractiveSceneWarp(InteractiveScene): @@ -21,6 +21,23 @@ class InteractiveSceneWarp(InteractiveScene): avoiding the need to convert between env_ids and masks. """ + def __init__(self, cfg: InteractiveSceneCfg): + """Initialize the Warp scene and cache non-terrain environment origins. + + Args: + cfg: Configuration for the interactive scene. + """ + super().__init__(cfg) + if self.terrain is None: + self._env_origins_wp = wp.from_torch(self.env_origins, dtype=wp.vec3f) + + @property + def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): + """Cached zero-copy Warp view of environment origins [m], shape ``(num_envs,)``.""" + if self.terrain is not None: + return self.terrain.env_origins_wp + return self._env_origins_wp + def reset( self, env_ids: Sequence[int] | None = None, diff --git a/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py b/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py index cc38406ad190..edb4051aec58 100644 --- a/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py +++ b/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py @@ -17,7 +17,6 @@ from isaaclab.scene import InteractiveScene from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils.warp import ProxyArray class TestInteractiveSceneWarp: @@ -30,8 +29,9 @@ def test_frontend_uses_canonical_terrain_storage_without_mutating_cfg( terrain_cfg = TerrainImporterCfg(prim_path="/World/Ground", terrain_type="plane") cfg = SimpleNamespace(terrain=terrain_cfg) class_type = terrain_cfg.class_type - origins_wp = wp.zeros(2, dtype=wp.vec3f, device="cpu") - terrain = SimpleNamespace(env_origins=ProxyArray(origins_wp)) + origins = torch.zeros((2, 3), dtype=torch.float32) + origins_wp = wp.from_torch(origins, dtype=wp.vec3f) + terrain = SimpleNamespace(env_origins=origins, env_origins_wp=origins_wp) def initialize_scene(scene: InteractiveScene, scene_cfg: SimpleNamespace) -> None: scene.cfg = scene_cfg @@ -43,10 +43,10 @@ def initialize_scene(scene: InteractiveScene, scene_cfg: SimpleNamespace) -> Non assert terrain_cfg.class_type is class_type assert isinstance(scene.env_origins, torch.Tensor) - assert scene.env_origins is terrain.env_origins.torch + assert scene.env_origins is terrain.env_origins assert scene.env_origins_wp is origins_wp assert InteractiveSceneWarp.env_origins is InteractiveScene.env_origins - assert InteractiveSceneWarp.env_origins_wp is InteractiveScene.env_origins_wp + assert "env_origins_wp" not in InteractiveScene.__dict__ def test_mask_reset_stays_mask_based_for_supported_entities(self): """Mask-based reset should not pass environment IDs to Warp-capable entities.""" diff --git a/source/isaaclab_experimental/test/managers/test_curriculum_manager.py b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py index 73c33a28fce7..74a1b847ff92 100644 --- a/source/isaaclab_experimental/test/managers/test_curriculum_manager.py +++ b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py @@ -21,7 +21,6 @@ from isaaclab.managers import ManagerTermBase as StableManagerTermBase from isaaclab.terrains import TerrainImporter -from isaaclab.utils.warp import ProxyArray class _GlobalCurriculumTerm(StableManagerTermBase): @@ -124,11 +123,11 @@ def __init__(self, robot: _Robot, terrain: TerrainImporter): @property def env_origins(self) -> torch.Tensor: - return self.terrain.env_origins.torch + return self.terrain.env_origins @property def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): - return self.terrain.env_origins.warp + return self.terrain.env_origins_wp def __getitem__(self, name: str): return self._entities[name] @@ -175,53 +174,48 @@ def _make_terrain(levels: list[int]) -> TerrainImporter: float(terrain_type), ] terrain.configure_env_origins(origin_values) - terrain.terrain_levels.torch.copy_(torch.from_numpy(level_values)) - terrain.terrain_types.torch.copy_(torch.from_numpy(type_values)) - terrain.env_origins.torch.copy_( - terrain.terrain_origins.torch[terrain.terrain_levels.torch, terrain.terrain_types.torch] - ) + terrain.terrain_levels.copy_(torch.from_numpy(level_values)) + terrain.terrain_types.copy_(torch.from_numpy(type_values)) + terrain.env_origins.copy_(terrain.terrain_origins[terrain.terrain_levels, terrain.terrain_types]) return terrain -def test_terrain_importer_exposes_persistent_proxy_array_views(): - """Canonical terrain storage should expose persistent zero-copy Torch and Warp views.""" +def test_terrain_importer_exposes_persistent_warp_views(): + """Terrain importer should retain one zero-copy Warp view for each Torch buffer.""" terrain = _make_terrain([0, 1, 2, 1]) - assert isinstance(terrain.terrain_origins, ProxyArray) - assert isinstance(terrain.env_origins, ProxyArray) - assert isinstance(terrain.terrain_levels, ProxyArray) - assert isinstance(terrain.terrain_types, ProxyArray) - assert terrain.terrain_origins.dtype == wp.vec3f - assert terrain.env_origins.dtype == wp.vec3f - assert terrain.terrain_levels.dtype == wp.int64 - assert terrain.terrain_types.dtype == wp.int64 - assert terrain.terrain_origins.torch.data_ptr() == terrain.terrain_origins.warp.ptr - assert terrain.env_origins.torch.data_ptr() == terrain.env_origins.warp.ptr - assert terrain.terrain_levels.torch.data_ptr() == terrain.terrain_levels.warp.ptr - assert terrain.terrain_types.torch.data_ptr() == terrain.terrain_types.warp.ptr - - terrain.terrain_levels.torch[0] = 2 - assert terrain.terrain_levels.warp.numpy()[0] == 2 + assert isinstance(terrain.terrain_origins, torch.Tensor) + assert isinstance(terrain.env_origins, torch.Tensor) + assert isinstance(terrain.terrain_levels, torch.Tensor) + assert isinstance(terrain.terrain_types, torch.Tensor) + assert terrain.terrain_origins.data_ptr() == terrain._terrain_origins_wp.ptr + assert terrain.env_origins.data_ptr() == terrain.env_origins_wp.ptr + assert terrain.terrain_levels.data_ptr() == terrain.terrain_levels_wp.ptr + assert terrain.terrain_types.data_ptr() == terrain._terrain_types_wp.ptr + + terrain.terrain_levels[0] = 2 + assert terrain.terrain_levels_wp.numpy()[0] == 2 replacement_levels = wp.array([1, 0, 2, 1], dtype=wp.int64, device="cpu") - wp.copy(terrain.terrain_levels.warp, replacement_levels) + wp.copy(terrain.terrain_levels_wp, replacement_levels) wp.synchronize() - torch.testing.assert_close(terrain.terrain_levels.torch, torch.tensor([1, 0, 2, 1])) + torch.testing.assert_close(terrain.terrain_levels, torch.tensor([1, 0, 2, 1])) -def test_terrain_importer_grid_origins_use_proxy_array_without_curriculum_buffers(): - """Grid origins should use the same proxy contract without curriculum buffers.""" +def test_terrain_importer_grid_origins_cache_only_environment_view(): + """Grid origins should cache a Warp view without allocating curriculum buffers.""" terrain = TerrainImporter.__new__(TerrainImporter) terrain.device = "cpu" terrain.cfg = SimpleNamespace(num_envs=4, env_spacing=2.0) terrain.configure_env_origins() assert terrain.terrain_origins is None - assert isinstance(terrain.env_origins, ProxyArray) - assert terrain.env_origins.shape == (4,) - assert terrain.env_origins.torch.shape == (4, 3) - assert terrain.env_origins.torch.data_ptr() == terrain.env_origins.warp.ptr - assert terrain.terrain_levels is None - assert terrain.terrain_types is None + assert isinstance(terrain.env_origins, torch.Tensor) + assert terrain.env_origins.shape == (4, 3) + assert terrain.env_origins_wp.shape == (4,) + assert terrain.env_origins.data_ptr() == terrain.env_origins_wp.ptr + assert terrain._terrain_levels_wp is None + assert terrain._terrain_types_wp is None + assert terrain._terrain_origins_wp is None def test_terrain_importer_mask_update_preserves_sparse_and_stable_semantics(): @@ -232,36 +226,36 @@ def test_terrain_importer_mask_update_preserves_sparse_and_stable_semantics(): move_down = wp.array([True, False, False, True, False, True], dtype=wp.bool, device="cpu") rng_state = wp.array(np.arange(6, dtype=np.uint32) + 71, device="cpu") rng_before = rng_state.numpy().copy() - origins_before = terrain.env_origins.torch.clone() - levels_wp = terrain.terrain_levels.warp - origins_wp = terrain.env_origins.warp + origins_before = terrain.env_origins.clone() + levels_wp = terrain.terrain_levels_wp + origins_wp = terrain.env_origins_wp levels_ptr = levels_wp.ptr origins_ptr = origins_wp.ptr - assert terrain.terrain_levels.torch.data_ptr() == levels_ptr - assert terrain.env_origins.torch.data_ptr() == origins_ptr + assert terrain.terrain_levels.data_ptr() == levels_ptr + assert terrain.env_origins.data_ptr() == origins_ptr terrain.update_env_origins_mask(env_mask, move_up, move_down, rng_state) wp.synchronize() - levels = terrain.terrain_levels.torch.tolist() + levels = terrain.terrain_levels.tolist() assert levels[0] == 0 assert levels[1] == 2 assert 0 <= levels[2] < terrain.max_terrain_level assert levels[3] == 1 assert levels[4:] == [1, 0] torch.testing.assert_close( - terrain.env_origins.torch[:4], - terrain.terrain_origins.torch[terrain.terrain_levels.torch[:4], terrain.terrain_types.torch[:4]], + terrain.env_origins[:4], + terrain.terrain_origins[terrain.terrain_levels[:4], terrain.terrain_types[:4]], ) - torch.testing.assert_close(terrain.env_origins.torch[4:], origins_before[4:]) + torch.testing.assert_close(terrain.env_origins[4:], origins_before[4:]) assert np.all(rng_state.numpy()[:4] != rng_before[:4]) np.testing.assert_array_equal(rng_state.numpy()[4:], rng_before[4:]) - assert terrain.terrain_levels.warp is levels_wp - assert terrain.env_origins.warp is origins_wp - assert terrain.terrain_levels.warp.ptr == levels_ptr - assert terrain.env_origins.warp.ptr == origins_ptr + assert terrain.terrain_levels_wp is levels_wp + assert terrain.env_origins_wp is origins_wp + assert terrain.terrain_levels_wp.ptr == levels_ptr + assert terrain.env_origins_wp.ptr == origins_ptr - levels_before_stable_update = terrain.terrain_levels.torch.clone() + levels_before_stable_update = terrain.terrain_levels.clone() env_ids = torch.tensor([1, 3], dtype=torch.int64) terrain.update_env_origins( env_ids, @@ -270,18 +264,18 @@ def test_terrain_importer_mask_update_preserves_sparse_and_stable_semantics(): ) expected_levels = levels_before_stable_update.clone() expected_levels[env_ids] -= 1 - torch.testing.assert_close(terrain.terrain_levels.torch, expected_levels) - assert terrain.terrain_levels.warp is levels_wp - assert terrain.env_origins.warp is origins_wp - assert terrain.terrain_levels.warp.ptr == levels_ptr - assert terrain.env_origins.warp.ptr == origins_ptr + torch.testing.assert_close(terrain.terrain_levels, expected_levels) + assert terrain.terrain_levels_wp is levels_wp + assert terrain.env_origins_wp is origins_wp + assert terrain.terrain_levels_wp.ptr == levels_ptr + assert terrain.env_origins_wp.ptr == origins_ptr def test_curriculum_manager_updates_masked_levels_and_logs_without_host_compaction(monkeypatch: pytest.MonkeyPatch): """Manager compute/reset should remain mask-native and expose persistent scalar logging.""" assert terrain_levels_vel is not TerrainLevelsVel terrain = _make_terrain([0, 1, 2, 2, 1, 0]) - root_positions = terrain.env_origins.torch.numpy().copy() + root_positions = terrain.env_origins.numpy().copy() root_positions[:, 0] += np.array([5.0, 1.0, 3.0, 5.0, 9.0, 0.0], dtype=np.float32) commands = np.array( [[0.1, 0.0, 0.0], [1.0, 0.0, 0.0], [0.2, 0.0, 0.0], [2.0, 0.0, 0.0], [0.1, 0.0, 0.0], [0.1, 0.0, 0.0]], @@ -318,7 +312,7 @@ def _fail_host_compaction(*args, **kwargs): extras = manager.reset(env_mask) wp.synchronize() - levels = terrain.terrain_levels.torch.tolist() + levels = terrain.terrain_levels.tolist() assert levels[0] == 1 assert levels[1] == 0 assert levels[2] == 2 @@ -328,7 +322,7 @@ def _fail_host_compaction(*args, **kwargs): np.testing.assert_array_equal(term._move_down_wp.numpy(), [False, True, False, False, False, False]) assert extras is extras_ref assert extras["Curriculum/terrain_levels"] is manager.reset_extras["Curriculum/terrain_levels"] - torch.testing.assert_close(extras["Curriculum/terrain_levels"], terrain.terrain_levels.torch.float().mean()) + torch.testing.assert_close(extras["Curriculum/terrain_levels"], terrain.terrain_levels.float().mean()) torch.testing.assert_close(extras["Curriculum/global_state"], torch.tensor(7.0)) assert term._move_up_wp.ptr == move_up_ptr assert term._move_down_wp.ptr == move_down_ptr @@ -364,7 +358,7 @@ def test_registered_reach_curricula_update_weights_without_a_host_boundary(): terrain = _make_terrain([0, 1, 2, 2]) env = _Env( terrain, - terrain.env_origins.torch.numpy().copy(), + terrain.env_origins.numpy().copy(), np.zeros((4, 3), dtype=np.float32), ) env.reward_manager = _RewardManager() @@ -394,7 +388,7 @@ def test_registered_reach_curricula_update_weights_without_a_host_boundary(): def test_curriculum_manager_rejects_structured_legacy_state() -> None: """Unsupported fallback logging should fail explicitly instead of narrowing silently.""" terrain = _make_terrain([0, 1]) - env = _Env(terrain, terrain.env_origins.torch.numpy().copy(), np.zeros((2, 3), dtype=np.float32)) + env = _Env(terrain, terrain.env_origins.numpy().copy(), np.zeros((2, 3), dtype=np.float32)) manager = CurriculumManager( {"structured": CurriculumTermCfg(func=_structured_state, requires_host_ids=False)}, env, diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/anymal_c_direct/anymal_c_env.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/anymal_c_direct/anymal_c_env.py index 11ff1287a23b..f932022481e0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/anymal_c_direct/anymal_c_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/anymal_c_direct/anymal_c_env.py @@ -205,7 +205,7 @@ def _reset_idx(self, env_ids: torch.Tensor | None): joint_vel = self._robot.data.default_joint_vel.torch[env_ids] default_root_pose = self._robot.data.default_root_pose.torch[env_ids] default_root_vel = self._robot.data.default_root_vel.torch[env_ids] - default_root_pose[:, :3] += self._terrain.env_origins.torch[env_ids] + default_root_pose[:, :3] += self._terrain.env_origins[env_ids] self._robot.write_root_pose_to_sim_index(root_pose=default_root_pose, env_ids=env_ids) self._robot.write_root_velocity_to_sim_index(root_velocity=default_root_vel, env_ids=env_ids) self._robot.write_joint_position_to_sim_index(position=joint_pos, env_ids=env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/events.py index 5c6691d5f6df..9f10f18a2ec0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/dexsuite/mdp/events.py @@ -580,7 +580,7 @@ def __init__(self, cfg: SlabClearanceCfg, env: ManagerBasedEnv): wp.array(np.array([x is not None for x, _, _ in slabs], dtype=np.int32), device=device), wp.array(np.array([y is not None for _, y, _ in slabs], dtype=np.int32), device=device), ] - self._env_origins = env.scene.env_origins_wp + self._env_origins = wp.from_torch(env.scene.env_origins.contiguous(), dtype=wp.vec3) def __call__(self, env: ManagerBasedEnv, env_ids: torch.Tensor) -> torch.Tensor: num = len(env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/curriculums.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/curriculums.py index 6b05db44348b..1438cbf03801 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/curriculums.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/curriculums.py @@ -53,4 +53,4 @@ def terrain_levels_vel( # update terrain levels terrain.update_env_origins(env_ids, move_up, move_down) # return the mean terrain level - return torch.mean(terrain.terrain_levels.torch.float()) + return torch.mean(terrain.terrain_levels.float()) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py index 734b496ad18b..f55af9254efa 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py @@ -77,7 +77,7 @@ def __init__(self, cfg: CurriculumTermCfg, env: ManagerBasedRLEnv): self._root_pos_w_wp = self._asset.data.root_pos_w.warp self._env_origins_wp = env.env_origins_wp self._command_wp = env.command_manager.get_command_wp("base_velocity") - self._terrain_levels_wp = self._terrain.terrain_levels.warp + self._terrain_levels_wp = self._terrain.terrain_levels_wp self._move_up_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) self._move_down_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) self._terrain_half_length = float(self._terrain.cfg.terrain_generator.size[0]) * 0.5 @@ -141,4 +141,4 @@ def terrain_levels_vel( move_down = distance < torch.norm(command[env_ids, :2], dim=1) * env.max_episode_length_s * 0.5 move_down *= ~move_up terrain.update_env_origins(env_ids, move_up, move_down) - return torch.mean(terrain.terrain_levels.torch.float()) + return torch.mean(terrain.terrain_levels.float()) From 6f19326bdebba5b7dd83b2748f4ed239b62945f1 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Mon, 20 Jul 2026 03:24:13 -0700 Subject: [PATCH 43/58] Simplify Warp frontend execution Remove the experimental manager routing layer and make Warp managers the direct frontend implementation. Centralize capture eligibility in the environment-owned graph cache while preserving eager fallbacks for unsafe boundaries. --- .../isaaclab/scene/interactive_scene.py | 33 +- .../isaaclab/terrains/terrain_importer.py | 2 +- .../test/scene/test_interactive_scene.py | 55 ++ .../jichuanh-warp-frontend-cleanup.minor.rst | 14 +- .../isaaclab_experimental/envs/__init__.pyi | 2 - .../envs/direct_rl_env_warp.py | 11 +- .../envs/interactive_scene_warp.py | 90 ---- .../envs/manager_based_env_warp.py | 165 +++--- .../envs/manager_based_rl_env_warp.py | 148 +++-- .../isaaclab_experimental/envs/mdp/events.py | 95 +--- .../managers/action_manager.py | 1 + .../managers/command_manager.py | 1 + .../managers/manager_base.py | 16 +- .../isaaclab_experimental/utils/__init__.pyi | 3 - .../utils/manager_call_switch.py | 372 ------------- .../utils/warp_graph_cache.py | 150 +++--- .../test/envs/mdp/test_events_warp_parity.py | 56 +- .../test/envs/test_interactive_scene_warp.py | 94 ---- .../envs/test_manager_based_rl_env_warp.py | 111 ++-- .../test/managers/test_manager_base.py | 19 + .../test/utils/test_manager_call_switch.py | 507 ------------------ .../test/utils/test_warp_graph_cache.py | 101 +++- .../manager_based/classic/ant/ant_env_cfg.py | 2 +- .../classic/humanoid/humanoid_env_cfg.py | 2 +- .../velocity/config/anymal_b/rough_env_cfg.py | 1 - .../velocity/config/anymal_c/rough_env_cfg.py | 1 - .../velocity/config/anymal_d/rough_env_cfg.py | 1 - .../velocity/config/go1/rough_env_cfg.py | 1 - .../locomotion/velocity/velocity_env_cfg.py | 8 +- .../manipulation/reach/reach_env_cfg.py | 2 +- 30 files changed, 506 insertions(+), 1558 deletions(-) delete mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py delete mode 100644 source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py delete mode 100644 source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py delete mode 100644 source/isaaclab_experimental/test/utils/test_manager_call_switch.py diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index 15b693197a46..9a7ab623bbfa 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -190,6 +190,10 @@ def __init__(self, cfg: InteractiveSceneCfg): self._aggregate_scene_data_requirements(requested_viz_types) + # The clone plan and terrain importer own the Torch origins buffer. Cache + # one zero-copy Warp view so frontend code does not rebuild it per access. + self._env_origins_wp = wp.from_torch(self.env_origins, dtype=wp.vec3f) + # Collision filtering is PhysX-only (matches both physx and ovphysx). if self.cfg.filter_collisions and "physx" in self.physics_backend and self._is_scene_setup_from_cfg(): self.filter_collisions(self._global_prim_paths) @@ -371,6 +375,11 @@ def env_origins(self) -> torch.Tensor: return self._terrain.env_origins return self.sim.get_clone_plan().positions + @property + def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): + """Cached zero-copy Warp view of environment origins [m], shape ``(num_envs,)``.""" + return self._env_origins_wp + @property def terrain(self) -> TerrainImporter | None: """The terrain in the scene. If None, then the scene has no terrain. @@ -450,27 +459,39 @@ def state(self) -> dict[str, dict[str, dict[str, torch.Tensor]]]: Operations. """ - def reset(self, env_ids: Sequence[int] | None = None): + def reset( + self, + env_ids: Sequence[int] | None = None, + env_mask: wp.array(dtype=wp.bool) | None = None, + ): """Resets the scene entities. Args: env_ids: The indices of the environments to reset. Defaults to None (all instances). + env_mask: Boolean Warp mask selecting environments. When provided, + it takes precedence over :paramref:`env_ids`. + """ + reset_kwargs = {"env_mask": env_mask} if env_mask is not None else {"env_ids": env_ids} # -- assets for articulation in self._articulations.values(): - articulation.reset(env_ids) + articulation.reset(**reset_kwargs) for deformable_object in self._deformable_objects.values(): - deformable_object.reset(env_ids) + deformable_object.reset(**reset_kwargs) for rigid_object in self._rigid_objects.values(): - rigid_object.reset(env_ids) + rigid_object.reset(**reset_kwargs) + if env_mask is not None and self._surface_grippers: + # Surface grippers expose only the legacy ID API. Materialize IDs at + # this optional PhysX boundary without penalizing Warp-native scenes. + env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) for surface_gripper in self._surface_grippers.values(): surface_gripper.reset(env_ids) for rigid_object_collection in self._rigid_object_collections.values(): - rigid_object_collection.reset(env_ids) + rigid_object_collection.reset(**reset_kwargs) # -- sensors for sensor in self._sensors.values(): - sensor.reset(env_ids) + sensor.reset(**reset_kwargs) def write_data_to_sim(self): """Writes the data of the scene entities to the simulation.""" diff --git a/source/isaaclab/isaaclab/terrains/terrain_importer.py b/source/isaaclab/isaaclab/terrains/terrain_importer.py index d9523497016f..570ecede58e0 100644 --- a/source/isaaclab/isaaclab/terrains/terrain_importer.py +++ b/source/isaaclab/isaaclab/terrains/terrain_importer.py @@ -345,7 +345,7 @@ def configure_env_origins(self, origins: np.ndarray | torch.Tensor | None = None if isinstance(origins, np.ndarray): origins = torch.from_numpy(origins) # store the origins - self.terrain_origins = origins.to(self.device, dtype=torch.float).contiguous() + self.terrain_origins = origins.to(self.device, dtype=torch.float) # compute environment origins self.env_origins = self._compute_env_origins_curriculum(self.cfg.num_envs, self.terrain_origins) else: diff --git a/source/isaaclab/test/scene/test_interactive_scene.py b/source/isaaclab/test/scene/test_interactive_scene.py index 41aaaf7f9c3d..909077bdbb39 100644 --- a/source/isaaclab/test/scene/test_interactive_scene.py +++ b/source/isaaclab/test/scene/test_interactive_scene.py @@ -13,9 +13,11 @@ """Rest everything follows.""" from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch +import warp as wp import isaaclab.sim as sim_utils from isaaclab.actuators import ImplicitActuatorCfg @@ -132,6 +134,59 @@ def test_reset_to_env_ids_input_types(device, setup_scene): assert_state_equal(prev_state, scene.get_state()) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_env_origins_wp_is_cached_zero_copy_view(device, setup_scene): + """The scene should expose one persistent Warp view of its Torch origins.""" + make_scene, sim = setup_scene + scene = InteractiveScene(make_scene(num_envs=4)) + sim.reset() + + assert scene.env_origins_wp is scene.env_origins_wp + assert scene.env_origins_wp.dtype == wp.vec3f + assert scene.env_origins_wp.shape == (scene.num_envs,) + assert scene.env_origins_wp.ptr == scene.env_origins.data_ptr() + + +def test_reset_dispatches_warp_mask_to_supported_entities() -> None: + """Mask reset should remain mask-based across assets and sensors.""" + scene = object.__new__(InteractiveScene) + articulation = Mock() + deformable = Mock() + rigid_object = Mock() + rigid_collection = Mock() + sensor = Mock() + scene._articulations = {"articulation": articulation} + scene._deformable_objects = {"deformable": deformable} + scene._rigid_objects = {"rigid_object": rigid_object} + scene._rigid_object_collections = {"collection": rigid_collection} + scene._sensors = {"sensor": sensor} + scene._surface_grippers = {} + env_mask = wp.array([True, False, True], dtype=wp.bool, device="cpu") + + scene.reset(env_mask=env_mask) + + for entity in (articulation, deformable, rigid_object, rigid_collection, sensor): + entity.reset.assert_called_once_with(env_mask=env_mask) + + +def test_mask_reset_compacts_ids_for_surface_grippers() -> None: + """Mask reset should materialize IDs only for the legacy gripper boundary.""" + scene = object.__new__(InteractiveScene) + gripper = Mock() + scene._articulations = {} + scene._deformable_objects = {} + scene._rigid_objects = {} + scene._rigid_object_collections = {} + scene._sensors = {} + scene._surface_grippers = {"gripper": gripper} + env_mask = wp.array([True, False, True], dtype=wp.bool, device="cpu") + + scene.reset(env_mask=env_mask) + + env_ids = gripper.reset.call_args.args[0] + assert torch.equal(env_ids, torch.tensor([0, 2], dtype=torch.int64)) + + def test_scene_publishes_plan_via_replicate(monkeypatch: pytest.MonkeyPatch): """A cfg-driven scene forwards the right plan and stage to cloner.replicate. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst index 97452d26225d..bfc3976ac1d3 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -12,16 +12,10 @@ Added Changed ^^^^^^^ -* Changed Warp manager stages to execute eagerly by default while stateful - capture semantics are validated. Set - ``MANAGER_CALL_CONFIG='{"default": 2}'`` to opt eligible manager stages into - CUDA graph capture. Direct-environment stages remained eager. - -Deprecated -^^^^^^^^^^ - -* Deprecated selecting stable managers inside ``ManagerBasedEnvWarp``. Use - ``ManagerBasedEnv`` for Torch managers or mode ``1`` for the Warp frontend. +* **Breaking:** Removed the experimental ``ManagerCallSwitch`` and + ``MANAGER_CALL_CONFIG`` interface. Warp environments now instantiate Warp + managers directly and automatically keep managers with non-capturable terms + on eager execution. Fixed ^^^^^ diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/envs/__init__.pyi index ef3d01d6b4ec..109fa5f8d06c 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/__init__.pyi +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/__init__.pyi @@ -6,13 +6,11 @@ __all__ = [ "mdp", "DirectRLEnvWarp", - "InteractiveSceneWarp", "ManagerBasedEnvWarp", "ManagerBasedRLEnvWarp", ] from . import mdp from .direct_rl_env_warp import DirectRLEnvWarp -from .interactive_scene_warp import InteractiveSceneWarp from .manager_based_env_warp import ManagerBasedEnvWarp from .manager_based_rl_env_warp import ManagerBasedRLEnvWarp diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py index f6f327b9ada2..fde6fa6f0ed2 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py @@ -31,13 +31,13 @@ # from isaaclab.envs.ui import ViewportCameraController from isaaclab.managers import EventManager +from isaaclab.scene import InteractiveScene from isaaclab.sim import SimulationContext from isaaclab.sim.utils import use_stage from isaaclab.utils.noise import NoiseModel from isaaclab.utils.seed import configure_seed from isaaclab.utils.timer import Timer -from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp from isaaclab_experimental.utils.warp import increment_all_int32, zero_masked_int32 from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache @@ -144,7 +144,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs with Timer("[INFO]: Time taken for scene creation", "scene_creation"): # set the stage context for scene creation steps which use the stage with use_stage(self.sim.stage): - self.scene = InteractiveSceneWarp(self.cfg.scene) + self.scene = InteractiveScene(self.cfg.scene) self._setup_scene() self.scene.initialize_renderers() # attach_stage_to_usd_context() @@ -225,10 +225,9 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.torch_reset_time_outs: torch.Tensor = None self.torch_episode_length_buf: torch.Tensor = None - # Direct stages stay eager until their complete backend boundaries are - # verified capture-safe. The owner-held executor keeps that policy out of - # individual stage call sites. - self._warp_graph_cache = WarpGraphCache(enabled=False) + # Direct-task stages include Python reset hooks and stay eager until those + # complete boundaries are validated for capture. + self._warp_graph_cache = WarpGraphCache(enabled=False, device=self.device) # setup the action and observation spaces for Gym self._configure_gym_env_spaces() diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py deleted file mode 100644 index 6905988499e3..000000000000 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py +++ /dev/null @@ -1,90 +0,0 @@ -# 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 - -"""Warp-native interactive scene with env_mask support for reset.""" - -from __future__ import annotations - -from collections.abc import Sequence - -import warp as wp - -from isaaclab.scene import InteractiveScene, InteractiveSceneCfg - - -class InteractiveSceneWarp(InteractiveScene): - """Interactive scene with warp-native env_mask support for reset. - - Extends :class:`InteractiveScene` to accept a boolean warp mask for selective resets, - avoiding the need to convert between env_ids and masks. - """ - - def __init__(self, cfg: InteractiveSceneCfg): - """Initialize the Warp scene and cache non-terrain environment origins. - - Args: - cfg: Configuration for the interactive scene. - """ - super().__init__(cfg) - if self.terrain is None: - self._env_origins_wp = wp.from_torch(self.env_origins, dtype=wp.vec3f) - - @property - def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): - """Cached zero-copy Warp view of environment origins [m], shape ``(num_envs,)``.""" - if self.terrain is not None: - return self.terrain.env_origins_wp - return self._env_origins_wp - - def reset( - self, - env_ids: Sequence[int] | None = None, - env_mask: wp.array(dtype=wp.bool) | None = None, - ) -> None: - """Reset scene entities using environment IDs or a boolean mask. - - Mask-based calls stay on the Warp data path for every entity that supports - masks. Surface grippers remain an explicit ID-based host boundary and are - reset only when :paramref:`env_ids` is also provided. Calls that only pass - IDs preserve the behavior of :meth:`InteractiveScene.reset`. - - Args: - env_ids: The indices of the environments to reset. Defaults to None (all instances). - env_mask: Boolean warp mask of shape (num_envs,). Defaults to None. - """ - if env_mask is None: - super().reset(env_ids) - return - if env_mask.dtype != wp.bool or env_mask.ndim != 1: - raise TypeError(f"env_mask must be a one-dimensional Warp boolean array, got {env_mask}.") - if env_mask.shape[0] != self.cfg.num_envs: - raise ValueError(f"env_mask must have shape ({self.cfg.num_envs},), received {env_mask.shape}.") - if env_mask.device != wp.get_device(self.sim.device): - raise ValueError(f"env_mask must be on device {self.sim.device}, received {env_mask.device}.") - - # -- Warp-capable assets - for articulation in self._articulations.values(): - articulation.reset(env_mask=env_mask) - for deformable_object in self._deformable_objects.values(): - deformable_object.reset(env_mask=env_mask) - for rigid_object in self._rigid_objects.values(): - rigid_object.reset(env_mask=env_mask) - for rigid_object_collection in self._rigid_object_collections.values(): - rigid_object_collection.reset(env_mask=env_mask) - # -- Warp-capable sensors - for sensor in self._sensors.values(): - sensor.reset(env_mask=env_mask) - - if env_ids is not None: - self.reset_host(env_ids) - - def reset_host(self, env_ids: Sequence[int] | None = None) -> None: - """Reset ID-based scene entities at an explicit host boundary. - - Args: - env_ids: The indices of the environments to reset. Defaults to all instances. - """ - for surface_gripper in self._surface_grippers.values(): - surface_gripper.reset(env_ids) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py index f1f390e8f163..df27295ea933 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py @@ -28,6 +28,8 @@ from isaaclab.envs.manager_based_env_cfg import ManagerBasedEnvCfg from isaaclab.envs.ui import ViewportCameraController from isaaclab.envs.utils.io_descriptors import export_articulations_data, export_scene_data +from isaaclab.managers import RecorderManager +from isaaclab.scene import InteractiveScene from isaaclab.sim import SimulationContext from isaaclab.sim.utils import use_stage from isaaclab.ui.widgets import ManagerLiveVisualizer @@ -35,8 +37,8 @@ from isaaclab.utils.timer import Timer from isaaclab.utils.version import has_kit -from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp as InteractiveScene -from isaaclab_experimental.utils.manager_call_switch import ManagerCallSwitch +from isaaclab_experimental.managers import ActionManager, EventManager, ObservationManager +from isaaclab_experimental.utils.torch_utils import clone_obs_buffer from isaaclab_experimental.utils.warp import resolve_1d_mask from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache @@ -78,18 +80,8 @@ def __init__(self, cfg: ManagerBasedEnvCfg): self.cfg = cfg # initialize internal variables self._is_closed = False - # Keep the execution policy on the environment. ManagerCallSwitch is a - # compatibility router that can be removed without replacing this owner. - self._warp_graph_cache = WarpGraphCache(enabled=True) - # Temporary runtime config for stable/Warp manager routing. - cfg_source: dict[str, int] | str | None = getattr(self.cfg, "manager_call_config", None) - max_modes: dict[str, int] | None = getattr(self.cfg, "manager_call_max_mode", None) - self._manager_call_switch = ManagerCallSwitch( - cfg_source, - max_modes=max_modes, - graph_cache=self._warp_graph_cache, - ) - self._manager_call_switch.apply_term_cfg_profile(self.cfg) + # Manager stages register their capture safety with this environment-owned cache. + self._warp_graph_cache = WarpGraphCache(device=self.cfg.sim.device) # set the seed for the environment if self.cfg.seed is not None: @@ -160,6 +152,9 @@ def __init__(self, cfg: ManagerBasedEnvCfg): # Pre-allocated env masks (shared across managers/terms via `env`). self.ALL_ENV_MASK = wp.ones((self.num_envs,), dtype=wp.bool, device=self.device) self.ENV_MASK = wp.zeros((self.num_envs,), dtype=wp.bool, device=self.device) + # Reset graphs always read this owner-held pointer, while callers may + # provide a different mask object on each reset. + self.reset_mask_wp = wp.zeros((self.num_envs,), dtype=wp.bool, device=self.device) # Persistent scalar buffer for global env step count (stable pointer for capture). self._global_env_step_count_wp = wp.zeros((1,), dtype=wp.int32, device=self.device) @@ -178,7 +173,7 @@ def __init__(self, cfg: ManagerBasedEnvCfg): # create event manager # note: this is needed here (rather than after simulation play) to allow USD-related randomization events # that must happen before the simulation starts. Example: randomizing mesh scale - self.event_manager = self._manager_call_switch.resolve_manager_class("EventManager")(self.cfg.events, self) + self.event_manager = EventManager(self.cfg.events, self) # apply USD-related randomization events if "prestartup" in self.event_manager.available_modes: @@ -360,22 +355,23 @@ def load_managers(self): :meth:`SimulationContext.reset_async` and it isn't possible to call async functions in the constructor. """ + # Reloading managers replaces pointer-stable buffers captured by prior stages. + self._warp_graph_cache.invalidate() # prepare the managers # -- event manager (we print it here to make the logging consistent) print("[INFO] Event Manager: ", self.event_manager) # -- recorder manager - self.recorder_manager = self._manager_call_switch.resolve_manager_class("RecorderManager")( - self.cfg.recorders, self - ) + self.recorder_manager = RecorderManager(self.cfg.recorders, self) self._has_recorders = bool(self.recorder_manager.active_terms) print("[INFO] Recorder Manager: ", self.recorder_manager) # -- action manager - self.action_manager = self._manager_call_switch.resolve_manager_class("ActionManager")(self.cfg.actions, self) + self.action_manager = ActionManager(self.cfg.actions, self) + self._action_in_wp = wp.zeros( + (self.num_envs, self.action_manager.total_action_dim), dtype=wp.float32, device=self.device + ) print("[INFO] Action Manager: ", self.action_manager) # -- observation manager - self.observation_manager = self._manager_call_switch.resolve_manager_class("ObservationManager")( - self.cfg.observations, self - ) + self.observation_manager = ObservationManager(self.cfg.observations, self) print("[INFO] Observation Manager:", self.observation_manager) # perform events at the start of the simulation @@ -425,11 +421,11 @@ def reset( A tuple containing the observations and extras. """ reset_mask = self.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) - host_env_ids = self._resolve_reset_host_ids(env_ids=env_ids, env_mask=reset_mask) - if host_env_ids is not None: - if self._has_recorders: - self.recorder_manager.record_pre_reset(host_env_ids) - self._reset_host_pre(host_env_ids) + if self._has_recorders: + recorder_env_ids = env_ids + if recorder_env_ids is None or isinstance(recorder_env_ids, wp.array): + recorder_env_ids = wp.to_torch(reset_mask).nonzero(as_tuple=False).squeeze(-1) + self.recorder_manager.record_pre_reset(recorder_env_ids) # set the seed if seed is not None: @@ -445,8 +441,8 @@ def reset( ) self._reset_idx(env_mask=reset_mask) - if host_env_ids is not None: - self.extras["log"].update(self._reset_host_post(host_env_ids)) + if self._has_recorders: + self.extras["log"].update(self.recorder_manager.reset(recorder_env_ids)) # update articulation kinematics self.scene.write_data_to_sim() @@ -456,11 +452,17 @@ def reset( for _ in range(self.cfg.num_rerenders_on_reset): self.sim.render() - if self._has_recorders and host_env_ids is not None: - self.recorder_manager.record_post_reset(host_env_ids) + if self._has_recorders: + self.recorder_manager.record_post_reset(recorder_env_ids) # compute observations - self.obs_buf = self.observation_manager.compute(update_history=True) + self.obs_buf = self._warp_graph_cache.call( + "ObservationManager_compute_reset", + self.observation_manager.compute, + update_history=True, + return_cloned_output=False, + output=clone_obs_buffer, + ) # return observations return self.obs_buf, self.extras @@ -501,9 +503,9 @@ def reset_to( self.seed(seed) env_mask = self.resolve_env_mask(env_ids=env_ids) - self._reset_host_pre(env_ids) self._reset_idx(env_mask=env_mask) - self.extras["log"].update(self._reset_host_post(env_ids)) + if self._has_recorders: + self.extras["log"].update(self.recorder_manager.reset(env_ids)) # set the state self.scene.reset_to(state, env_ids, is_relative=is_relative) @@ -520,7 +522,13 @@ def reset_to( self.recorder_manager.record_post_reset(env_ids) # compute observations - self.obs_buf = self.observation_manager.compute(update_history=True) + self.obs_buf = self._warp_graph_cache.call( + "ObservationManager_compute_reset", + self.observation_manager.compute, + update_history=True, + return_cloned_output=False, + output=clone_obs_buffer, + ) # return observations return self.obs_buf, self.extras @@ -541,13 +549,13 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: A tuple containing the observations and extras. """ # process actions - action_device = action.to(self.device) - if action_device.dtype != torch.float32: - action_device = action_device.float() - if not action_device.is_contiguous(): - action_device = action_device.contiguous() - action_wp = wp.from_torch(action_device, dtype=wp.float32) - self.action_manager.process_action(action_wp) + action_device = action.to(device=self.device, dtype=torch.float32).contiguous() + wp.copy(self._action_in_wp, wp.from_torch(action_device, dtype=wp.float32)) + self._warp_graph_cache.call( + "ActionManager_process_action", + self.action_manager.process_action, + action=self._action_in_wp, + ) if self._has_recorders: self.recorder_manager.record_pre_step() @@ -560,7 +568,7 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: for _ in range(self.cfg.decimation): self._sim_step_counter += 1 # set actions into buffers - self.action_manager.apply_action() + self._warp_graph_cache.call("ActionManager_apply_action", self.action_manager.apply_action) # set actions into simulator self.scene.write_data_to_sim() # simulate @@ -575,10 +583,21 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: # post-step: step interval event if "interval" in self.event_manager.available_modes: - self.event_manager.apply(mode="interval", dt=self.step_dt) + self._warp_graph_cache.call( + "EventManager_apply_interval", + self.event_manager.apply, + mode="interval", + dt=self.step_dt, + ) # -- compute observations - self.obs_buf = self.observation_manager.compute(update_history=True) + self.obs_buf = self._warp_graph_cache.call( + "ObservationManager_compute_update_history", + self.observation_manager.compute, + update_history=True, + return_cloned_output=False, + output=clone_obs_buffer, + ) if self._has_recorders: self.recorder_manager.record_post_step() @@ -639,6 +658,10 @@ def _reset_idx( Args: env_mask: Boolean Warp mask selecting environments to reset. """ + if env_mask is not self.reset_mask_wp: + wp.copy(self.reset_mask_wp, env_mask) + env_mask = self.reset_mask_wp + # reset the internal buffers of the scene elements self.scene.reset(env_mask=env_mask) @@ -646,8 +669,12 @@ def _reset_idx( if "reset" in self.event_manager.available_modes: env_step_count = self._sim_step_counter // self.cfg.decimation self._global_env_step_count_wp.fill_(env_step_count) - self.event_manager.apply( - mode="reset", env_mask_wp=env_mask, global_env_step_count=self._global_env_step_count_wp + self._warp_graph_cache.call( + "EventManager_apply_reset", + self.event_manager.apply, + mode="reset", + env_mask_wp=env_mask, + global_env_step_count=self._global_env_step_count_wp, ) # iterate over all managers and reset them @@ -655,47 +682,13 @@ def _reset_idx( # note: This is order-sensitive! Certain things need be reset before others. self.extras["log"] = dict() # -- observation manager - info = self.observation_manager.reset(env_mask=env_mask) + info = self._warp_graph_cache.call( + "ObservationManager_reset", self.observation_manager.reset, env_mask=env_mask + ) self.extras["log"].update(info) # -- action manager - info = self.action_manager.reset(env_mask=env_mask) + info = self._warp_graph_cache.call("ActionManager_reset", self.action_manager.reset, env_mask=env_mask) self.extras["log"].update(info) # -- event manager - info = self.event_manager.reset(env_mask=env_mask) + info = self._warp_graph_cache.call("EventManager_reset", self.event_manager.reset, env_mask=env_mask) self.extras["log"].update(info) - - def _reset_host_pre(self, env_ids: Sequence[int] | torch.Tensor) -> None: - """Reset ID-only scene state before the Warp reset stage.""" - self.scene.reset_host(env_ids) - - def _reset_host_post(self, env_ids: Sequence[int] | torch.Tensor) -> dict[str, Any]: - """Reset host-only managers after the Warp reset stage.""" - if self._has_recorders: - return self.recorder_manager.reset(env_ids) - return {} - - def _reset_requires_host_selection(self) -> bool: - """Return whether configured reset features require a host-visible selection.""" - return bool(self._has_recorders or self.scene.surface_grippers) - - def _resolve_reset_host_ids( - self, - *, - env_ids: Sequence[int] | wp.array | torch.Tensor | None, - env_mask: wp.array(dtype=wp.bool), - ) -> Sequence[int] | torch.Tensor | None: - """Materialize reset IDs only when a configured host feature consumes them. - - Args: - env_ids: Original environment ID selection, if provided. - env_mask: Canonical Warp reset mask. - - Returns: - Host-readable environment IDs, or ``None`` when the reset remains - entirely on the Warp data path. - """ - if not self._reset_requires_host_selection(): - return None - if env_ids is not None and not isinstance(env_ids, wp.array): - return env_ids - return wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index 696bf001b42b..04567703e618 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -27,7 +27,7 @@ from isaaclab.ui.widgets import ManagerLiveVisualizer from isaaclab.utils.timer import Timer -from isaaclab_experimental.managers import CommandManager, CurriculumManager +from isaaclab_experimental.managers import CommandManager, CurriculumManager, RewardManager, TerminationManager from isaaclab_experimental.utils.torch_utils import clone_obs_buffer from isaaclab_experimental.utils.warp import increment_all_int64, zero_masked_int64 @@ -105,16 +105,6 @@ def __init__(self, cfg: ManagerBasedRLEnvCfg, render_mode: str | None = None, ** # store the render mode self.render_mode = render_mode - # Reset stages may be captured independently, so every stage must see - # one owner-held mask pointer. Copy each new selection into this buffer - # before dispatch instead of recording caller-owned mask addresses. - self.reset_mask_wp = wp.zeros(cfg.scene.num_envs, dtype=wp.bool, device=cfg.sim.device) - - # Persistent action input buffer to keep pointer stable for captured graphs. - self._action_in_wp: wp.array = wp.zeros( - (self.num_envs, self.action_manager.total_action_dim), dtype=wp.float32, device=self.device - ) - # initialize data and constants # -- set the framerate of the gym video recorder wrapper so that the playback speed # of the produced video matches the simulation @@ -163,12 +153,10 @@ def load_managers(self): # prepare the managers # -- termination manager - self.termination_manager = self._manager_call_switch.resolve_manager_class("TerminationManager")( - self.cfg.terminations, self - ) + self.termination_manager = TerminationManager(self.cfg.terminations, self) print("[INFO] Termination Manager: ", self.termination_manager) - # -- reward manager (experimental fork; Warp-compatible rewards) - self.reward_manager = self._manager_call_switch.resolve_manager_class("RewardManager")(self.cfg.rewards, self) + # -- reward manager + self.reward_manager = RewardManager(self.cfg.rewards, self) print("[INFO] Reward Manager: ", self.reward_manager) # -- Warp-first curriculum manager self.curriculum_manager = CurriculumManager(self.cfg.curriculum, self) @@ -235,16 +223,14 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: # IMPORTANT: Do NOT re-wrap/replace the `wp.array` used by captured graphs each step. # Instead, copy the latest actions into the persistent buffer. with Timer(name="action_preprocess", msg="Action preprocessing took:", enable=DEBUG_TIMER_STEP, time_unit="us"): - if self._action_in_wp is None: - raise RuntimeError("Action buffer not initialized. Call reset() before step().") - action_device = action.to(self.device) + action_device = action.to(device=self.device, dtype=torch.float32).contiguous() wp.copy(self._action_in_wp, wp.from_torch(action_device, dtype=wp.float32)) - self._manager_call_switch.call( + self._warp_graph_cache.call( "ActionManager_process_action", self.action_manager.process_action, action=self._action_in_wp, - _timer=DEBUG_TIMER_STEP, + timer=DEBUG_TIMER_STEP, ) if self._has_recorders: @@ -258,16 +244,18 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: for _ in range(self.cfg.decimation): self._sim_step_counter += 1 # set actions into buffers - self._manager_call_switch.call( + self._warp_graph_cache.call( "ActionManager_apply_action", self.action_manager.apply_action, - _timer=DEBUG_TIMER_STEP, - ) - self._manager_call_switch.call( - "Scene_write_data_to_sim", - self.scene.write_data_to_sim, - _timer=DEBUG_TIMER_STEP, + timer=DEBUG_TIMER_STEP, ) + with Timer( + name="Scene_write_data_to_sim", + msg="Scene write took:", + enable=DEBUG_TIMER_STEP, + time_unit="us", + ): + self.scene.write_data_to_sim() # simulate with Timer(name="simulate", msg="Newton simulation step took:", enable=DEBUG_TIMER_STEP, time_unit="us"): @@ -299,57 +287,57 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self.common_step_counter += 1 # total step (common for all envs) # -- post-processing (termination + reward) as independently configurable stages - self._manager_call_switch.call( + self._warp_graph_cache.call( "TerminationManager_compute", self.step_warp_termination_compute, - _timer=DEBUG_TIMER_STEP, + timer=DEBUG_TIMER_STEP, ) - self.reward_buf = self._manager_call_switch.call( + self.reward_buf = self._warp_graph_cache.call( "RewardManager_compute", self.reward_manager.compute, dt=float(self.step_dt), - _timer=DEBUG_TIMER_STEP, + timer=DEBUG_TIMER_STEP, ) if self._has_recorders: # update observations for recording if needed - self._manager_call_switch.call( + self._warp_graph_cache.call( "ObservationManager_compute_no_history", self.observation_manager.compute, return_cloned_output=False, - _timer=DEBUG_TIMER_STEP, + timer=DEBUG_TIMER_STEP, ) self.recorder_manager.record_post_step() self._reset_terminated_envs() # -- update command - self._manager_call_switch.call( + self._warp_graph_cache.call( "CommandManager_compute", self.command_manager.compute, dt=float(self.step_dt), - _timer=DEBUG_TIMER_STEP, + timer=DEBUG_TIMER_STEP, ) # -- step interval events if "interval" in self.event_manager.available_modes: - self._manager_call_switch.call( + self._warp_graph_cache.call( "EventManager_apply_interval", self.event_manager.apply, mode="interval", dt=float(self.step_dt), - _timer=DEBUG_TIMER_STEP, + timer=DEBUG_TIMER_STEP, ) # -- compute observations # note: done after reset to get the correct observations for reset envs - self.obs_buf = self._manager_call_switch.call( + self.obs_buf = self._warp_graph_cache.call( "ObservationManager_compute_update_history", self.observation_manager.compute, update_history=True, return_cloned_output=False, - _output=clone_obs_buffer, - _timer=DEBUG_TIMER_STEP, + output=clone_obs_buffer, + timer=DEBUG_TIMER_STEP, ) # return observations, rewards, resets and extras return self.obs_buf, self.reward_buf, self.reset_terminated, self.reset_time_outs, self.extras @@ -475,78 +463,74 @@ def _reset_idx( wp.copy(self.reset_mask_wp, env_mask) env_mask = self.reset_mask_wp - self._manager_call_switch.call( + self._warp_graph_cache.call( "CurriculumManager_compute", self.curriculum_manager.compute, env_mask=env_mask, - _timer=DEBUG_TIMER_RESET, + timer=DEBUG_TIMER_RESET, ) - self._manager_call_switch.call( - "Scene_reset", - self.scene.reset, - env_mask=env_mask, - _timer=DEBUG_TIMER_RESET, - ) + with Timer(name="Scene_reset", msg="Scene reset took:", enable=DEBUG_TIMER_RESET, time_unit="us"): + self.scene.reset(env_mask=env_mask) if "reset" in self.event_manager.available_modes: self._global_env_step_count_wp.fill_(self._sim_step_counter // self.cfg.decimation) - self._manager_call_switch.call( + self._warp_graph_cache.call( "EventManager_apply_reset", self.event_manager.apply, mode="reset", env_mask_wp=env_mask, global_env_step_count=self._global_env_step_count_wp, - _timer=DEBUG_TIMER_RESET, + timer=DEBUG_TIMER_RESET, ) # iterate over all managers and reset them # this returns a dictionary of information which is stored in the extras # note: This is order-sensitive! Certain things need be reset before others. # -- observation manager + action + reward managers - obs_info = self._manager_call_switch.call( + obs_info = self._warp_graph_cache.call( "ObservationManager_reset", self.observation_manager.reset, env_mask=env_mask, - _timer=DEBUG_TIMER_RESET, + timer=DEBUG_TIMER_RESET, ) - action_info = self._manager_call_switch.call( + action_info = self._warp_graph_cache.call( "ActionManager_reset", self.action_manager.reset, env_mask=env_mask, - _timer=DEBUG_TIMER_RESET, + timer=DEBUG_TIMER_RESET, ) - reward_info = self._manager_call_switch.call( + reward_info = self._warp_graph_cache.call( "RewardManager_reset", self.reward_manager.reset, env_mask=env_mask, - _timer=DEBUG_TIMER_RESET, + timer=DEBUG_TIMER_RESET, ) - curriculum_info = self._manager_call_switch.call( + curriculum_info = self._warp_graph_cache.call( "CurriculumManager_reset", self.curriculum_manager.reset, env_mask=env_mask, - _timer=DEBUG_TIMER_RESET, + timer=DEBUG_TIMER_RESET, ) # -- command + event + termination managers - command_info = self._manager_call_switch.call( + command_info = self._warp_graph_cache.call( "CommandManager_reset", self.command_manager.reset, env_mask=env_mask, - _timer=DEBUG_TIMER_RESET, + timer=DEBUG_TIMER_RESET, ) - event_info = self._manager_call_switch.call( + event_info = self._warp_graph_cache.call( "EventManager_reset", self.event_manager.reset, env_mask=env_mask, - _timer=DEBUG_TIMER_RESET, + timer=DEBUG_TIMER_RESET, ) - termination_info = self._manager_call_switch.call( + termination_info = self._warp_graph_cache.call( "TerminationManager_reset", self.termination_manager.reset, env_mask=env_mask, - _timer=DEBUG_TIMER_RESET, + timer=DEBUG_TIMER_RESET, ) # reset the episode length buffer @@ -572,28 +556,27 @@ def _reset_idx( self.extras["log"] = log def _reset_terminated_envs(self) -> None: - """Reset terminated environments, compacting IDs only for host consumers.""" + """Reset terminated environments using the canonical Warp mask.""" reset_mask = self.termination_manager.dones_wp - # The eager reset pipeline contains many small launches. Keep the mask as - # the canonical selection, but use one explicit host predicate to avoid - # dispatching the entire pipeline when it is empty. Recorded launches - # can reduce that cost; capture can later make this predicate unnecessary. + # Keep the mask as the canonical selection, but use one host predicate + # to avoid dispatching the complete reset pipeline when it is empty. if not self.reset_buf.any().item(): return - reset_env_ids = None - if self._reset_requires_host_selection(): + + # Same-step autoreset exposes terminal observations before any selected + # environment is reset, matching the stable manager-based environment. + if self.cfg.compute_final_obs: + self.extras["final_obs"] = self.observation_manager.compute() + + if self._has_recorders: with Timer( name="reset_selection_host", - msg="Host reset selection took:", + msg="Recorder reset selection took:", enable=DEBUG_TIMER_STEP, time_unit="us", ): - reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) - if reset_env_ids.numel() == 0: - return - if self._has_recorders: - self.recorder_manager.record_pre_reset(reset_env_ids) - self._reset_host_pre(reset_env_ids) + recorder_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) + self.recorder_manager.record_pre_reset(recorder_env_ids) with Timer( name="reset_idx", @@ -603,10 +586,9 @@ def _reset_terminated_envs(self) -> None: ): self._reset_idx(env_mask=reset_mask) - if reset_env_ids is not None and reset_env_ids.numel() > 0: - self.extras["log"].update(self._reset_host_post(reset_env_ids)) - if self._has_recorders: - self.recorder_manager.record_post_reset(reset_env_ids) + if self._has_recorders: + self.extras["log"].update(self.recorder_manager.reset(recorder_env_ids)) + self.recorder_manager.record_post_reset(recorder_env_ids) if self.has_rtx_sensors and self.cfg.num_rerenders_on_reset > 0: for _ in range(self.cfg.num_rerenders_on_reset): self.sim.render() diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py index edff91465a9b..7f586633b764 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py @@ -27,7 +27,6 @@ from __future__ import annotations -import warnings from typing import TYPE_CHECKING import warp as wp @@ -36,10 +35,6 @@ from isaaclab_experimental.utils.warp import WarpCapturable, warp_capturable __all__ = [ - "ApplyExternalForceTorque", - "PushBySettingVelocity", - "RandomizeRigidBodyCom", - "ResetRootStateUniform", "apply_external_force_torque", "push_by_setting_velocity", "randomize_rigid_body_com", @@ -94,7 +89,7 @@ def _randomize_com_kernel( @warp_capturable(False) -class RandomizeRigidBodyCom(ManagerTermBase): +class randomize_rigid_body_com(ManagerTermBase): """Randomize rigid-body centers of mass from a persistent default baseline. This term is not CUDA-graph capturable because notifying the solver of changed @@ -154,26 +149,6 @@ def __call__( self._asset.set_coms_mask(coms=self._asset.data.body_com_pos_b.warp, env_mask=env_mask) -@WarpCapturable(False, reason="set_coms_mask calls SimulationManager.add_model_change") -def randomize_rigid_body_com( - env: ManagerBasedEnv, - env_mask: wp.array(dtype=wp.bool), - com_range: dict[str, tuple[float, float]], - asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -) -> None: - """Deprecated compatibility adapter for :class:`RandomizeRigidBodyCom`.""" - warnings.warn( - "'randomize_rigid_body_com' is deprecated; use 'RandomizeRigidBodyCom' in event configurations.", - DeprecationWarning, - stacklevel=2, - ) - params = {"com_range": com_range, "asset_cfg": asset_cfg} - cfg = EventTermCfg(func=RandomizeRigidBodyCom, mode="reset", params={}) - cfg.params = params - term = RandomizeRigidBodyCom(cfg, env) - term(env, env_mask, **params) - - # --------------------------------------------------------------------------- # Apply external force and torque # --------------------------------------------------------------------------- @@ -211,7 +186,7 @@ def _apply_external_force_torque_kernel( rng_state[env_id] = state -class ApplyExternalForceTorque(ManagerTermBase): +class apply_external_force_torque(ManagerTermBase): """Apply random external forces and torques using persistent wrench buffers.""" def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv) -> None: @@ -283,27 +258,6 @@ def __call__( ) -@WarpCapturable(False, reason="deprecated adapter constructs term state; use ApplyExternalForceTorque") -def apply_external_force_torque( - env: ManagerBasedEnv, - env_mask: wp.array(dtype=wp.bool), - force_range: tuple[float, float], - torque_range: tuple[float, float], - asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -) -> None: - """Deprecated compatibility adapter for :class:`ApplyExternalForceTorque`.""" - warnings.warn( - "'apply_external_force_torque' is deprecated; use 'ApplyExternalForceTorque' in event configurations.", - DeprecationWarning, - stacklevel=2, - ) - params = {"force_range": force_range, "torque_range": torque_range, "asset_cfg": asset_cfg} - cfg = EventTermCfg(func=ApplyExternalForceTorque, mode="reset", params={}) - cfg.params = params - term = ApplyExternalForceTorque(cfg, env) - term(env, env_mask, **params) - - # --------------------------------------------------------------------------- # Push by velocity # --------------------------------------------------------------------------- @@ -339,7 +293,7 @@ def _push_by_setting_velocity_kernel( rng_state[env_id] = state -class PushBySettingVelocity(ManagerTermBase): +class push_by_setting_velocity(ManagerTermBase): """Push an asset by sampling into a persistent root-velocity buffer.""" def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv) -> None: @@ -395,26 +349,6 @@ def __call__( self._asset.write_root_velocity_to_sim_mask(root_velocity=self._velocity, env_mask=env_mask) -@WarpCapturable(False, reason="deprecated adapter constructs term state; use PushBySettingVelocity") -def push_by_setting_velocity( - env: ManagerBasedEnv, - env_mask: wp.array(dtype=wp.bool), - velocity_range: dict[str, tuple[float, float]], - asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -) -> None: - """Deprecated compatibility adapter for :class:`PushBySettingVelocity`.""" - warnings.warn( - "'push_by_setting_velocity' is deprecated; use 'PushBySettingVelocity' in event configurations.", - DeprecationWarning, - stacklevel=2, - ) - params = {"velocity_range": velocity_range, "asset_cfg": asset_cfg} - cfg = EventTermCfg(func=PushBySettingVelocity, mode="interval", params={}) - cfg.params = params - term = PushBySettingVelocity(cfg, env) - term(env, env_mask, **params) - - # --------------------------------------------------------------------------- # Reset root state uniform # --------------------------------------------------------------------------- @@ -484,7 +418,7 @@ def _reset_root_state_uniform_kernel( rng_state[env_id] = state -class ResetRootStateUniform(ManagerTermBase): +class reset_root_state_uniform(ManagerTermBase): """Reset root pose and velocity using persistent Warp output buffers.""" def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv) -> None: @@ -561,27 +495,6 @@ def __call__( self._asset.write_root_velocity_to_sim_mask(root_velocity=self._velocity, env_mask=env_mask) -@WarpCapturable(False, reason="deprecated adapter constructs term state; use ResetRootStateUniform") -def reset_root_state_uniform( - env: ManagerBasedEnv, - env_mask: wp.array(dtype=wp.bool), - pose_range: dict[str, tuple[float, float]], - velocity_range: dict[str, tuple[float, float]], - asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -) -> None: - """Deprecated compatibility adapter for :class:`ResetRootStateUniform`.""" - warnings.warn( - "'reset_root_state_uniform' is deprecated; use 'ResetRootStateUniform' in event configurations.", - DeprecationWarning, - stacklevel=2, - ) - params = {"pose_range": pose_range, "velocity_range": velocity_range, "asset_cfg": asset_cfg} - cfg = EventTermCfg(func=ResetRootStateUniform, mode="reset", params={}) - cfg.params = params - term = ResetRootStateUniform(cfg, env) - term(env, env_mask, **params) - - # --------------------------------------------------------------------------- # Reset joints by offset # --------------------------------------------------------------------------- diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py index 0ccee0e6d381..1a2b654b76dc 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py @@ -473,6 +473,7 @@ def _prepare_terms(self): ) # create the action term term = term_cfg.class_type(term_cfg, self._env) + self._register_term_capturability(term_cfg.class_type) # sanity check if term is valid type if not isinstance(term, ActionTerm): raise TypeError(f"Returned object for the term '{term_name}' is not of type ActionType.") diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py index 496bc17b87fd..9ebb23ad2492 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py @@ -593,6 +593,7 @@ def _prepare_terms(self): ) # create the action term term = term_cfg.class_type(term_cfg, self._env) + self._register_term_capturability(term_cfg.class_type) # sanity check if term is valid type if not isinstance(term, CommandTerm): raise TypeError(f"Returned object for the term '{term_name}' is not of type CommandType.") diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py index 291d807ddadc..0524d6efd7c6 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py @@ -20,7 +20,7 @@ import inspect import logging from abc import ABC, abstractmethod -from collections.abc import Sequence +from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, Any import warp as wp @@ -446,18 +446,20 @@ def _resolve_common_term_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg, f" and optional parameters: {args_with_defaults}, but received: {term_params}." ) - # Register capture safety with the environment-owned executor. The - # manager call switch only consumes this policy while it still exists. - if not is_warp_capturable(term_cfg.func): - graph_cache = getattr(self._env, "_warp_graph_cache", None) - if graph_cache is not None: - graph_cache.register_capturability(type(self).__name__, False) + self._register_term_capturability(term_cfg.func) # process attributes at runtime # these properties are only resolvable once the simulation starts playing if self._env.sim.is_playing(): self._process_term_cfg_at_play(term_name, term_cfg) + def _register_term_capturability(self, term: Callable) -> None: + """Keep the complete manager eager when a configured term is unsafe.""" + if not is_warp_capturable(term): + graph_cache = getattr(self._env, "_warp_graph_cache", None) + if graph_cache is not None: + graph_cache.register_capturability(type(self).__name__, False) + def _process_term_cfg_at_play(self, term_name: str, term_cfg: ManagerTermBaseCfg): """Process the term configuration at runtime. diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi index de2860732514..00660ba9775e 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi @@ -4,8 +4,6 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ - "ManagerCallMode", - "ManagerCallSwitch", "WarpGraphCache", "clone_obs_buffer", "buffers", @@ -14,7 +12,6 @@ __all__ = [ "warp", ] -from .manager_call_switch import ManagerCallMode, ManagerCallSwitch from .torch_utils import clone_obs_buffer from .warp_graph_cache import WarpGraphCache from . import buffers, modifiers, noise, warp diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py b/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py deleted file mode 100644 index 4999646e402b..000000000000 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py +++ /dev/null @@ -1,372 +0,0 @@ -# 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 - -"""Manager call switch for routing manager stage calls through stable/warp/captured paths.""" - -from __future__ import annotations - -import importlib -import json -import logging -import os -import warnings -from collections.abc import Callable -from copy import deepcopy -from enum import IntEnum -from typing import TYPE_CHECKING, Any - -from isaaclab.utils.timer import Timer - -from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache - -if TYPE_CHECKING: - from isaaclab.envs import ManagerBasedEnvCfg - -logger = logging.getLogger(__name__) - - -class ManagerCallMode(IntEnum): - """Execution mode for manager stage calls. - - * ``STABLE`` (0): Call stable Python manager implementations from :mod:`isaaclab.managers`. - This mode is deprecated inside ``ManagerBasedEnvWarp``; use the stable environment instead. - * ``WARP_NOT_CAPTURED`` (1): Call Warp-compatible implementations without CUDA graph capture. - * ``WARP_CAPTURED`` (2): Call Warp implementations with CUDA graph capture/replay. - """ - - STABLE = 0 - WARP_NOT_CAPTURED = 1 - WARP_CAPTURED = 2 - - -class ManagerCallSwitch: - """Compatibility router for stable and Warp manager implementations. - - This temporary layer selects stable or Warp manager classes and forwards - Warp stages to an environment-owned :class:`WarpGraphCache`. Execution and - capture policy therefore remain reusable after mixed stable-manager routing - is removed. Calls may optionally use a :class:`Timer` for profiling. - """ - - # Warp eager is the correctness-first default. Capture remains an explicit - # optimization while stage state and pointer contracts are validated. - DEFAULT_CONFIG: dict[str, int] = {"default": 1} - DEFAULT_KEY = "default" - MANAGER_NAMES: tuple[str, ...] = ( - "ActionManager", - "ObservationManager", - "EventManager", - "RecorderManager", - "CommandManager", - "TerminationManager", - "RewardManager", - "CurriculumManager", - "Scene", - ) - # Scene stages remain eager until scene, sensor, and actuator graphability is - # validated together. Warp-first execution does not depend on that later step. - MAX_MODE_OVERRIDES: dict[str, int] = {"Scene": ManagerCallMode.WARP_NOT_CAPTURED} - - ENV_VAR = "MANAGER_CALL_CONFIG" - """Environment variable name for the JSON config string. - - Example usage:: - - MANAGER_CALL_CONFIG='{"RewardManager": 0, "default": 1}' python train.py ... - """ - - def __init__( - self, - cfg_source: dict[str, int] | str | None = None, - *, - max_modes: dict[str, int] | None = None, - graph_cache: WarpGraphCache | None = None, - ): - # The environment normally owns this durable executor. Creating one here - # keeps the compatibility router independently usable in tests and tools. - self._graph_cache = graph_cache if graph_cache is not None else WarpGraphCache(enabled=True) - # Merge caller-supplied max_modes with the class-level MAX_MODE_OVERRIDES. - self._max_modes = dict(self.MAX_MODE_OVERRIDES) - if max_modes is not None: - self._max_modes.update(max_modes) - # Resolve config: prefer explicit cfg_source, fall back to env var. - if cfg_source is None: - cfg_source = os.environ.get(self.ENV_VAR) - self._cfg = self._load_cfg(cfg_source) - print("[INFO] ManagerCallSwitch configuration:") - print(f" - {self.DEFAULT_KEY}: {self._cfg[self.DEFAULT_KEY]}") - for manager_name in self.MANAGER_NAMES: - mode = int(self.get_mode_for_manager(manager_name)) - cap = self._max_modes.get(manager_name) - cap_str = f" (cap={cap})" if cap is not None else "" - print(f" - {manager_name}: {mode}{cap_str}") - - # ------------------------------------------------------------------ - # Graph management - # ------------------------------------------------------------------ - - def invalidate_graphs(self) -> None: - """Invalidate cached capture graphs and their cached return values.""" - self._graph_cache.invalidate() - - def apply_term_cfg_profile(self, cfg: ManagerBasedEnvCfg) -> None: - """Apply the deprecated mixed stable-manager configuration profile. - - This compatibility hook is intentionally isolated from the Warp - environment. It can be removed together with stable manager routing - once all task configurations use the Warp frontend natively. - - Args: - cfg: Experimental environment configuration to update in place. - """ - manager_to_cfg_attr = { - "ActionManager": "actions", - "ObservationManager": "observations", - "EventManager": "events", - "RecorderManager": "recorders", - "CommandManager": "commands", - "TerminationManager": "terminations", - "RewardManager": "rewards", - "CurriculumManager": "curriculum", - } - stable_managers = [ - name for name in manager_to_cfg_attr if self.get_mode_for_manager(name) == ManagerCallMode.STABLE - ] - if not stable_managers: - return - - warnings.warn( - "Selecting STABLE managers inside ManagerBasedEnvWarp is deprecated. Use ManagerBasedEnv for Torch " - "managers or WARP_NOT_CAPTURED for the Warp frontend.", - DeprecationWarning, - stacklevel=2, - ) - stable_cfg = self._resolve_stable_cfg_counterpart(cfg) - if stable_cfg is None: - logger.warning( - "Stable managers requested (%s), but no stable cfg counterpart could be resolved. Keeping " - "experimental term configs.", - ", ".join(stable_managers), - ) - return - - for manager_name, cfg_attr in manager_to_cfg_attr.items(): - if self.get_mode_for_manager(manager_name) != ManagerCallMode.STABLE: - continue - if hasattr(cfg, cfg_attr) and hasattr(stable_cfg, cfg_attr): - setattr(cfg, cfg_attr, deepcopy(getattr(stable_cfg, cfg_attr))) - - # ------------------------------------------------------------------ - # Stage dispatch - # ------------------------------------------------------------------ - - def call( - self, - stage: str, - fn: Callable[..., Any], - /, - *args: Any, - _output: Callable[[Any], Any] | None = None, - _timer: bool = False, - **kwargs: Any, - ) -> Any: - """Run a Warp frontend stage eagerly or through its cached CUDA graph. - - The call site always supplies the same callable and arguments. The - configured manager mode only controls how that call is executed. - - Args: - stage: Stage identifier in the form ``"ManagerName_function_name"``. - fn: Callable implementing the stage. - *args: Positional arguments forwarded to :paramref:`fn`. - _output: Optional transform applied to the stage result after execution. - _timer: Whether to time the stage. - **kwargs: Keyword arguments forwarded to :paramref:`fn`. - - Returns: - The stage result, optionally transformed by :paramref:`_output`. - """ - with Timer(name=stage, msg=f"{stage} took:", enable=_timer, time_unit="us"): - manager_name = self._manager_name_from_stage(stage) - mode = self.get_mode_for_manager(manager_name) - result = self._graph_cache.call( - stage, - fn, - args=args, - kwargs=kwargs, - capture=mode == ManagerCallMode.WARP_CAPTURED, - group=manager_name, - ) - return _output(result) if _output is not None else result - - def call_stage( - self, - *, - stage: str, - warp_call: dict[str, Any], - stable_call: dict[str, Any] | None = None, - timer: bool = False, - ) -> Any: - """Run the stage according to configured mode, optionally wrapped in a :class:`Timer`. - - A call spec dict supports the following keys: - - * ``fn`` (required): The callable to invoke. - * ``args`` (optional): Positional arguments tuple. - * ``kwargs`` (optional): Keyword arguments dict. - * ``output`` (optional): A ``Callable[[Any], Any]`` that transforms the raw - return value into the final output. For captured stages the raw value is - ``None``. When omitted, the raw return value is used as-is. - - Args: - stage: Stage identifier in the form ``"ManagerName_function_name"``. - warp_call: Call spec for the warp path (eager or captured). - stable_call: Call spec for the stable (torch) path. Defaults to ``None``. - timer: Whether to wrap execution in a :class:`Timer`. Defaults to ``False`` - (controlled by the global :attr:`Timer.enable` class-level toggle). - Pass a module-level flag like ``TIMER_ENABLED_STEP`` to make timing - conditional on that flag. - - Returns: - The (possibly transformed) return value of the stage. - """ - mode = self.get_mode_for_manager(self._manager_name_from_stage(stage)) - if mode == ManagerCallMode.STABLE: - if stable_call is None: - raise ValueError(f"Stage '{stage}' is configured as STABLE (mode=0) but no stable_call was provided.") - call = stable_call - else: - call = warp_call - return self.call( - stage, - call["fn"], - *call.get("args", ()), - _output=call.get("output"), - _timer=timer, - **call.get("kwargs", {}), - ) - - # ------------------------------------------------------------------ - # Manager resolution - # ------------------------------------------------------------------ - - def _manager_name_from_stage(self, stage: str) -> str: - if "_" not in stage: - raise ValueError(f"Invalid stage '{stage}'. Expected '{{manager_name}}_{{function_name}}'.") - return stage.split("_", 1)[0] - - def get_mode_for_manager(self, manager_name: str) -> ManagerCallMode: - """Return the resolved execution mode for the given manager. - - Looks up the manager in the config dict, falls back to the default, - then caps by :attr:`_max_modes` (static overrides + dynamic registrations). - """ - default_key = next(iter(self.DEFAULT_CONFIG)) - mode_value = self._cfg.get(manager_name, self._cfg[default_key]) - cap = self._max_modes.get(manager_name) - if cap is not None: - mode_value = min(mode_value, cap) - if mode_value == ManagerCallMode.WARP_CAPTURED and not self._graph_cache.is_capturable(manager_name): - mode_value = ManagerCallMode.WARP_NOT_CAPTURED - return ManagerCallMode(mode_value) - - def resolve_manager_class(self, manager_name: str, mode_override: ManagerCallMode | int | None = None) -> type: - """Import and return the manager class for the configured mode.""" - mode = self.get_mode_for_manager(manager_name) if mode_override is None else ManagerCallMode(mode_override) - module_name = "isaaclab.managers" if mode == ManagerCallMode.STABLE else "isaaclab_experimental.managers" - module = importlib.import_module(module_name) - if not hasattr(module, manager_name): - raise AttributeError(f"Manager '{manager_name}' not found in module '{module_name}'.") - return getattr(module, manager_name) - - def register_manager_capturability(self, manager_name: str, capturable: bool) -> None: - """Register that a manager has non-capturable terms, capping its mode. - - Called by :class:`ManagerBase` during term preparation when a term - is decorated with ``@warp_capturable(False)``. - """ - self._graph_cache.register_capturability(manager_name, capturable) - - # ------------------------------------------------------------------ - # Internal helpers - # ------------------------------------------------------------------ - - def _load_cfg(self, cfg_source: dict[str, int] | str | None) -> dict[str, int]: - if cfg_source is None: - cfg = dict(self.DEFAULT_CONFIG) - elif isinstance(cfg_source, dict): - cfg = dict(cfg_source) - if self.DEFAULT_KEY not in cfg: - cfg[self.DEFAULT_KEY] = self.DEFAULT_CONFIG[self.DEFAULT_KEY] - elif isinstance(cfg_source, str): - if cfg_source.strip() == "": - cfg = dict(self.DEFAULT_CONFIG) - else: - parsed = json.loads(cfg_source) - if not isinstance(parsed, dict): - raise TypeError("manager_call_config must decode to a dict.") - cfg = dict(parsed) - if self.DEFAULT_KEY not in cfg: - cfg[self.DEFAULT_KEY] = self.DEFAULT_CONFIG[self.DEFAULT_KEY] - else: - raise TypeError(f"cfg_source must be a dict, string, or None, got: {type(cfg_source)}") - - # validation - for manager_name, mode_value in cfg.items(): - if not isinstance(mode_value, int): - raise TypeError( - f"manager_call_config value for '{manager_name}' must be int (0/1/2), got: {type(mode_value)}" - ) - try: - ManagerCallMode(mode_value) - except ValueError as exc: - raise ValueError( - f"Invalid manager_call_config value for '{manager_name}': {mode_value}. Expected 0/1/2." - ) from exc - - # Apply max mode caps: bake caps into the resolved config so - # get_mode_for_manager never needs per-call branching. - default_mode = cfg[self.DEFAULT_KEY] - for name, max_mode in self._max_modes.items(): - resolved = cfg.get(name, default_mode) - if resolved > max_mode: - cfg[name] = max_mode - - return cfg - - @staticmethod - def _resolve_stable_cfg_counterpart(cfg: ManagerBasedEnvCfg) -> ManagerBasedEnvCfg | None: - """Resolve the legacy stable task configuration counterpart.""" - cfg_cls = cfg.__class__ - cfg_module_name = cfg_cls.__module__ - if "isaaclab_tasks_experimental" not in cfg_module_name: - return None - - stable_module_name = cfg_module_name.replace("isaaclab_tasks_experimental", "isaaclab_tasks", 1) - try: - stable_module = importlib.import_module(stable_module_name) - except Exception as exc: - logger.warning( - "Failed to import stable task cfg module '%s' for manager_call_config stable mode: %s", - stable_module_name, - exc, - ) - return None - - stable_cfg_cls = getattr(stable_module, cfg_cls.__name__, None) - if stable_cfg_cls is None: - logger.warning("Stable task cfg class '%s' not found in module '%s'.", cfg_cls.__name__, stable_module_name) - return None - - try: - return stable_cfg_cls() - except Exception as exc: - logger.warning( - "Failed to instantiate stable task cfg '%s.%s': %s", - stable_module_name, - cfg_cls.__name__, - exc, - ) - return None diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py index dd9f2d2f7229..437bd2fe1031 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py @@ -10,15 +10,17 @@ import warp as wp +from isaaclab.utils.timer import Timer + class WarpGraphCache: """Execute Warp stages eagerly or through cached CUDA graphs. - On the very first call for a given stage, an **eager warm-up** run - executes *before* graph capture. This lets one-time initialisation - code (memory allocations, torch dtype casts, ``hasattr`` guards, etc.) - run outside the capture context. Only the steady-state kernel - launches are then recorded into the graph. + :meth:`call` uses the first invocation as an eager warm-up, captures the + second invocation, and replays later invocations. This lets one-time + initialization run outside capture without executing stateful manager work + twice in one environment step. :meth:`capture_or_replay` retains the direct + environment behavior of warming up and capturing on its first invocation. The return value from the capture run is cached and returned on every subsequent replay, ensuring captured stages return the same references @@ -27,54 +29,71 @@ class WarpGraphCache: Usage:: cache = WarpGraphCache() - result = cache.call("my_stage", my_warp_function) + result = cache.call("MyManager_step", my_warp_function) # uncaptured work here ... - result2 = cache.call("my_stage_post", my_other_function) + result2 = cache.call("OtherManager_step", my_other_function) """ - def __init__(self, *, enabled: bool = True): - """Initialize the execution cache. + def __init__(self, *, enabled: bool = True, device: wp.DeviceLike = None): + """Initialize the cache. Args: - enabled: Whether to capture stages. When false, stages execute eagerly. + enabled: Whether eligible stages may use CUDA graph capture. + device: Warp device used by cached stages. CPU devices always run eagerly. """ self._enabled = enabled + self._device = wp.get_device(device) if device is not None else None self._graphs: dict[str, Any] = {} self._results: dict[str, Any] = {} self._capturable: dict[str, bool] = {} + self._warmed: set[str] = set() def call( self, stage: str, fn: Callable[..., Any], - args: tuple[Any, ...] = (), - kwargs: dict[str, Any] | None = None, - *, - capture: bool | None = None, - group: str | None = None, + /, + *args: Any, + output: Callable[[Any], Any] | None = None, + timer: bool = False, + **kwargs: Any, ) -> Any: - """Execute a stage using its registered capture policy. + """Run a Warp frontend stage eagerly or through its cached CUDA graph. + + The stage prefix before the first underscore identifies its capturability + group. A group stays eager once any of its terms is registered as unsafe. + Eligible stages run eagerly once, capture on their second invocation, and + replay thereafter. Args: - stage: Unique stage identifier. + stage: Stage identifier in the form ``"GroupName_function_name"``. fn: Callable implementing the stage. - args: Positional arguments forwarded to :paramref:`fn`. - kwargs: Keyword arguments forwarded to :paramref:`fn`. - capture: Whether capture is requested for this call. When omitted, - follows the cache-wide :attr:`_enabled` setting. - group: Optional capturability group. This lets an owner register one - decision for several stages, such as all calls on one manager. + *args: Positional arguments forwarded to :paramref:`fn`. + output: Optional transform applied to the stage result after execution. + timer: Whether to time the stage execution. + **kwargs: Keyword arguments forwarded to :paramref:`fn`. Returns: - The eager stage result or cached capture result. + The stage result, optionally transformed by :paramref:`output`. """ - if kwargs is None: - kwargs = {} - capture_requested = self._enabled if capture is None else capture - capture_key = stage if group is None else group - if not self._enabled or not capture_requested or not self.is_capturable(capture_key): - return fn(*args, **kwargs) - return self._capture_or_replay(stage, fn, args, kwargs) + group = stage.partition("_")[0] + with Timer(name=stage, msg=f"{stage} took:", enable=timer, time_unit="us"): + if not self._capture_enabled or not self.is_capturable(group): + result = fn(*args, **kwargs) + elif stage in self._graphs: + wp.capture_launch(self._graphs[stage]) + result = self._results[stage] + elif stage not in self._warmed: + # The first real call doubles as warm-up. Stateful manager stages + # must execute exactly once per environment step. + result = fn(*args, **kwargs) + self._warmed.add(stage) + else: + result = self._capture(stage, fn, args, kwargs) + # CUDA graph capture records launches without executing them. Launch + # the new graph so this logical manager call still advances once. + wp.capture_launch(self._graphs[stage]) + return output(result) if output is not None else result def capture_or_replay( self, @@ -94,53 +113,62 @@ def capture_or_replay( kwargs: Keyword arguments forwarded to *fn*. Defaults to ``None``. Returns: - The eager result when disabled, otherwise the cached result from capture. + The cached return value from the first (capture) invocation. """ - return self.call(stage, fn, args, kwargs, capture=True) + if kwargs is None: + kwargs = {} + if not self._capture_enabled: + return fn(*args, **kwargs) + graph = self._graphs.get(stage) + if graph is not None: + wp.capture_launch(graph) + return self._results[stage] + # Warm-up: run eagerly to flush first-call allocations / hasattr guards. + fn(*args, **kwargs) + return self._capture(stage, fn, args, kwargs) - def register_capturability(self, key: str, capturable: bool) -> None: - """Register capture eligibility for a stage or owner group. + def _capture( + self, + stage: str, + fn: Callable[..., Any], + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> Any: + """Capture one already-warmed stage invocation.""" + with wp.ScopedCapture(device=self._device) as capture: + result = fn(*args, **kwargs) + self._graphs[stage] = capture.graph + self._results[stage] = result + self._warmed.add(stage) + return result - Repeated registrations are conservative: once a key is non-capturable, - later registrations cannot silently re-enable it. + def register_capturability(self, key: str, capturable: bool) -> None: + """Register conservative capture eligibility for a stage group. Args: - key: Stage or group identifier. - capturable: Whether every operation registered under the key is safe + key: Stage group identifier. + capturable: Whether every operation registered under the group is safe during CUDA graph capture. """ self._capturable[key] = self._capturable.get(key, True) and capturable def is_capturable(self, key: str) -> bool: - """Return whether a stage or group is eligible for capture.""" + """Return whether a stage group is eligible for capture.""" return self._capturable.get(key, True) + @property + def _capture_enabled(self) -> bool: + """Return whether this cache can capture on its configured device.""" + device = self._device if self._device is not None else wp.get_device() + return self._enabled and device.is_cuda + def invalidate(self, stage: str | None = None) -> None: """Drop cached graph(s). If *stage* is ``None``, drop all.""" if stage is None: self._graphs.clear() self._results.clear() + self._warmed.clear() else: self._graphs.pop(stage, None) self._results.pop(stage, None) - - def _capture_or_replay( - self, - stage: str, - fn: Callable[..., Any], - args: tuple[Any, ...], - kwargs: dict[str, Any], - ) -> Any: - """Capture a stage on first use and replay it afterward.""" - graph = self._graphs.get(stage) - if graph is not None: - wp.capture_launch(graph) - return self._results[stage] - # Warm-up: run eagerly to flush first-call allocations / hasattr guards. - fn(*args, **kwargs) - # Capture: allocations already done, only wp.launch calls are recorded. - with wp.ScopedCapture() as capture: - result = fn(*args, **kwargs) - self._graphs[stage] = capture.graph - self._results[stage] = result - return result + self._warmed.discard(stage) diff --git a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py index 402ea2bf44ec..480ec829cd97 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py @@ -293,7 +293,7 @@ def test_push_by_setting_velocity(self, warp_env, art_data, all_joints_cfg): "yaw": (0.0, 0.0), } term = _make_event_term( - warp_evt.PushBySettingVelocity, + warp_evt.push_by_setting_velocity, warp_env, mode="interval", velocity_range=zero_range, @@ -325,7 +325,7 @@ def test_apply_external_force_torque(self, warp_env, art_data, all_joints_cfg): ) zero_range = (0.0, 0.0) term = _make_event_term( - warp_evt.ApplyExternalForceTorque, + warp_evt.apply_external_force_torque, warp_env, force_range=zero_range, torque_range=zero_range, @@ -377,27 +377,7 @@ class TestEventStateOwnership: """Verify event scratch and parsed ranges are owned by one environment configuration.""" def test_com_term_is_marked_non_capturable(self): - assert warp_evt.RandomizeRigidBodyCom._warp_capturable is False - - def test_lowercase_adapter_warns_and_forwards(self): - env, data, asset = _make_event_env(88) - copy_np_to_wp(data.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32)) - captured = {} - asset.write_root_velocity_to_sim_mask = lambda **kwargs: captured.update(kwargs) - env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) - asset_cfg = SceneEntityCfg("robot") - asset_cfg.body_ids_wp = wp.array([0], dtype=wp.int32, device=DEVICE) - - with pytest.warns(DeprecationWarning, match="PushBySettingVelocity"): - warp_evt.push_by_setting_velocity( - env, - env_mask, - velocity_range={"x": (1.0, 1.0)}, - asset_cfg=asset_cfg, - ) - wp.synchronize() - - assert_close(wp.to_torch(captured["root_velocity"])[:, 0], torch.ones(NUM_ENVS, device=DEVICE)) + assert warp_evt.randomize_rigid_body_com._warp_capturable is False def test_push_state_is_not_shared_between_environments(self): env_a, data_a, asset_a = _make_event_env(101) @@ -411,8 +391,8 @@ def test_push_state_is_not_shared_between_environments(self): env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) range_a = {"x": (1.0, 1.0)} range_b = {"x": (2.0, 2.0)} - term_a = _make_event_term(warp_evt.PushBySettingVelocity, env_a, mode="interval", velocity_range=range_a) - term_b = _make_event_term(warp_evt.PushBySettingVelocity, env_b, mode="interval", velocity_range=range_b) + term_a = _make_event_term(warp_evt.push_by_setting_velocity, env_a, mode="interval", velocity_range=range_a) + term_b = _make_event_term(warp_evt.push_by_setting_velocity, env_b, mode="interval", velocity_range=range_b) term_a(env_a, env_mask, **term_a.cfg.params) term_b(env_b, env_mask, **term_b.cfg.params) @@ -432,8 +412,8 @@ def test_push_state_is_not_shared_between_configurations(self): env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) range_a = {"x": (1.0, 1.0)} range_b = {"x": (2.0, 2.0)} - term_a = _make_event_term(warp_evt.PushBySettingVelocity, env, mode="interval", velocity_range=range_a) - term_b = _make_event_term(warp_evt.PushBySettingVelocity, env, mode="interval", velocity_range=range_b) + term_a = _make_event_term(warp_evt.push_by_setting_velocity, env, mode="interval", velocity_range=range_a) + term_b = _make_event_term(warp_evt.push_by_setting_velocity, env, mode="interval", velocity_range=range_b) term_a(env, env_mask, **term_a.cfg.params) term_b(env, env_mask, **term_b.cfg.params) @@ -453,7 +433,7 @@ def test_push_term_snapshots_range_during_initialization(self): env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) velocity_range = {"x": (1.0, 1.0)} term = _make_event_term( - warp_evt.PushBySettingVelocity, + warp_evt.push_by_setting_velocity, env, mode="interval", velocity_range=velocity_range, @@ -477,7 +457,7 @@ def test_push_state_respects_sparse_mask(self): mask_np = np.arange(NUM_ENVS) % 3 == 0 env_mask = wp.array(mask_np, dtype=wp.bool, device=DEVICE) term = _make_event_term( - warp_evt.PushBySettingVelocity, + warp_evt.push_by_setting_velocity, env, mode="interval", velocity_range={"x": (3.0, 3.0)}, @@ -503,13 +483,13 @@ def test_external_wrench_state_is_not_shared_between_environments(self): force_range_b = (2.0, 2.0) torque_range = (0.0, 0.0) term_a = _make_event_term( - warp_evt.ApplyExternalForceTorque, + warp_evt.apply_external_force_torque, env_a, force_range=force_range_a, torque_range=torque_range, ) term_b = _make_event_term( - warp_evt.ApplyExternalForceTorque, + warp_evt.apply_external_force_torque, env_b, force_range=force_range_b, torque_range=torque_range, @@ -532,7 +512,7 @@ def test_external_wrench_respects_body_selection(self): asset_cfg = SceneEntityCfg("robot", body_ids=[1]) env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) term = _make_event_term( - warp_evt.ApplyExternalForceTorque, + warp_evt.apply_external_force_torque, env, force_range=(2.0, 2.0), torque_range=(0.0, 0.0), @@ -567,13 +547,13 @@ def test_root_reset_state_is_not_shared_between_environments(self): pose_range_b = {"x": (2.0, 2.0)} velocity_range = {} term_a = _make_event_term( - warp_evt.ResetRootStateUniform, + warp_evt.reset_root_state_uniform, env_a, pose_range=pose_range_a, velocity_range=velocity_range, ) term_b = _make_event_term( - warp_evt.ResetRootStateUniform, + warp_evt.reset_root_state_uniform, env_b, pose_range=pose_range_b, velocity_range=velocity_range, @@ -602,13 +582,13 @@ def test_com_ranges_are_not_shared_between_environments(self): com_range_a = {"x": (1.0, 1.0)} com_range_b = {"x": (2.0, 2.0)} term_a = _make_event_term( - warp_evt.RandomizeRigidBodyCom, + warp_evt.randomize_rigid_body_com, env_a, com_range=com_range_a, asset_cfg=asset_cfg_a, ) term_b = _make_event_term( - warp_evt.RandomizeRigidBodyCom, + warp_evt.randomize_rigid_body_com, env_b, com_range=com_range_b, asset_cfg=asset_cfg_b, @@ -632,7 +612,7 @@ def test_com_randomization_uses_persistent_baseline(self): env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) com_range = {"x": (1.0, 1.0)} term = _make_event_term( - warp_evt.RandomizeRigidBodyCom, + warp_evt.randomize_rigid_body_com, env, com_range=com_range, asset_cfg=asset_cfg, @@ -654,7 +634,7 @@ def test_com_randomization_broadcasts_one_offset_across_selected_bodies(self): asset_cfg = SceneEntityCfg("robot", body_ids=[0, 2]) env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) term = _make_event_term( - warp_evt.RandomizeRigidBodyCom, + warp_evt.randomize_rigid_body_com, env, com_range={"x": (-1.0, 1.0), "y": (-2.0, 2.0), "z": (-3.0, 3.0)}, asset_cfg=asset_cfg, diff --git a/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py b/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py deleted file mode 100644 index edb4051aec58..000000000000 --- a/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py +++ /dev/null @@ -1,94 +0,0 @@ -# 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 - -"""Tests for mask-first scene reset dispatch.""" - -from __future__ import annotations - -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest -import torch -import warp as wp -from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp - -from isaaclab.scene import InteractiveScene -from isaaclab.terrains import TerrainImporterCfg - - -class TestInteractiveSceneWarp: - """Tests for :class:`InteractiveSceneWarp`.""" - - def test_frontend_uses_canonical_terrain_storage_without_mutating_cfg( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Warp scenes should inherit canonical origin storage without replacing the configured class.""" - terrain_cfg = TerrainImporterCfg(prim_path="/World/Ground", terrain_type="plane") - cfg = SimpleNamespace(terrain=terrain_cfg) - class_type = terrain_cfg.class_type - origins = torch.zeros((2, 3), dtype=torch.float32) - origins_wp = wp.from_torch(origins, dtype=wp.vec3f) - terrain = SimpleNamespace(env_origins=origins, env_origins_wp=origins_wp) - - def initialize_scene(scene: InteractiveScene, scene_cfg: SimpleNamespace) -> None: - scene.cfg = scene_cfg - scene._terrain = terrain - - monkeypatch.setattr(InteractiveScene, "__init__", initialize_scene) - - scene = InteractiveSceneWarp(cfg) - - assert terrain_cfg.class_type is class_type - assert isinstance(scene.env_origins, torch.Tensor) - assert scene.env_origins is terrain.env_origins - assert scene.env_origins_wp is origins_wp - assert InteractiveSceneWarp.env_origins is InteractiveScene.env_origins - assert "env_origins_wp" not in InteractiveScene.__dict__ - - def test_mask_reset_stays_mask_based_for_supported_entities(self): - """Mask-based reset should not pass environment IDs to Warp-capable entities.""" - scene = InteractiveSceneWarp.__new__(InteractiveSceneWarp) - articulation = Mock() - deformable = Mock() - rigid_object = Mock() - rigid_collection = Mock() - sensor = Mock() - surface_gripper = Mock() - scene._articulations = {"articulation": articulation} - scene._deformable_objects = {"deformable": deformable} - scene._rigid_objects = {"rigid_object": rigid_object} - scene._rigid_object_collections = {"collection": rigid_collection} - scene._sensors = {"sensor": sensor} - scene._surface_grippers = {"gripper": surface_gripper} - scene.cfg = SimpleNamespace(num_envs=3) - scene.sim = SimpleNamespace(device="cpu") - env_mask = wp.array([True, False, True], dtype=wp.bool, device="cpu") - - scene.reset(env_mask=env_mask) - - for entity in (articulation, deformable, rigid_object, rigid_collection, sensor): - entity.reset.assert_called_once_with(env_mask=env_mask) - surface_gripper.reset.assert_not_called() - - def test_surface_gripper_reset_is_an_explicit_id_boundary(self): - """Surface grippers should be reset separately because their API is ID-based.""" - scene = InteractiveSceneWarp.__new__(InteractiveSceneWarp) - surface_gripper = Mock() - scene._surface_grippers = {"gripper": surface_gripper} - - scene.reset_host([0, 2]) - - surface_gripper.reset.assert_called_once_with([0, 2]) - - def test_mask_shape_is_validated_before_entity_dispatch(self): - """A malformed mask should fail before reaching scene entities.""" - scene = InteractiveSceneWarp.__new__(InteractiveSceneWarp) - scene.cfg = SimpleNamespace(num_envs=3) - scene.sim = SimpleNamespace(device="cpu") - env_mask = wp.array([True, False], dtype=wp.bool, device="cpu") - - with pytest.raises(ValueError, match="shape"): - scene.reset(env_mask=env_mask) diff --git a/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py index 49224c43708b..87a6e4c2b50d 100644 --- a/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py @@ -11,26 +11,9 @@ from unittest.mock import Mock import numpy as np -import pytest import torch import warp as wp from isaaclab_experimental.envs.manager_based_rl_env_warp import ManagerBasedRLEnvWarp -from isaaclab_experimental.utils.manager_call_switch import ManagerCallSwitch - - -def test_warp_environment_deprecates_and_preserves_stable_manager_profile() -> None: - """The legacy stable-manager profile should warn while retaining its config bridge.""" - cfg = SimpleNamespace(rewards="warp_rewards", actions="warp_actions") - switch = ManagerCallSwitch(cfg_source={"default": 1, "RewardManager": 0}) - switch._resolve_stable_cfg_counterpart = Mock( - return_value=SimpleNamespace(rewards="stable_rewards", actions="stable_actions") - ) - - with pytest.warns(DeprecationWarning, match="STABLE managers"): - switch.apply_term_cfg_profile(cfg) - - assert cfg.rewards == "stable_rewards" - assert cfg.actions == "warp_actions" def test_reset_without_host_consumers_does_not_compact_ids() -> None: @@ -39,11 +22,12 @@ def test_reset_without_host_consumers_does_not_compact_ids() -> None: reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") env.termination_manager = SimpleNamespace(dones_wp=reset_mask) env.reset_buf = Mock() + env.reset_buf.any.return_value.item.return_value = True env.reset_buf.nonzero.side_effect = AssertionError("reset IDs should not be materialized") - env._reset_requires_host_selection = Mock(return_value=False) env._reset_idx = Mock() + env._has_recorders = False env.has_rtx_sensors = False - env.cfg = SimpleNamespace(num_rerenders_on_reset=0) + env.cfg = SimpleNamespace(compute_final_obs=False, num_rerenders_on_reset=0) env.extras = {"log": {}} env._reset_terminated_envs() @@ -52,28 +36,6 @@ def test_reset_without_host_consumers_does_not_compact_ids() -> None: env._reset_idx.assert_called_once_with(env_mask=reset_mask) -def test_public_reset_selection_stays_as_mask_without_host_consumers() -> None: - """Mask-native reset selection should not be compacted preemptively.""" - env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) - env._reset_requires_host_selection = Mock(return_value=False) - reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") - - result = env._resolve_reset_host_ids(env_ids=None, env_mask=reset_mask) - - assert result is None - - -def test_public_reset_compacts_mask_only_for_host_consumer() -> None: - """A mask should cross to IDs only at an enabled host reset boundary.""" - env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) - env._reset_requires_host_selection = Mock(return_value=True) - reset_mask = wp.array([False, True, False, True], dtype=wp.bool, device="cpu") - - result = env._resolve_reset_host_ids(env_ids=None, env_mask=reset_mask) - - torch.testing.assert_close(result, torch.tensor([1, 3])) - - def test_curriculum_uses_mask_before_scene_reset() -> None: """Curriculum updates should stay mask-native and precede scene reset.""" env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) @@ -82,12 +44,16 @@ def test_curriculum_uses_mask_before_scene_reset() -> None: def call(stage, fn, /, *args, **kwargs): del fn, args stages.append((stage, kwargs)) - return None if stage.endswith("_compute") or stage == "Scene_reset" else {} + return None if stage.endswith("_compute") else {} - env._manager_call_switch = SimpleNamespace(call=call) + env._warp_graph_cache = SimpleNamespace(call=call) env.sim = SimpleNamespace(device="cpu") env.reset_mask_wp = wp.zeros(3, dtype=wp.bool, device="cpu") - env.scene = SimpleNamespace(num_envs=3, reset=Mock(), surface_grippers={}) + env.scene = SimpleNamespace( + num_envs=3, + reset=Mock(side_effect=lambda **kwargs: stages.append(("Scene_reset", kwargs))), + surface_grippers={}, + ) env.curriculum_manager = SimpleNamespace(compute=Mock(), reset=Mock()) env.event_manager = SimpleNamespace(available_modes=[], reset=Mock()) env.observation_manager = SimpleNamespace(reset=Mock()) @@ -119,7 +85,7 @@ def call(stage, fn, /, *args, **kwargs): seen.append((mask, mask.numpy().copy())) return None if stage.endswith("_compute") or stage == "Scene_reset" else {} - env._manager_call_switch = SimpleNamespace(call=call) + env._warp_graph_cache = SimpleNamespace(call=call) env.sim = SimpleNamespace(device="cpu") env.reset_mask_wp = wp.zeros(3, dtype=wp.bool, device="cpu") env.scene = SimpleNamespace(num_envs=3, reset=Mock(), surface_grippers={}) @@ -142,28 +108,28 @@ def call(stage, fn, /, *args, **kwargs): np.testing.assert_array_equal(seen[1][1], [False, False, False]) -def test_reset_compacts_ids_for_enabled_host_consumer() -> None: - """Host-only reset hooks should receive IDs while the Warp stage keeps its mask.""" +def test_reset_compacts_ids_only_for_active_recorders() -> None: + """Recorder callbacks should receive IDs while the Warp reset keeps its mask.""" env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") env.termination_manager = SimpleNamespace(dones_wp=reset_mask) env.reset_buf = torch.tensor([False, True, False]) - env._reset_requires_host_selection = Mock(return_value=True) env._reset_idx = Mock() - env._reset_host_pre = Mock() - env._reset_host_post = Mock(return_value={"host": 1.0}) env.recorder_manager = Mock() - env._has_recorders = False + env.recorder_manager.reset.return_value = {"recorder": 1.0} + env._has_recorders = True env.has_rtx_sensors = False - env.cfg = SimpleNamespace(num_rerenders_on_reset=0) + env.cfg = SimpleNamespace(compute_final_obs=False, num_rerenders_on_reset=0) env.extras = {"log": {}} env._reset_terminated_envs() - reset_ids = env._reset_host_pre.call_args.args[0] + reset_ids = env.recorder_manager.record_pre_reset.call_args.args[0] torch.testing.assert_close(reset_ids, torch.tensor([1])) env._reset_idx.assert_called_once_with(env_mask=reset_mask) - assert env.extras["log"]["host"] == 1.0 + env.recorder_manager.reset.assert_called_once_with(reset_ids) + env.recorder_manager.record_post_reset.assert_called_once_with(reset_ids) + assert env.extras["log"]["recorder"] == 1.0 def test_rtx_rerender_does_not_require_compact_reset_ids() -> None: @@ -172,56 +138,51 @@ def test_rtx_rerender_does_not_require_compact_reset_ids() -> None: reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") env.termination_manager = SimpleNamespace(dones_wp=reset_mask) env.reset_buf = torch.tensor([False, True, False]) - env._reset_requires_host_selection = Mock(return_value=False) env._reset_idx = Mock() - env._reset_host_pre = Mock() - env._reset_host_post = Mock() env._has_recorders = False env.has_rtx_sensors = True - env.cfg = SimpleNamespace(num_rerenders_on_reset=2) + env.cfg = SimpleNamespace(compute_final_obs=False, num_rerenders_on_reset=2) env.sim = SimpleNamespace(render=Mock()) env.extras = {"log": {}} env._reset_terminated_envs() env._reset_idx.assert_called_once_with(env_mask=reset_mask) - env._reset_host_pre.assert_not_called() - env._reset_host_post.assert_not_called() assert env.sim.render.call_count == 2 -def test_empty_reset_skips_legacy_host_boundary_and_preserves_logs() -> None: - """An empty host-visible selection should not execute reset side effects.""" +def test_empty_reset_skips_warp_pipeline_without_compacting_ids() -> None: + """An empty mask should skip eager reset stages without materializing IDs.""" env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) reset_mask = wp.zeros(3, dtype=wp.bool, device="cpu") env.termination_manager = SimpleNamespace(dones_wp=reset_mask) env.reset_buf = torch.zeros(3, dtype=torch.bool) - env._reset_requires_host_selection = Mock(return_value=True) env._reset_idx = Mock() - env._reset_host_pre = Mock() - env._reset_host_post = Mock() env._has_recorders = False env.extras = {"log": {"previous": 1.0}} env._reset_terminated_envs() env._reset_idx.assert_not_called() - env._reset_host_pre.assert_not_called() - env._reset_host_post.assert_not_called() assert env.extras["log"] == {"previous": 1.0} -def test_empty_reset_skips_warp_pipeline_without_compacting_ids() -> None: - """An empty mask should skip eager reset stages without materializing IDs.""" +def test_reset_captures_final_observation_before_reset() -> None: + """Same-step autoreset should expose the terminal observation before reset.""" env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) - reset_mask = wp.zeros(3, dtype=wp.bool, device="cpu") + reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") + final_obs = {"policy": torch.tensor([[1.0], [2.0], [3.0]])} env.termination_manager = SimpleNamespace(dones_wp=reset_mask) - env.reset_buf = torch.zeros(3, dtype=torch.bool) - env._reset_requires_host_selection = Mock(return_value=False) + env.reset_buf = torch.tensor([False, True, False]) + env.observation_manager = SimpleNamespace(compute=Mock(return_value=final_obs)) env._reset_idx = Mock() - env.extras = {"log": {"previous": 1.0}} + env._has_recorders = False + env.has_rtx_sensors = False + env.cfg = SimpleNamespace(compute_final_obs=True, num_rerenders_on_reset=0) + env.extras = {"log": {}} env._reset_terminated_envs() - env._reset_idx.assert_not_called() - assert env.extras["log"] == {"previous": 1.0} + assert env.extras["final_obs"] is final_obs + env.observation_manager.compute.assert_called_once_with() + env._reset_idx.assert_called_once_with(env_mask=reset_mask) diff --git a/source/isaaclab_experimental/test/managers/test_manager_base.py b/source/isaaclab_experimental/test/managers/test_manager_base.py index d2cdef5e2553..5450a49835ed 100644 --- a/source/isaaclab_experimental/test/managers/test_manager_base.py +++ b/source/isaaclab_experimental/test/managers/test_manager_base.py @@ -6,10 +6,13 @@ """Tests for Warp manager reset-mask contracts.""" from types import SimpleNamespace +from unittest.mock import Mock import pytest import warp as wp +from isaaclab_experimental.managers.action_manager import ActionManager from isaaclab_experimental.managers.manager_base import _resolve_reset_mask +from isaaclab_experimental.utils.warp import warp_capturable def test_reset_mask_rejects_wrong_device(monkeypatch: pytest.MonkeyPatch) -> None: @@ -21,3 +24,19 @@ def test_reset_mask_rejects_wrong_device(monkeypatch: pytest.MonkeyPatch) -> Non with pytest.raises(ValueError, match="env_mask must be on"): _resolve_reset_mask(owner, None, env_mask) + + +def test_class_term_capturability_is_registered() -> None: + """Class-based manager terms should honor explicit capture metadata.""" + + @warp_capturable(False) + class HostTerm: + pass + + graph_cache = Mock() + manager = object.__new__(ActionManager) + manager._env = SimpleNamespace(_warp_graph_cache=graph_cache) + + manager._register_term_capturability(HostTerm) + + graph_cache.register_capturability.assert_called_once_with("ActionManager", False) diff --git a/source/isaaclab_experimental/test/utils/test_manager_call_switch.py b/source/isaaclab_experimental/test/utils/test_manager_call_switch.py deleted file mode 100644 index 029c2eb236cf..000000000000 --- a/source/isaaclab_experimental/test/utils/test_manager_call_switch.py +++ /dev/null @@ -1,507 +0,0 @@ -# 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 - -"""Tests for ManagerCallSwitch config loading, mode resolution, and dispatch.""" - -from __future__ import annotations - -import json -import os -import unittest - -import warp as wp -from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode, ManagerCallSwitch - - -@wp.kernel -def _add_one(a: wp.array(dtype=wp.float32), b: wp.array(dtype=wp.float32)): - i = wp.tid() - b[i] = a[i] + 1.0 - - -class TestManagerCallMode(unittest.TestCase): - """Tests for the ManagerCallMode enum.""" - - def test_enum_values(self): - self.assertEqual(ManagerCallMode.STABLE, 0) - self.assertEqual(ManagerCallMode.WARP_NOT_CAPTURED, 1) - self.assertEqual(ManagerCallMode.WARP_CAPTURED, 2) - - def test_ordering(self): - self.assertLess(ManagerCallMode.STABLE, ManagerCallMode.WARP_NOT_CAPTURED) - self.assertLess(ManagerCallMode.WARP_NOT_CAPTURED, ManagerCallMode.WARP_CAPTURED) - - -# ====================================================================== -# Config loading -# ====================================================================== - - -class TestConfigLoading(unittest.TestCase): - """Tests for ManagerCallSwitch config parsing from dict, JSON, env var, and None.""" - - def test_none_uses_default(self): - switch = ManagerCallSwitch(cfg_source=None) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_NOT_CAPTURED) - - def test_dict_config(self): - switch = ManagerCallSwitch(cfg_source={"default": 1}) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_NOT_CAPTURED) - - def test_dict_per_manager_override(self): - switch = ManagerCallSwitch(cfg_source={"default": 2, "RewardManager": 0}) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.STABLE) - self.assertEqual(switch.get_mode_for_manager("ActionManager"), ManagerCallMode.WARP_CAPTURED) - - def test_dict_without_default_key(self): - """A dict missing 'default' should inherit from DEFAULT_CONFIG.""" - switch = ManagerCallSwitch(cfg_source={"RewardManager": 0}) - # default should be Warp eager (from DEFAULT_CONFIG) - self.assertEqual(switch.get_mode_for_manager("ActionManager"), ManagerCallMode.WARP_NOT_CAPTURED) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.STABLE) - - def test_json_string_config(self): - cfg_str = json.dumps({"default": 1, "ObservationManager": 0}) - switch = ManagerCallSwitch(cfg_source=cfg_str) - self.assertEqual(switch.get_mode_for_manager("ObservationManager"), ManagerCallMode.STABLE) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_NOT_CAPTURED) - - def test_empty_string_uses_default(self): - switch = ManagerCallSwitch(cfg_source="") - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_NOT_CAPTURED) - - def test_env_var_fallback(self): - """When cfg_source is None, should read from MANAGER_CALL_CONFIG env var.""" - old = os.environ.get(ManagerCallSwitch.ENV_VAR) - try: - os.environ[ManagerCallSwitch.ENV_VAR] = json.dumps({"default": 0}) - switch = ManagerCallSwitch(cfg_source=None) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.STABLE) - finally: - if old is None: - os.environ.pop(ManagerCallSwitch.ENV_VAR, None) - else: - os.environ[ManagerCallSwitch.ENV_VAR] = old - - def test_invalid_json_raises(self): - with self.assertRaises(json.JSONDecodeError): - ManagerCallSwitch(cfg_source="not valid json") - - def test_invalid_mode_value_raises(self): - with self.assertRaises(ValueError): - ManagerCallSwitch(cfg_source={"default": 99}) - - def test_non_int_mode_raises(self): - with self.assertRaises(TypeError): - ManagerCallSwitch(cfg_source={"default": "fast"}) - - def test_invalid_cfg_type_raises(self): - with self.assertRaises(TypeError): - ManagerCallSwitch(cfg_source=42) - - -# ====================================================================== -# Mode resolution and capping -# ====================================================================== - - -class TestModeResolution(unittest.TestCase): - """Tests for mode resolution, MAX_MODE_OVERRIDES, and dynamic capturability.""" - - def test_max_mode_override_caps_default(self): - """Scene is capped to WARP_NOT_CAPTURED by MAX_MODE_OVERRIDES.""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - # Scene has a static cap of WARP_NOT_CAPTURED - self.assertEqual( - switch.get_mode_for_manager("Scene"), - ManagerCallMode.WARP_NOT_CAPTURED, - ) - # Other managers should not be capped - self.assertEqual( - switch.get_mode_for_manager("RewardManager"), - ManagerCallMode.WARP_CAPTURED, - ) - - def test_max_modes_kwarg(self): - """max_modes kwarg should cap managers beyond MAX_MODE_OVERRIDES.""" - switch = ManagerCallSwitch( - cfg_source={"default": 2}, - max_modes={"RewardManager": ManagerCallMode.WARP_NOT_CAPTURED}, - ) - self.assertEqual( - switch.get_mode_for_manager("RewardManager"), - ManagerCallMode.WARP_NOT_CAPTURED, - ) - - def test_register_manager_capturability_downgrades(self): - """register_manager_capturability(False) should cap a manager to WARP_NOT_CAPTURED.""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_CAPTURED) - - switch.register_manager_capturability("RewardManager", capturable=False) - self.assertEqual( - switch.get_mode_for_manager("RewardManager"), - ManagerCallMode.WARP_NOT_CAPTURED, - ) - - def test_register_capturability_true_is_noop(self): - """register_manager_capturability(True) should not change anything.""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - switch.register_manager_capturability("RewardManager", capturable=True) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_CAPTURED) - - def test_register_capturability_does_not_upgrade(self): - """If a manager is already capped to NOT_CAPTURED, registering capturable=False again shouldn't change it.""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - switch.register_manager_capturability("RewardManager", capturable=False) - switch.register_manager_capturability("RewardManager", capturable=False) - self.assertEqual( - switch.get_mode_for_manager("RewardManager"), - ManagerCallMode.WARP_NOT_CAPTURED, - ) - - def test_capturability_interacts_with_static_cap(self): - """Dynamic capturability should respect existing static caps.""" - switch = ManagerCallSwitch( - cfg_source={"default": 2}, - max_modes={"RewardManager": ManagerCallMode.WARP_NOT_CAPTURED}, - ) - # Already capped, register_capturability(False) should be harmless - switch.register_manager_capturability("RewardManager", capturable=False) - self.assertEqual( - switch.get_mode_for_manager("RewardManager"), - ManagerCallMode.WARP_NOT_CAPTURED, - ) - - def test_mode_0_not_affected_by_cap(self): - """A manager explicitly set to STABLE (0) should stay at 0 even with cap at 1.""" - switch = ManagerCallSwitch(cfg_source={"default": 2, "RewardManager": 0}) - switch.register_manager_capturability("RewardManager", capturable=False) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.STABLE) - - -# ====================================================================== -# Stage dispatch -# ====================================================================== - - -class TestStageDispatch(unittest.TestCase): - """Tests for call_stage routing through stable / warp-eager / warp-captured paths.""" - - def test_call_forwards_normal_arguments_and_output(self): - """The concise frontend call should not require a call-spec dictionary.""" - switch = ManagerCallSwitch(cfg_source={"default": 1}) - - result = switch.call( - "RewardManager_compute", - lambda value, *, scale: value * scale, - 3, - scale=4, - _output=lambda value: value + 1, - ) - - self.assertEqual(result, 13) - - def test_stable_mode_calls_stable_fn(self): - switch = ManagerCallSwitch(cfg_source={"default": 0}) - called = {"stable": False, "warp": False} - - def stable_fn(): - called["stable"] = True - return "stable_result" - - def warp_fn(): - called["warp"] = True - return "warp_result" - - result = switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": warp_fn}, - stable_call={"fn": stable_fn}, - ) - self.assertTrue(called["stable"]) - self.assertFalse(called["warp"]) - self.assertEqual(result, "stable_result") - - def test_stable_mode_without_stable_call_raises(self): - switch = ManagerCallSwitch(cfg_source={"default": 0}) - with self.assertRaises(ValueError): - switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": lambda: None}, - stable_call=None, - ) - - def test_warp_eager_mode_calls_warp_fn(self): - switch = ManagerCallSwitch(cfg_source={"default": 1}) - called = {"stable": False, "warp": False} - - def stable_fn(): - called["stable"] = True - return "stable_result" - - def warp_fn(): - called["warp"] = True - return "warp_result" - - result = switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": warp_fn}, - stable_call={"fn": stable_fn}, - ) - self.assertFalse(called["stable"]) - self.assertTrue(called["warp"]) - self.assertEqual(result, "warp_result") - - def test_output_transform(self): - """The 'output' key in call spec should transform the return value.""" - switch = ManagerCallSwitch(cfg_source={"default": 1}) - - def warp_fn(): - return [1, 2, 3] - - result = switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": warp_fn, "output": len}, - ) - self.assertEqual(result, 3) - - def test_args_and_kwargs_forwarded(self): - """call_stage should forward args and kwargs from the call spec.""" - switch = ManagerCallSwitch(cfg_source={"default": 1}) - received = {} - - def warp_fn(a, b, key=None): - received["a"] = a - received["b"] = b - received["key"] = key - - switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": warp_fn, "args": (1, 2), "kwargs": {"key": "val"}}, - ) - self.assertEqual(received, {"a": 1, "b": 2, "key": "val"}) - - -class TestStageDispatchCaptured(unittest.TestCase): - """Tests for WARP_CAPTURED mode (requires GPU).""" - - def setUp(self): - self.device = "cuda:0" - - def test_captured_mode_produces_correct_output(self): - """WARP_CAPTURED should capture and replay a warp kernel correctly.""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - src = wp.full(4, value=5.0, dtype=wp.float32, device=self.device) - dst = wp.zeros(4, dtype=wp.float32, device=self.device) - - def warp_fn(): - wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device) - return dst - - # First call: warm-up + capture - result = switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": warp_fn}, - ) - self.assertAlmostEqual(result.numpy()[0], 6.0, places=5) - - # Replay - wp.copy(src, wp.full(4, value=10.0, dtype=wp.float32, device=self.device)) - result2 = switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": warp_fn}, - ) - self.assertIs(result, result2, "Replay must return same reference") - self.assertAlmostEqual(result2.numpy()[0], 11.0, places=5) - - def test_call_uses_the_same_signature_when_captured(self): - """The concise call should capture without changing the function signature.""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - src = wp.full(4, value=2.0, dtype=wp.float32, device=self.device) - dst = wp.zeros(4, dtype=wp.float32, device=self.device) - - def warp_fn(source, destination): - wp.launch(_add_one, dim=4, inputs=[source, destination], device=self.device) - return destination - - result = switch.call("RewardManager_compute", warp_fn, src, dst) - self.assertAlmostEqual(result.numpy()[0], 3.0, places=5) - - wp.copy(src, wp.full(4, value=8.0, dtype=wp.float32, device=self.device)) - replay = switch.call("RewardManager_compute", warp_fn, src, dst) - self.assertIs(result, replay) - self.assertAlmostEqual(replay.numpy()[0], 9.0, places=5) - - def test_captured_warmup_call_count(self): - """WARP_CAPTURED first call should invoke fn exactly 2 times (warm-up + capture).""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - src = wp.zeros(4, dtype=wp.float32, device=self.device) - dst = wp.zeros(4, dtype=wp.float32, device=self.device) - call_count = [0] - - def warp_fn(): - call_count[0] += 1 - wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device) - return dst - - # First call_stage: warm-up (1) + capture (2) - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 2, "First captured call should invoke fn twice (warm-up + capture)") - - # Second call_stage: replay only — no new fn invocation - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 2, "Replay should not invoke fn again") - - # Third call_stage: still replay - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 2) - - def test_captured_warmup_handles_hasattr_guard(self): - """Warm-up should flush first-call allocations so capture doesn't record them. - - This simulates the real-world pattern where MDP terms allocate scratch - buffers on first call using hasattr guards. - """ - switch = ManagerCallSwitch(cfg_source={"default": 2}) - src = wp.ones(8, dtype=wp.float32, device=self.device) - holder = {} - - def fn_with_guard(): - if "buf" not in holder: - # First-call allocation — must happen during warm-up, not capture - holder["buf"] = wp.zeros(8, dtype=wp.float32, device=self.device) - wp.launch(_add_one, dim=8, inputs=[src, holder["buf"]], device=self.device) - return holder["buf"] - - # Should not raise — warm-up handles the allocation outside capture context - result = switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": fn_with_guard}, - ) - result_np = result.numpy() - for val in result_np: - self.assertAlmostEqual(val, 2.0, places=5) - - # Replay should also work (allocation already done, only kernel replays) - wp.copy(src, wp.full(8, value=5.0, dtype=wp.float32, device=self.device)) - result2 = switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": fn_with_guard}, - ) - for val in result2.numpy(): - self.assertAlmostEqual(val, 6.0, places=5) - - def test_warp_eager_no_warmup(self): - """WARP_NOT_CAPTURED mode should call fn exactly once per call (no warm-up).""" - switch = ManagerCallSwitch(cfg_source={"default": 1}) - call_count = [0] - - def warp_fn(): - call_count[0] += 1 - - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 1, "Eager mode should call fn exactly once") - - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 2, "Eager mode should call fn again each time") - - -# ====================================================================== -# Graph invalidation -# ====================================================================== - - -class TestGraphInvalidation(unittest.TestCase): - """Tests for invalidate_graphs on ManagerCallSwitch.""" - - def setUp(self): - self.device = "cuda:0" - - def test_invalidate_forces_recapture(self): - switch = ManagerCallSwitch(cfg_source={"default": 2}) - src = wp.full(4, value=1.0, dtype=wp.float32, device=self.device) - dst = wp.zeros(4, dtype=wp.float32, device=self.device) - - call_count = [0] - - def warp_fn(): - call_count[0] += 1 - wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device) - return dst - - # First capture - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 2) # warm-up + capture - - # Replay — no new calls - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 2) - - # Invalidate - switch.invalidate_graphs() - - # Should re-warm-up + re-capture - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 4) - - -# ====================================================================== -# resolve_manager_class -# ====================================================================== - - -class TestResolveManagerClass(unittest.TestCase): - """Tests for resolve_manager_class.""" - - def test_stable_resolves_to_isaaclab_managers(self): - switch = ManagerCallSwitch(cfg_source={"default": 0}) - cls = switch.resolve_manager_class("RewardManager") - self.assertTrue(cls.__module__.startswith("isaaclab.managers")) - - def test_warp_resolves_to_experimental_managers(self): - switch = ManagerCallSwitch(cfg_source={"default": 1}) - cls = switch.resolve_manager_class("RewardManager") - self.assertTrue(cls.__module__.startswith("isaaclab_experimental")) - - def test_mode_override(self): - """mode_override should bypass the config for class resolution.""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - cls = switch.resolve_manager_class("RewardManager", mode_override=ManagerCallMode.STABLE) - self.assertTrue(cls.__module__.startswith("isaaclab.managers")) - - def test_invalid_manager_raises(self): - switch = ManagerCallSwitch(cfg_source={"default": 0}) - with self.assertRaises(AttributeError): - switch.resolve_manager_class("NonexistentManager") - - -# ====================================================================== -# Manager name parsing -# ====================================================================== - - -class TestManagerNameParsing(unittest.TestCase): - """Tests for stage name → manager name extraction.""" - - def test_valid_stage_name(self): - switch = ManagerCallSwitch(cfg_source={"default": 1}) - # Dispatch a stage with valid name format - called = [False] - - def fn(): - called[0] = True - - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": fn}) - self.assertTrue(called[0]) - - def test_invalid_stage_name_raises(self): - switch = ManagerCallSwitch(cfg_source={"default": 1}) - with self.assertRaises(ValueError): - switch.call_stage(stage="nounderscore", warp_call={"fn": lambda: None}) - - -if __name__ == "__main__": - unittest.main() diff --git a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py index e7136c19ab54..62585d30d62f 100644 --- a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py +++ b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py @@ -20,48 +20,119 @@ def _add_one(a: wp.array(dtype=wp.float32), b: wp.array(dtype=wp.float32)): b[i] = a[i] + 1.0 +@wp.kernel +def _increment(value: wp.array(dtype=wp.int32)): + """Increment a scalar state buffer once.""" + value[0] += 1 + + class TestWarpGraphCache(unittest.TestCase): """Tests for :class:`WarpGraphCache`.""" def setUp(self): self.device = "cuda:0" - self.cache = WarpGraphCache() + self.cache = WarpGraphCache(device=self.device) + + # ------------------------------------------------------------------ + # Capture policy + # ------------------------------------------------------------------ - def test_disabled_cache_runs_eager_once_per_call(self): - """A disabled cache should preserve eager execution semantics.""" - cache = WarpGraphCache(enabled=False) + def test_non_capturable_group_stays_eager(self): + """Every stage in a non-capturable group should execute eagerly.""" call_count = 0 - def counted_call(value): + def counted_call(): nonlocal call_count call_count += 1 - return value - self.assertEqual(cache.capture_or_replay("stage", counted_call, args=(1,)), 1) - self.assertEqual(cache.capture_or_replay("stage", counted_call, args=(2,)), 2) + self.cache.register_capturability("RewardManager", False) + self.cache.call("RewardManager_compute", counted_call) + self.cache.call("RewardManager_reset", counted_call) + self.assertEqual(call_count, 2) - def test_non_capturable_group_stays_eager(self): - """A group-level safety decision should keep every member stage eager.""" + def test_cpu_device_stays_eager(self): + """CPU stages should never attempt CUDA graph capture.""" + cache = WarpGraphCache(device="cpu") call_count = 0 def counted_call(): nonlocal call_count call_count += 1 - self.cache.register_capturability("RewardManager", False) - self.cache.call("RewardManager_compute", counted_call, capture=True, group="RewardManager") - self.cache.call("RewardManager_reset", counted_call, capture=True, group="RewardManager") + for _ in range(3): + cache.call("RewardManager_compute", counted_call) - self.assertEqual(call_count, 2) + self.assertEqual(call_count, 3) + + def test_disabled_cache_stays_eager(self): + """An owner should be able to keep an unvalidated stage boundary eager.""" + cache = WarpGraphCache(enabled=False, device=self.device) + call_count = 0 + + def counted_call(): + nonlocal call_count + call_count += 1 + + for _ in range(3): + cache.call("direct_stage", counted_call) + + self.assertEqual(call_count, 3) def test_capturability_registration_is_conservative(self): - """A later positive registration must not override a known unsafe operation.""" + """A positive registration must not override a known unsafe operation.""" self.cache.register_capturability("EventManager", False) self.cache.register_capturability("EventManager", True) self.assertFalse(self.cache.is_capturable("EventManager")) + def test_call_forwards_arguments_and_applies_output(self): + """The frontend call should forward normal arguments and transform output.""" + self.cache.register_capturability("RewardManager", False) + + result = self.cache.call( + "RewardManager_compute", + lambda value, *, scale: value * scale, + 3, + scale=4, + output=lambda value: value + 1, + ) + + self.assertEqual(result, 13) + + def test_call_captures_by_default(self): + """A capturable stage should warm up, capture, and then replay across calls.""" + call_count = 0 + src = wp.full(4, value=2.0, dtype=wp.float32, device=self.device) + dst = wp.zeros(4, dtype=wp.float32, device=self.device) + + def counted_launch(): + nonlocal call_count + call_count += 1 + wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device) + return dst + + result = self.cache.call("RewardManager_compute", counted_launch) + captured = self.cache.call("RewardManager_compute", counted_launch) + replay = self.cache.call("RewardManager_compute", counted_launch) + + self.assertEqual(call_count, 2) + self.assertIs(result, captured) + self.assertIs(result, replay) + self.assertAlmostEqual(replay.numpy()[0], 3.0, places=5) + + def test_call_executes_stateful_stage_once_per_logical_call(self): + """Warm-up and capture should not advance manager state twice in one call.""" + state = wp.zeros(1, dtype=wp.int32, device=self.device) + + def increment_state(): + wp.launch(_increment, dim=1, inputs=[state], device=self.device) + return state + + for expected in (1, 2, 3): + result = self.cache.call("RewardManager_stateful", increment_state) + self.assertEqual(result.numpy()[0], expected) + # ------------------------------------------------------------------ # Warm-up # ------------------------------------------------------------------ diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py index a8b7964617a4..66a9131011f3 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py @@ -100,7 +100,7 @@ class EventCfg: """Configuration for events.""" reset_base = EventTerm( - func=mdp.ResetRootStateUniform, + func=mdp.reset_root_state_uniform, mode="reset", params={"pose_range": {}, "velocity_range": {}}, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py index 49bdb600e4e2..9de689ffb98e 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py @@ -109,7 +109,7 @@ class EventCfg: """Configuration for events.""" reset_base = EventTerm( - func=mdp.ResetRootStateUniform, + func=mdp.reset_root_state_uniform, mode="reset", params={"pose_range": {}, "velocity_range": {}}, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py index 3a7d80af470f..6b5e1a1dda00 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py @@ -15,7 +15,6 @@ class AnymalBRoughEnvCfg(LocomotionVelocityRoughEnvCfg): def __post_init__(self): super().__post_init__() self.scene.robot = ANYMAL_B_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.manager_call_max_mode = {"Scene": 1} @configclass diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py index 1d2f2676c64a..686a5b362e12 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py @@ -21,7 +21,6 @@ def __post_init__(self): # switch robot to anymal-c self.scene.robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") self.scene.robot.actuators["legs"].armature = 0.01 - self.manager_call_max_mode = {"Scene": 1} @configclass diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py index 71f49edd2e68..0f5a11ece015 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py @@ -18,7 +18,6 @@ class AnymalDRoughEnvCfg(LocomotionVelocityRoughEnvCfg): def __post_init__(self): super().__post_init__() self.scene.robot = ANYMAL_D_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.manager_call_max_mode = {"Scene": 1} @configclass diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py index cb3a602394a6..82e724a654d4 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py @@ -15,7 +15,6 @@ class UnitreeGo1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): def __post_init__(self): super().__post_init__() self.scene.robot = UNITREE_GO1_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.manager_call_max_mode = {"Scene": 1} self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1) self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06) self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py index 8999c4a01a64..2807bf2c5a31 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py @@ -146,7 +146,7 @@ class EventCfg: # FIXME(warp-migration): COM randomization in exp manager-based locomotion currently causes # NaNs and is temporarily disabled. # base_com = EventTerm( - # func=mdp.RandomizeRigidBodyCom, + # func=mdp.randomize_rigid_body_com, # mode="startup", # params={ # "asset_cfg": SceneEntityCfg("robot", body_names="base"), @@ -157,7 +157,7 @@ class EventCfg: # reset base_external_force_torque = EventTerm( - func=mdp.ApplyExternalForceTorque, + func=mdp.apply_external_force_torque, mode="reset", params={ "asset_cfg": SceneEntityCfg("robot", body_names="base"), @@ -167,7 +167,7 @@ class EventCfg: ) reset_base = EventTerm( - func=mdp.ResetRootStateUniform, + func=mdp.reset_root_state_uniform, mode="reset", params={ "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, @@ -193,7 +193,7 @@ class EventCfg: # interval push_robot = EventTerm( - func=mdp.PushBySettingVelocity, + func=mdp.push_by_setting_velocity, mode="interval", interval_range_s=(10.0, 15.0), params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}}, diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py index e7c2684b7f96..ebb94cab0c5c 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py @@ -109,7 +109,7 @@ class EventCfg: """Configuration for events.""" reset_base = EventTerm( - func=mdp.ResetRootStateUniform, + func=mdp.reset_root_state_uniform, mode="reset", params={"pose_range": {}, "velocity_range": {}}, ) From a3bba1d918d44ea4daece0635a59a428014df8c4 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 21 Jul 2026 00:53:40 -0700 Subject: [PATCH 44/58] Tighten mask boundaries and changelog accuracy Remove the orphan isaaclab_tasks fragment and move the pose-command quaternion fix entry to isaaclab_experimental. Route surface-gripper mask resets through reset_mask, delegate the scene's Warp origins view to the terrain owner, and make terrain origin reconfiguration pointer-stable. Skip Lab actuator resets in Newton-native mode, share one compact-ID materialization per reset between the recorder and legacy curriculum terms, and downgrade the collection's mask-write docstrings to match their host compaction. Add a static gate that requires a reviewed mask-boundary marker on production nonzero sites. --- .../jichuanh-warp-frontend-cleanup.minor.rst | 7 +- .../isaaclab/scene/interactive_scene.py | 37 ++++++--- .../isaaclab/terrains/terrain_importer.py | 60 +++++++++++--- .../test/scene/test_interactive_scene.py | 10 ++- .../jichuanh-warp-frontend-cleanup.minor.rst | 2 + .../envs/direct_rl_env_warp.py | 7 +- .../envs/manager_based_env_warp.py | 3 +- .../envs/manager_based_rl_env_warp.py | 20 ++++- .../managers/curriculum_manager.py | 12 ++- .../envs/test_manager_based_rl_env_warp.py | 43 ++++++++-- .../test/managers/test_curriculum_manager.py | 56 +++++++++++++ .../test/utils/test_mask_boundary_gate.py | 39 ++++++++++ .../assets/articulation/articulation.py | 25 +++--- .../rigid_object_collection.py | 78 ++++++++++--------- .../cloner/newton_clone_utils.py | 4 +- .../envs/mdp/actions/newton_ik_actions.py | 2 +- .../isaaclab_newton/sensors/pva/pva.py | 3 +- .../test_newton_mask_reset_forwarding.py | 25 +++++- .../jichuanh-warp-frontend-cleanup.rst | 5 -- 19 files changed, 331 insertions(+), 107 deletions(-) create mode 100644 source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py delete mode 100644 source/isaaclab_tasks/changelog.d/jichuanh-warp-frontend-cleanup.rst diff --git a/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst index cf5e9d3e6841..32fbd10d6748 100644 --- a/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst +++ b/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -5,8 +5,5 @@ Added :attr:`~isaaclab.terrains.TerrainImporter.terrain_levels_wp`, and :meth:`~isaaclab.terrains.TerrainImporter.update_env_origins_mask` for zero-copy Warp terrain access and mask-based curriculum updates. - -Fixed -^^^^^ - -* Fixed identity quaternion initialization for uniform pose commands. +* Added :attr:`~isaaclab.scene.InteractiveScene.env_origins_wp` and a boolean ``env_mask`` + option to :meth:`~isaaclab.scene.InteractiveScene.reset` for mask-based scene resets. diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index 9a7ab623bbfa..538cbc19d33a 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -190,9 +190,11 @@ def __init__(self, cfg: InteractiveSceneCfg): self._aggregate_scene_data_requirements(requested_viz_types) - # The clone plan and terrain importer own the Torch origins buffer. Cache - # one zero-copy Warp view so frontend code does not rebuild it per access. - self._env_origins_wp = wp.from_torch(self.env_origins, dtype=wp.vec3f) + # The terrain importer owns its origins buffer together with its Warp view. + # For clone-plan origins, :attr:`env_origins_wp` caches a zero-copy view on + # first access (manager/term initialization, outside capture), so scene + # construction does not require a published clone plan. + self._clone_origins_wp = None # Collision filtering is PhysX-only (matches both physx and ovphysx). if self.cfg.filter_collisions and "physx" in self.physics_backend and self._is_scene_setup_from_cfg(): @@ -377,8 +379,17 @@ def env_origins(self) -> torch.Tensor: @property def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): - """Cached zero-copy Warp view of environment origins [m], shape ``(num_envs,)``.""" - return self._env_origins_wp + """Zero-copy Warp view of environment origins [m], shape ``(num_envs,)``. + + Terrain-backed scenes forward the terrain importer's cached view so both + accessors always share one storage owner. Clone-plan origins are wrapped + lazily on first access and cached. + """ + if self._terrain is not None: + return self._terrain.env_origins_wp + if self._clone_origins_wp is None: + self._clone_origins_wp = wp.from_torch(self.env_origins, dtype=wp.vec3f) + return self._clone_origins_wp @property def terrain(self) -> TerrainImporter | None: @@ -481,12 +492,16 @@ def reset( deformable_object.reset(**reset_kwargs) for rigid_object in self._rigid_objects.values(): rigid_object.reset(**reset_kwargs) - if env_mask is not None and self._surface_grippers: - # Surface grippers expose only the legacy ID API. Materialize IDs at - # this optional PhysX boundary without penalizing Warp-native scenes. - env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) - for surface_gripper in self._surface_grippers.values(): - surface_gripper.reset(env_ids) + if env_mask is not None: + # Surface grippers are CPU-only assets driven through a Torch API. + # Hand them the boolean view and let their compatibility layer own + # any index materialization. + gripper_env_mask = wp.to_torch(env_mask) if self._surface_grippers else None + for surface_gripper in self._surface_grippers.values(): + surface_gripper.reset_mask(gripper_env_mask) + else: + for surface_gripper in self._surface_grippers.values(): + surface_gripper.reset(env_ids) for rigid_object_collection in self._rigid_object_collections.values(): rigid_object_collection.reset(**reset_kwargs) # -- sensors diff --git a/source/isaaclab/isaaclab/terrains/terrain_importer.py b/source/isaaclab/isaaclab/terrains/terrain_importer.py index 570ecede58e0..5f9fee1d4d3a 100644 --- a/source/isaaclab/isaaclab/terrains/terrain_importer.py +++ b/source/isaaclab/isaaclab/terrains/terrain_importer.py @@ -336,6 +336,9 @@ def import_usd(self, name: str, usd_path: str): def configure_env_origins(self, origins: np.ndarray | torch.Tensor | None = None) -> None: """Configure the origins of the environments based on the added terrain. + Reconfiguration copies into the existing origin buffers when layouts match, + so cached Torch tensors and Warp views held by consumers stay valid. + Args: origins: The origins [m] of the sub-terrains. Shape is (num_rows, num_cols, 3). """ @@ -345,16 +348,22 @@ def configure_env_origins(self, origins: np.ndarray | torch.Tensor | None = None if isinstance(origins, np.ndarray): origins = torch.from_numpy(origins) # store the origins - self.terrain_origins = origins.to(self.device, dtype=torch.float) + self.terrain_origins = self._assign_pointer_stable( + self.terrain_origins, origins.to(self.device, dtype=torch.float) + ) # compute environment origins - self.env_origins = self._compute_env_origins_curriculum(self.cfg.num_envs, self.terrain_origins) + self.env_origins = self._assign_pointer_stable( + self.env_origins, self._compute_env_origins_curriculum(self.cfg.num_envs, self.terrain_origins) + ) else: self.terrain_origins = None # check if env spacing is valid if self.cfg.env_spacing is None: raise ValueError("Environment spacing must be specified for configuring grid-like origins.") # compute environment origins - self.env_origins = self._compute_env_origins_grid(self.cfg.num_envs, self.cfg.env_spacing) + self.env_origins = self._assign_pointer_stable( + self.env_origins, self._compute_env_origins_grid(self.cfg.num_envs, self.cfg.env_spacing) + ) self._configure_warp_origin_views() def update_env_origins(self, env_ids: torch.Tensor, move_up: torch.Tensor, move_down: torch.Tensor) -> None: @@ -440,17 +449,38 @@ def update_env_origins_mask( Internal helpers. """ + @staticmethod + def _assign_pointer_stable(current: torch.Tensor | None, new: torch.Tensor) -> torch.Tensor: + """Copy into the existing buffer when layouts match so cached views stay valid.""" + if ( + current is not None + and current.shape == new.shape + and current.dtype == new.dtype + and current.device == new.device + ): + current.copy_(new) + return current + return new + def _configure_warp_origin_views(self) -> None: - """Create persistent zero-copy Warp views after configuring the Torch buffers.""" - self._env_origins_wp = wp.from_torch(self.env_origins, dtype=wp.vec3f) + """Create persistent zero-copy Warp views after configuring the Torch buffers. + + Existing views are kept when reconfiguration reuses the underlying Torch + storage, so consumers holding these views stay valid. + """ + if self._env_origins_wp is None or self._env_origins_wp.ptr != self.env_origins.data_ptr(): + self._env_origins_wp = wp.from_torch(self.env_origins, dtype=wp.vec3f) if self.terrain_origins is None: self._terrain_levels_wp = None self._terrain_types_wp = None self._terrain_origins_wp = None return - self._terrain_levels_wp = wp.from_torch(self.terrain_levels, dtype=wp.int64) - self._terrain_types_wp = wp.from_torch(self.terrain_types, dtype=wp.int64) - self._terrain_origins_wp = wp.from_torch(self.terrain_origins, dtype=wp.vec3f) + if self._terrain_levels_wp is None or self._terrain_levels_wp.ptr != self.terrain_levels.data_ptr(): + self._terrain_levels_wp = wp.from_torch(self.terrain_levels, dtype=wp.int64) + if self._terrain_types_wp is None or self._terrain_types_wp.ptr != self.terrain_types.data_ptr(): + self._terrain_types_wp = wp.from_torch(self.terrain_types, dtype=wp.int64) + if self._terrain_origins_wp is None or self._terrain_origins_wp.ptr != self.terrain_origins.data_ptr(): + self._terrain_origins_wp = wp.from_torch(self.terrain_origins, dtype=wp.vec3f) def _compute_env_origins_curriculum(self, num_envs: int, origins: torch.Tensor) -> torch.Tensor: """Compute the origins of the environments defined by the sub-terrains origins.""" @@ -464,10 +494,16 @@ def _compute_env_origins_curriculum(self, num_envs: int, origins: torch.Tensor) # store maximum terrain level possible self.max_terrain_level = num_rows # define all terrain levels and types available - self.terrain_levels = torch.randint(0, max_init_level + 1, (num_envs,), device=self.device) - self.terrain_types = torch.div( - torch.arange(num_envs, device=self.device), (num_envs / num_cols), rounding_mode="floor" - ).to(torch.long) + self.terrain_levels = self._assign_pointer_stable( + getattr(self, "terrain_levels", None), + torch.randint(0, max_init_level + 1, (num_envs,), device=self.device), + ) + self.terrain_types = self._assign_pointer_stable( + getattr(self, "terrain_types", None), + torch.div(torch.arange(num_envs, device=self.device), (num_envs / num_cols), rounding_mode="floor").to( + torch.long + ), + ) # create tensor based on number of environments env_origins = torch.zeros(num_envs, 3, device=self.device) env_origins[:] = origins[self.terrain_levels, self.terrain_types] diff --git a/source/isaaclab/test/scene/test_interactive_scene.py b/source/isaaclab/test/scene/test_interactive_scene.py index 909077bdbb39..c9b4010cc0e2 100644 --- a/source/isaaclab/test/scene/test_interactive_scene.py +++ b/source/isaaclab/test/scene/test_interactive_scene.py @@ -169,8 +169,8 @@ def test_reset_dispatches_warp_mask_to_supported_entities() -> None: entity.reset.assert_called_once_with(env_mask=env_mask) -def test_mask_reset_compacts_ids_for_surface_grippers() -> None: - """Mask reset should materialize IDs only for the legacy gripper boundary.""" +def test_mask_reset_forwards_boolean_mask_to_surface_grippers() -> None: + """Mask reset should hand surface grippers a boolean Torch view, not compact IDs.""" scene = object.__new__(InteractiveScene) gripper = Mock() scene._articulations = {} @@ -183,8 +183,10 @@ def test_mask_reset_compacts_ids_for_surface_grippers() -> None: scene.reset(env_mask=env_mask) - env_ids = gripper.reset.call_args.args[0] - assert torch.equal(env_ids, torch.tensor([0, 2], dtype=torch.int64)) + gripper.reset.assert_not_called() + gripper_env_mask = gripper.reset_mask.call_args.args[0] + assert gripper_env_mask.dtype == torch.bool + assert torch.equal(gripper_env_mask, torch.tensor([True, False, True])) def test_scene_publishes_plan_via_replicate(monkeypatch: pytest.MonkeyPatch): diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst index bfc3976ac1d3..1bc8f2464ff5 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -24,3 +24,5 @@ Fixed corrected center-of-mass sampling and termination metrics for partial resets. * Fixed captured velocity-command terms to read current root state on every replay instead of stale lazy-derived buffers. +* Fixed identity quaternion initialization for the Warp-native uniform pose + command term. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py index fde6fa6f0ed2..928e2ebcbd78 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py @@ -225,8 +225,11 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.torch_reset_time_outs: torch.Tensor = None self.torch_episode_length_buf: torch.Tensor = None - # Direct-task stages include Python reset hooks and stay eager until those - # complete boundaries are validated for capture. + # Direct-task stages stay eager. The end-pre stage resets through + # scene.reset(), whose legacy Torch actuator boundary materializes + # compact IDs on the host, and CUDA graph replay would silently skip + # that Python-side reset work. The follow-up launch-replay execution + # path restores capture for validated stages. self._warp_graph_cache = WarpGraphCache(enabled=False, device=self.device) # setup the action and observation spaces for Gym diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py index df27295ea933..5e77a8ee8ad8 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py @@ -424,7 +424,8 @@ def reset( if self._has_recorders: recorder_env_ids = env_ids if recorder_env_ids is None or isinstance(recorder_env_ids, wp.array): - recorder_env_ids = wp.to_torch(reset_mask).nonzero(as_tuple=False).squeeze(-1) + torch_mask = wp.to_torch(reset_mask) + recorder_env_ids = torch_mask.nonzero(as_tuple=False).squeeze(-1) # mask-boundary: recorders self.recorder_manager.record_pre_reset(recorder_env_ids) # set the seed diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index 04567703e618..789435f3309c 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -453,21 +453,32 @@ def _reset_idx( self, *, env_mask: wp.array(dtype=wp.bool), + env_ids: torch.Tensor | None = None, ) -> None: """Reset Warp-owned RL state for selected environments. Args: env_mask: Boolean Warp mask selecting environments to reset. + env_ids: Compact environment IDs matching :paramref:`env_mask`, when a + host consumer (e.g. the recorder boundary) already materialized them. """ if env_mask is not self.reset_mask_wp: wp.copy(self.reset_mask_wp, env_mask) env_mask = self.reset_mask_wp + # Legacy curriculum terms consume compact IDs. Materialize them at most + # once per reset and share the recorder boundary's IDs when available. + curriculum_kwargs: dict[str, Any] = {"env_mask": env_mask} + if self.curriculum_manager.requires_host_ids: + if env_ids is None: + env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) # mask-boundary: legacy terms + curriculum_kwargs["env_ids"] = env_ids + self._warp_graph_cache.call( "CurriculumManager_compute", self.curriculum_manager.compute, - env_mask=env_mask, timer=DEBUG_TIMER_RESET, + **curriculum_kwargs, ) with Timer(name="Scene_reset", msg="Scene reset took:", enable=DEBUG_TIMER_RESET, time_unit="us"): @@ -509,8 +520,8 @@ def _reset_idx( curriculum_info = self._warp_graph_cache.call( "CurriculumManager_reset", self.curriculum_manager.reset, - env_mask=env_mask, timer=DEBUG_TIMER_RESET, + **curriculum_kwargs, ) # -- command + event + termination managers @@ -568,6 +579,7 @@ def _reset_terminated_envs(self) -> None: if self.cfg.compute_final_obs: self.extras["final_obs"] = self.observation_manager.compute() + recorder_env_ids = None if self._has_recorders: with Timer( name="reset_selection_host", @@ -575,7 +587,7 @@ def _reset_terminated_envs(self) -> None: enable=DEBUG_TIMER_STEP, time_unit="us", ): - recorder_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) + recorder_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) # mask-boundary: recorders self.recorder_manager.record_pre_reset(recorder_env_ids) with Timer( @@ -584,7 +596,7 @@ def _reset_terminated_envs(self) -> None: enable=DEBUG_TIMER_STEP, time_unit="us", ): - self._reset_idx(env_mask=reset_mask) + self._reset_idx(env_mask=reset_mask, env_ids=recorder_env_ids) if self._has_recorders: self.extras["log"].update(self.recorder_manager.reset(recorder_env_ids)) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py index 5f94d157c440..9d1868407ccd 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py @@ -91,17 +91,20 @@ def reset_extras(self) -> dict[str, torch.Tensor]: def reset( self, env_mask: wp.array(dtype=wp.bool), + env_ids: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Reset selected class terms and return persistent curriculum logging outputs. Args: env_mask: Boolean Warp mask selecting environments to reset. + env_ids: Precomputed compact environment IDs matching :paramref:`env_mask`. + When omitted, legacy terms materialize IDs once at this boundary. Returns: Persistent scalar curriculum states keyed by their logging paths. """ env_mask = self._resolve_reset_mask(None, env_mask) - compact_env_ids = None + compact_env_ids = env_ids for term_cfg, mode in zip(self._term_cfgs, self._term_modes): if isinstance(term_cfg.func, ManagerTermBase): term_cfg.func.reset(env_mask=env_mask) @@ -128,15 +131,18 @@ def requires_host_boundary(self) -> bool: def compute( self, env_mask: wp.array(dtype=wp.bool), + env_ids: torch.Tensor | None = None, ) -> None: """Update curriculum terms for selected environments. Args: env_mask: Boolean Warp mask selecting environments to update. + env_ids: Precomputed compact environment IDs matching :paramref:`env_mask`. + When omitted, legacy terms materialize IDs once at this boundary. """ env_mask = self._resolve_reset_mask(None, env_mask) self._term_states_wp.zero_() - compact_env_ids = None + compact_env_ids = env_ids for term_idx, (term_cfg, mode) in enumerate(zip(self._term_cfgs, self._term_modes)): if mode == "mask": term_cfg.func(self._env, env_mask, term_cfg.out, **term_cfg.params) @@ -240,4 +246,4 @@ def _resolve_legacy_term_cfg(self, term_name: str, term_cfg: StableCurriculumTer @staticmethod def _compact_legacy_env_ids(env_mask: wp.array(dtype=wp.bool)) -> torch.Tensor: """Materialize compact Torch IDs at the legacy curriculum boundary.""" - return wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) + return wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) # mask-boundary: legacy curriculum terms diff --git a/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py index 87a6e4c2b50d..e3855a981cc5 100644 --- a/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py @@ -33,7 +33,7 @@ def test_reset_without_host_consumers_does_not_compact_ids() -> None: env._reset_terminated_envs() env.reset_buf.nonzero.assert_not_called() - env._reset_idx.assert_called_once_with(env_mask=reset_mask) + env._reset_idx.assert_called_once_with(env_mask=reset_mask, env_ids=None) def test_curriculum_uses_mask_before_scene_reset() -> None: @@ -54,7 +54,7 @@ def call(stage, fn, /, *args, **kwargs): reset=Mock(side_effect=lambda **kwargs: stages.append(("Scene_reset", kwargs))), surface_grippers={}, ) - env.curriculum_manager = SimpleNamespace(compute=Mock(), reset=Mock()) + env.curriculum_manager = SimpleNamespace(compute=Mock(), reset=Mock(), requires_host_ids=False) env.event_manager = SimpleNamespace(available_modes=[], reset=Mock()) env.observation_manager = SimpleNamespace(reset=Mock()) env.action_manager = SimpleNamespace(reset=Mock()) @@ -89,7 +89,7 @@ def call(stage, fn, /, *args, **kwargs): env.sim = SimpleNamespace(device="cpu") env.reset_mask_wp = wp.zeros(3, dtype=wp.bool, device="cpu") env.scene = SimpleNamespace(num_envs=3, reset=Mock(), surface_grippers={}) - env.curriculum_manager = SimpleNamespace(compute=Mock(), reset=Mock()) + env.curriculum_manager = SimpleNamespace(compute=Mock(), reset=Mock(), requires_host_ids=False) env.event_manager = SimpleNamespace(available_modes=[], reset=Mock()) env.observation_manager = SimpleNamespace(reset=Mock()) env.action_manager = SimpleNamespace(reset=Mock()) @@ -108,6 +108,37 @@ def call(stage, fn, /, *args, **kwargs): np.testing.assert_array_equal(seen[1][1], [False, False, False]) +def test_reset_threads_one_id_materialization_to_legacy_curriculum() -> None: + """Legacy curriculum consumers should share one compact-ID materialization per reset.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + stages: dict[str, dict] = {} + + def call(stage, fn, /, *args, **kwargs): + del fn, args + stages[stage] = kwargs + return None if stage.endswith("_compute") else {} + + env._warp_graph_cache = SimpleNamespace(call=call) + env.sim = SimpleNamespace(device="cpu") + env.reset_mask_wp = wp.zeros(3, dtype=wp.bool, device="cpu") + env.scene = SimpleNamespace(num_envs=3, reset=Mock(), surface_grippers={}) + env.curriculum_manager = SimpleNamespace(compute=Mock(), reset=Mock(), requires_host_ids=True) + env.event_manager = SimpleNamespace(available_modes=[], reset=Mock()) + env.observation_manager = SimpleNamespace(reset=Mock()) + env.action_manager = SimpleNamespace(reset=Mock()) + env.reward_manager = SimpleNamespace(reset=Mock()) + env.command_manager = SimpleNamespace(reset=Mock()) + env.termination_manager = SimpleNamespace(reset=Mock()) + env._episode_length_buf_wp = wp.ones(3, dtype=wp.int64, device="cpu") + env.extras = {} + + env._reset_idx(env_mask=wp.array([False, True, True], dtype=wp.bool, device="cpu")) + + compute_ids = stages["CurriculumManager_compute"]["env_ids"] + torch.testing.assert_close(compute_ids, torch.tensor([1, 2])) + assert stages["CurriculumManager_reset"]["env_ids"] is compute_ids + + def test_reset_compacts_ids_only_for_active_recorders() -> None: """Recorder callbacks should receive IDs while the Warp reset keeps its mask.""" env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) @@ -126,7 +157,7 @@ def test_reset_compacts_ids_only_for_active_recorders() -> None: reset_ids = env.recorder_manager.record_pre_reset.call_args.args[0] torch.testing.assert_close(reset_ids, torch.tensor([1])) - env._reset_idx.assert_called_once_with(env_mask=reset_mask) + env._reset_idx.assert_called_once_with(env_mask=reset_mask, env_ids=reset_ids) env.recorder_manager.reset.assert_called_once_with(reset_ids) env.recorder_manager.record_post_reset.assert_called_once_with(reset_ids) assert env.extras["log"]["recorder"] == 1.0 @@ -147,7 +178,7 @@ def test_rtx_rerender_does_not_require_compact_reset_ids() -> None: env._reset_terminated_envs() - env._reset_idx.assert_called_once_with(env_mask=reset_mask) + env._reset_idx.assert_called_once_with(env_mask=reset_mask, env_ids=None) assert env.sim.render.call_count == 2 @@ -185,4 +216,4 @@ def test_reset_captures_final_observation_before_reset() -> None: assert env.extras["final_obs"] is final_obs env.observation_manager.compute.assert_called_once_with() - env._reset_idx.assert_called_once_with(env_mask=reset_mask) + env._reset_idx.assert_called_once_with(env_mask=reset_mask, env_ids=None) diff --git a/source/isaaclab_experimental/test/managers/test_curriculum_manager.py b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py index 74a1b847ff92..532a0fadeb85 100644 --- a/source/isaaclab_experimental/test/managers/test_curriculum_manager.py +++ b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py @@ -154,6 +154,16 @@ def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): return self.scene.env_origins_wp +def _init_terrain_buffers(terrain: TerrainImporter) -> None: + """Mirror the buffer initializers of :meth:`TerrainImporter.__init__`.""" + terrain.terrain_origins = None + terrain.env_origins = None + terrain._terrain_levels_wp = None + terrain._terrain_types_wp = None + terrain._terrain_origins_wp = None + terrain._env_origins_wp = None + + def _make_terrain(levels: list[int]) -> TerrainImporter: num_envs = len(levels) terrain = TerrainImporter.__new__(TerrainImporter) @@ -163,6 +173,7 @@ def _make_terrain(levels: list[int]) -> TerrainImporter: max_init_terrain_level=2, terrain_generator=SimpleNamespace(size=(8.0, 8.0)), ) + _init_terrain_buffers(terrain) level_values = np.asarray(levels, dtype=np.int64) type_values = np.asarray([index % 2 for index in range(num_envs)], dtype=np.int64) origin_values = np.zeros((3, 2, 3), dtype=np.float32) @@ -206,6 +217,7 @@ def test_terrain_importer_grid_origins_cache_only_environment_view(): terrain = TerrainImporter.__new__(TerrainImporter) terrain.device = "cpu" terrain.cfg = SimpleNamespace(num_envs=4, env_spacing=2.0) + _init_terrain_buffers(terrain) terrain.configure_env_origins() assert terrain.terrain_origins is None @@ -218,6 +230,22 @@ def test_terrain_importer_grid_origins_cache_only_environment_view(): assert terrain._terrain_origins_wp is None +def test_terrain_importer_reconfigure_preserves_cached_views(): + """Reconfiguring origins must reuse buffer storage so cached views stay valid.""" + terrain = _make_terrain([0, 1, 2, 1]) + env_origins_wp = terrain.env_origins_wp + levels_wp = terrain.terrain_levels_wp + env_origins_ptr = terrain.env_origins.data_ptr() + shifted_origins = terrain.terrain_origins.clone() + 1.0 + + terrain.configure_env_origins(shifted_origins) + + assert terrain.env_origins_wp is env_origins_wp + assert terrain.terrain_levels_wp is levels_wp + assert terrain.env_origins.data_ptr() == env_origins_ptr + torch.testing.assert_close(terrain.terrain_origins, shifted_origins) + + def test_terrain_importer_mask_update_preserves_sparse_and_stable_semantics(): """Mask updates should handle up/down/clamp/wrap while stable IDs remain supported.""" terrain = _make_terrain([0, 1, 2, 2, 1, 0]) @@ -353,6 +381,34 @@ def test_curriculum_manager_compacts_ids_inside_legacy_term_boundary(): torch.testing.assert_close(extras["Curriculum/id_count"], torch.tensor(4.0)) +def test_curriculum_manager_consumes_threaded_ids_without_compaction(monkeypatch: pytest.MonkeyPatch): + """Precomputed compact IDs should reach legacy terms without another host materialization.""" + env = SimpleNamespace( + num_envs=4, + device="cpu", + sim=SimpleNamespace(is_playing=lambda: True), + ) + manager = CurriculumManager( + {"id_count": CurriculumTermCfg(func=_LegacyIdCurriculumTerm, params={"scale": 2.0})}, + env, + ) + env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + precomputed_ids = torch.tensor([0, 2]) + term = manager._term_cfgs[0].func + + def _fail_host_compaction(*args, **kwargs): + raise AssertionError("threaded IDs must not be re-materialized") + + with monkeypatch.context() as context: + context.setattr(torch.Tensor, "nonzero", _fail_host_compaction) + manager.compute(env_mask, env_ids=precomputed_ids) + extras = manager.reset(env_mask, env_ids=precomputed_ids) + + torch.testing.assert_close(term.compute_env_ids, precomputed_ids) + torch.testing.assert_close(term.reset_env_ids, precomputed_ids) + torch.testing.assert_close(extras["Curriculum/id_count"], torch.tensor(4.0)) + + def test_registered_reach_curricula_update_weights_without_a_host_boundary(): """Registered Reach reward schedules should update only on a device-selected reset.""" terrain = _make_terrain([0, 1, 2, 2]) diff --git a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py new file mode 100644 index 000000000000..cc62a94e0a33 --- /dev/null +++ b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py @@ -0,0 +1,39 @@ +# 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 + +"""Static gate for the mask-first contract. + +Host-side mask-to-ID compaction is a synchronization point and must only happen +at explicitly reviewed boundaries (recorders, legacy term/actuator APIs, +init-time tooling, diagnostics). Every such call site carries a same-line +``# mask-boundary: `` marker; this gate fails when a new unmarked site +appears in the scanned production packages. +""" + +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[4] +_SCAN_ROOTS = ( + "source/isaaclab_experimental/isaaclab_experimental", + "source/isaaclab_newton/isaaclab_newton", + "source/isaaclab/isaaclab/scene", + "source/isaaclab/isaaclab/terrains", +) +_COMPACTION_PATTERN = "nonzero(" +_BOUNDARY_MARKER = "mask-boundary:" + + +def test_host_compaction_sites_carry_mask_boundary_markers(): + """Every production mask-to-ID compaction must be a marked, reviewed boundary.""" + violations = [] + for scan_root in _SCAN_ROOTS: + for path in sorted((_REPO_ROOT / scan_root).rglob("*.py")): + for line_number, line in enumerate(path.read_text().splitlines(), start=1): + if _COMPACTION_PATTERN in line and _BOUNDARY_MARKER not in line: + violations.append(f"{path.relative_to(_REPO_ROOT)}:{line_number}: {line.strip()}") + assert not violations, ( + "Host mask-to-ID compaction found outside marked boundaries. Make the call mask-native, or" + " append '# mask-boundary: ' after review:\n" + "\n".join(violations) + ) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index ec66a907cc1a..76246171593a 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -270,15 +270,18 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None # use ellipses object to skip initial indices. if (env_ids is None) or (env_ids == slice(None)): env_ids = slice(None) - actuator_env_ids = env_ids - if self.actuators and env_mask is not None and not getattr(self, "_has_newton_actuators", False): - # Explicit boundary for the legacy Torch actuator fallback. The Warp - # frontend enables Newton-native actuators, whose state resets below - # consume the mask directly without materializing compact IDs. - actuator_env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) - # reset Lab actuators registered on this articulation - for actuator in self.actuators.values(): - actuator.reset(actuator_env_ids) + # Reset Lab actuators only outside the Newton-native mode. In native mode + # the per-env actuator state (delay queues, network hidden states) lives + # in the global adapter, which consumes the mask below, and the Lab + # actuator objects only carry configuration; resetting them here would + # touch unselected environments. + if not getattr(self, "_has_newton_actuators", False): + actuator_env_ids = env_ids + if self.actuators and env_mask is not None: + torch_mask = wp.to_torch(env_mask) + actuator_env_ids = torch_mask.nonzero(as_tuple=False).squeeze(-1) # mask-boundary: legacy actuators + for actuator in self.actuators.values(): + actuator.reset(actuator_env_ids) # reset the global Newton actuator adapter (its ``_states_a/_b`` buffers # carry per-env state — delay queues, neural hidden states — that must # be cleared for the resetting envs). The adapter spans the whole model, @@ -3963,7 +3966,7 @@ def _validate_cfg(self): default_joint_pos = self._data.default_joint_pos.torch[0] out_of_range = default_joint_pos < joint_pos_limits_lower out_of_range |= default_joint_pos > joint_pos_limits_upper - violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) + violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) # mask-boundary: diagnostics # throw error if any of the default joint positions are out of the limits if len(violated_indices) > 0: # prepare message for violated joints @@ -3980,7 +3983,7 @@ def _validate_cfg(self): joint_max_vel = self._data.joint_vel_limits.torch[0] default_joint_vel = self._data.default_joint_vel.torch[0] out_of_range = torch.abs(default_joint_vel) > joint_max_vel - violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) + violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) # mask-boundary: diagnostics if len(violated_indices) > 0: # prepare message for violated joints msg = "The following joints have default velocities out of the limits: \n" diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py index 168c3b9aa347..7f6c523eb391 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py @@ -266,8 +266,8 @@ def write_body_pose_to_sim_index( This method expects partial data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_poses: Body poses in simulation frame. Shape is (len(env_ids), len(body_ids), 7) @@ -297,8 +297,8 @@ def write_body_pose_to_sim_mask( This method expects full data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_poses: Body poses in simulation frame. Shape is (num_instances, num_bodies, 7) @@ -339,8 +339,8 @@ def write_body_velocity_to_sim_index( This method expects partial data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_velocities: Body velocities in simulation frame. @@ -374,8 +374,8 @@ def write_body_velocity_to_sim_mask( This method expects full data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_velocities: Body velocities in simulation frame. @@ -422,8 +422,8 @@ def write_body_link_pose_to_sim_index( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_poses: Body link poses in simulation frame. @@ -479,8 +479,8 @@ def write_body_link_pose_to_sim_mask( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_poses: Body link poses in simulation frame. Shape is (num_instances, num_bodies, 7) @@ -520,8 +520,8 @@ def write_body_com_pose_to_sim_index( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_poses: Body center of mass poses in simulation frame. @@ -580,8 +580,8 @@ def write_body_com_pose_to_sim_mask( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_poses: Body center of mass poses in simulation frame. Shape is (num_instances, num_bodies, 7) @@ -624,8 +624,8 @@ def write_body_com_velocity_to_sim_index( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_velocities: Body center of mass velocities in simulation frame. @@ -689,8 +689,8 @@ def write_body_com_velocity_to_sim_mask( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_velocities: Body center of mass velocities in simulation frame. @@ -738,8 +738,8 @@ def write_body_link_velocity_to_sim_index( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_velocities: Body link velocities in simulation frame. @@ -806,8 +806,8 @@ def write_body_link_velocity_to_sim_mask( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_velocities: Body link velocities in simulation frame. Shape is (num_instances, num_bodies, 6) @@ -848,8 +848,8 @@ def set_masses_index( This method expects partial data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: masses: Masses of all bodies. Shape is (len(env_ids), len(body_ids)). @@ -891,8 +891,8 @@ def set_masses_mask( This method expects full data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: masses: Masses of all bodies. Shape is (num_instances, num_bodies). @@ -935,8 +935,8 @@ def set_coms_index( This method expects partial data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. .. caution:: Unlike the PhysX version of this method, this method does not set the center of mass orientation. @@ -985,8 +985,8 @@ def set_coms_mask( This method expects full data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. .. caution:: Unlike the PhysX version of this method, this method does not set the center of mass orientation. @@ -1036,8 +1036,8 @@ def set_inertias_index( This method expects partial data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: inertias: Inertias of all bodies. Shape is (len(env_ids), len(body_ids), 9). @@ -1079,8 +1079,8 @@ def set_inertias_mask( This method expects full data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: inertias: Inertias of all bodies. Shape is (num_instances, num_bodies, 9). @@ -1277,12 +1277,14 @@ def _resolve_body_ids(self, body_ids) -> wp.array: return wp.from_torch(body_ids, dtype=wp.int32) return body_ids + # TODO: Replace the host-side compaction below with mask-native write kernels so the + # write_*_to_sim_mask methods become CUDA-graph capturable. def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array | torch.Tensor: """Resolve environment mask to indices via torch.nonzero.""" if env_mask is not None: if isinstance(env_mask, wp.array): env_mask = wp.to_torch(env_mask) - env_ids = torch.nonzero(env_mask)[:, 0].to(torch.int32) + env_ids = torch.nonzero(env_mask)[:, 0].to(torch.int32) # mask-boundary: legacy index write path else: env_ids = self._ALL_ENV_INDICES return env_ids @@ -1292,7 +1294,7 @@ def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array | torch.Ten if body_mask is not None: if isinstance(body_mask, wp.array): body_mask = wp.to_torch(body_mask) - body_ids = torch.nonzero(body_mask)[:, 0].to(torch.int32) + body_ids = torch.nonzero(body_mask)[:, 0].to(torch.int32) # mask-boundary: legacy index write path else: body_ids = self._ALL_BODY_INDICES return body_ids diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py index 586746744398..e411f7550eef 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py @@ -177,7 +177,7 @@ def replicate_builder_mapping( site_idx = builder.add_site(body=-1, xform=wp.transform_multiply(world_xform, xform), label=label) local_site_map.setdefault(label, [[] for _ in range(num_worlds)])[col].append(site_idx) - for row in torch.nonzero(mapping[:, col], as_tuple=True)[0].tolist(): + for row in torch.nonzero(mapping[:, col], as_tuple=True)[0].tolist(): # mask-boundary: init-time cloning source_builder = source_builders[sources[int(row)]] offset = builder.shape_count source_col = int(source_world_indices[int(row)]) @@ -222,7 +222,7 @@ def rename_builder_labels( for source_index, source in enumerate(sources): source_root = source.rstrip("/") - world_cols = torch.nonzero(mapping[source_index], as_tuple=True)[0].tolist() + world_cols = torch.nonzero(mapping[source_index], as_tuple=True)[0].tolist() # mask-boundary: init-time cloning world_roots = {int(env_ids[col]): destinations[source_index].format(int(env_ids[col])) for col in world_cols} def _rename_pair(values, worlds, *, collect_body_bindings: bool = False): diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py index 5a91f113111a..715a5197b4ec 100644 --- a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py @@ -335,7 +335,7 @@ def _validate_matching_root_orientations(self) -> None: # q and -q represent the same orientation, so compare absolute dot products. same_orientation = torch.abs(torch.sum(root_quat_w * root_quat_w[0:1], dim=-1)) > 1.0 - 1e-5 if not torch.all(same_orientation): - bad_env_ids = torch.nonzero(~same_orientation, as_tuple=False).flatten().tolist() + bad_env_ids = torch.nonzero(~same_orientation, as_tuple=False).flatten().tolist() # mask-boundary: errors raise RuntimeError( "NewtonInverseKinematicsAction solves against the env 0 prototype root orientation, but " f"root orientations differ in env ids {bad_env_ids}. Use identical fixed-base root orientations " diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py b/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py index 3eb06e007201..3a82e8540738 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py @@ -219,7 +219,8 @@ def _debug_vis_callback(self, event): up_axis = UsdGeom.GetStageUpAxis(self.stage) pos_w_torch = self._data.pos_w.torch accel_w = math_utils.quat_apply(self._data.quat_w.torch, self._data.lin_acc_b.torch) - valid_indices = (torch.linalg.norm(accel_w, dim=-1) > 1e-5).nonzero(as_tuple=True)[0] + accel_valid = torch.linalg.norm(accel_w, dim=-1) > 1e-5 + valid_indices = accel_valid.nonzero(as_tuple=True)[0] # mask-boundary: debug vis if valid_indices.numel() == 0: return pos_filtered = pos_w_torch.index_select(0, valid_indices) diff --git a/source/isaaclab_newton/test/assets/test_newton_mask_reset_forwarding.py b/source/isaaclab_newton/test/assets/test_newton_mask_reset_forwarding.py index 414234293114..c01f89c64162 100644 --- a/source/isaaclab_newton/test/assets/test_newton_mask_reset_forwarding.py +++ b/source/isaaclab_newton/test/assets/test_newton_mask_reset_forwarding.py @@ -5,7 +5,7 @@ """Tests for Warp reset-mask forwarding to Newton asset wrench storage.""" -from unittest.mock import Mock +from unittest.mock import Mock, patch import warp as wp from isaaclab_newton.assets.articulation.articulation import Articulation @@ -61,3 +61,26 @@ def test_articulation_without_lab_actuators_keeps_mask_on_device() -> None: asset.reset(env_mask=env_mask) asset._instantaneous_wrench_composer.reset.assert_called_once_with(slice(None), env_mask) + + +def test_articulation_native_mode_masked_reset_skips_lab_actuators() -> None: + """Native-mode partial resets must not reset Lab actuator state for all environments.""" + asset = Articulation.__new__(Articulation) + asset._initialize_handle = None + asset._invalidate_initialize_handle = None + asset._prim_deletion_handle = None + lab_actuator = Mock() + asset.actuators = {"legs": lab_actuator} + asset._has_newton_actuators = True + asset._instantaneous_wrench_composer = Mock() + asset._permanent_wrench_composer = Mock() + env_mask = wp.array([False, True], dtype=wp.bool, device="cpu") + adapter = Mock() + + from isaaclab_newton.physics import NewtonManager + + with patch.object(NewtonManager, "_adapter", adapter): + asset.reset(env_mask=env_mask) + + lab_actuator.reset.assert_not_called() + adapter.reset.assert_called_once_with(slice(None), env_mask=env_mask) diff --git a/source/isaaclab_tasks/changelog.d/jichuanh-warp-frontend-cleanup.rst b/source/isaaclab_tasks/changelog.d/jichuanh-warp-frontend-cleanup.rst deleted file mode 100644 index 5157824a2c21..000000000000 --- a/source/isaaclab_tasks/changelog.d/jichuanh-warp-frontend-cleanup.rst +++ /dev/null @@ -1,5 +0,0 @@ -Fixed -^^^^^ - -* Fixed DexSuite slab-clearance setup to reuse pointer-stable scene origin - storage. From 43b2ec15d27eecd0279a92c443c1210835dbe4fc Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 21 Jul 2026 03:15:31 -0700 Subject: [PATCH 45/58] Address review: capture-preserving stages and Warp-native naming Split the direct-task step so pure-kernel stages capture as on develop and only the reset stage stays eager, with its rationale registered at the cache. Rename Warp command buffers to own the plain names with _torch view suffixes, section their initialization, and fold the command debug visualization into shared core mixins. Add a mask option to reset_to, rename the manager reset orchestrator to _reset_mask, reuse the action buffer across manager reloads, and drop redundant scalar casts. Add captured-stage introspection, an environment-variable force-eager toggle, curriculum manager section comments, and the mask-first boundary conventions to AGENTS.md. --- AGENTS.md | 19 ++++++ .../jichuanh-warp-frontend-cleanup.minor.rst | 6 ++ .../isaaclab}/envs/mdp/commands/_debug_vis.py | 63 ++++++++++++++---- .../envs/mdp/commands/pose_command.py | 33 +--------- .../envs/mdp/commands/velocity_command.py | 62 +---------------- .../jichuanh-warp-frontend-cleanup.minor.rst | 3 + .../envs/direct_rl_env_warp.py | 28 ++++---- .../envs/manager_based_env_warp.py | 29 +++++--- .../envs/manager_based_rl_env_warp.py | 14 ++-- .../envs/mdp/commands/pose_command.py | 39 ++++++----- .../envs/mdp/commands/velocity_command.py | 66 ++++++++++--------- .../managers/curriculum_manager.py | 8 +++ .../utils/warp_graph_cache.py | 21 +++++- .../test/envs/mdp/commands/test_commands.py | 42 ++++++------ .../envs/test_manager_based_rl_env_warp.py | 28 ++++---- .../test/utils/test_warp_graph_cache.py | 35 +++++++++- 16 files changed, 285 insertions(+), 211 deletions(-) rename source/{isaaclab_experimental/isaaclab_experimental => isaaclab/isaaclab}/envs/mdp/commands/_debug_vis.py (53%) diff --git a/AGENTS.md b/AGENTS.md index 427e0bc9b0b5..497448ae4026 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -322,3 +322,22 @@ To debug Warp kernel behavior: 1. **Write a standalone reproduction script** and run it directly with `./isaaclab.sh -p -c "..."` or `./isaaclab.sh -p script.py`. This keeps stdout visible and avoids the test framework entirely. 2. **Use high-precision format strings** for floating-point debugging (e.g., `wp.printf("val=%.15e\n", x)`) — the default `%f` format hides values smaller than ~1e-6 that can still affect control flow. 3. **Remove all `wp.printf` calls before committing.** + +### Warp mask-first boundaries + +Boolean environment masks are the canonical selection in Warp-first code; host-side +mask-to-ID compaction (`nonzero()` on device data) is a GPU-to-host synchronization +with a data-dependent shape, which both stalls the pipeline and permanently +disqualifies the containing code path from CUDA graph capture. + +1. **Compaction only at reviewed boundaries.** Any production `nonzero(` in the + gated source trees must carry a same-line `# mask-boundary: ` marker; + the static gate `source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py` + fails otherwise. Legitimate boundaries are host consumers: recorders, legacy + compatibility fallbacks, init-time tooling, diagnostics, and debug visualization. +2. **Mask-native Torch means elementwise masked ops.** Use `masked_fill_` or + `torch.where`; boolean advanced indexing (`buf[:, mask] = 0`) gathers indices + with a data-dependent allocation and is not capture-safe. +3. **Dual reset APIs.** Keep the ID-based `reset(env_ids)` signature untouched and + add a sibling `reset_mask(env_mask)`; the base-class fallback compacts (and is + the marked boundary), mask-native implementations override it. diff --git a/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst index 32fbd10d6748..6e203e14eb99 100644 --- a/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst +++ b/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -7,3 +7,9 @@ Added access and mask-based curriculum updates. * Added :attr:`~isaaclab.scene.InteractiveScene.env_origins_wp` and a boolean ``env_mask`` option to :meth:`~isaaclab.scene.InteractiveScene.reset` for mask-based scene resets. + +Changed +^^^^^^^ + +* Changed the uniform pose and velocity command debug visualization to shared private mixins + reused by the Warp-native command terms. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/_debug_vis.py b/source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py similarity index 53% rename from source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/_debug_vis.py rename to source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py index e58e4a42e8ba..6e6c6301f101 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/_debug_vis.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Private debug-visualization helpers for Warp-native command terms.""" +"""Private debug-visualization mixins for command terms.""" from __future__ import annotations @@ -14,60 +14,99 @@ class _VelocityCommandDebugVis: - """Debug visualization for Warp-native velocity commands.""" + """Debug-visualization mixin shared by uniform velocity command terms.""" def _set_debug_vis_impl(self, debug_vis: bool): + # set visibility of markers + # note: parent only deals with callbacks. not their visibility if debug_vis: + # create markers if necessary for the first time if not hasattr(self, "goal_vel_visualizer"): + # -- goal self.goal_vel_visualizer = VisualizationMarkers(self.cfg.goal_vel_visualizer_cfg) + # -- current self.current_vel_visualizer = VisualizationMarkers(self.cfg.current_vel_visualizer_cfg) + # set their visibility to true self.goal_vel_visualizer.set_visibility(True) self.current_vel_visualizer.set_visibility(True) - elif hasattr(self, "goal_vel_visualizer"): - self.goal_vel_visualizer.set_visibility(False) - self.current_vel_visualizer.set_visibility(False) + else: + if hasattr(self, "goal_vel_visualizer"): + self.goal_vel_visualizer.set_visibility(False) + self.current_vel_visualizer.set_visibility(False) def _debug_vis_callback(self, event): + # check if robot is initialized + # note: this is needed in-case the robot is de-initialized. we can't access the data if not self.robot.is_initialized: return + # get marker location + # -- base state base_pos_w = self.robot.data.root_pos_w.torch.clone() base_pos_w[:, 2] += 0.5 + # -- resolve the scales and quaternions vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2]) vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow( self.robot.data.root_lin_vel_b.torch[:, :2] ) + # display markers self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) self.current_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale) def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Convert an XY velocity [m/s] to arrow scale and orientation.""" + """Converts the XY base velocity command to arrow direction rotation.""" + # obtain default scale of the marker default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale + # arrow-scale arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1) arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0 + # arrow-direction heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0]) zeros = torch.zeros_like(heading_angle) arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) - arrow_quat = math_utils.quat_mul(self.robot.data.root_quat_w.torch, arrow_quat) + # convert everything back from base to world frame + base_quat_w = self.robot.data.root_quat_w.torch + arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) + return arrow_scale, arrow_quat class _PoseCommandDebugVis: - """Debug visualization for Warp-native pose commands.""" + """Debug-visualization mixin shared by uniform pose command terms.""" def _set_debug_vis_impl(self, debug_vis: bool): + # create markers if necessary for the first time if debug_vis: if not hasattr(self, "goal_pose_visualizer"): + # -- goal pose self.goal_pose_visualizer = VisualizationMarkers(self.cfg.goal_pose_visualizer_cfg) + # -- current body pose self.current_pose_visualizer = VisualizationMarkers(self.cfg.current_pose_visualizer_cfg) + # set their visibility to true self.goal_pose_visualizer.set_visibility(True) self.current_pose_visualizer.set_visibility(True) - elif hasattr(self, "goal_pose_visualizer"): - self.goal_pose_visualizer.set_visibility(False) - self.current_pose_visualizer.set_visibility(False) + else: + if hasattr(self, "goal_pose_visualizer"): + self.goal_pose_visualizer.set_visibility(False) + self.current_pose_visualizer.set_visibility(False) def _debug_vis_callback(self, event): + # check if robot is initialized + # note: this is needed in-case the robot is de-initialized. we can't access the data if not self.robot.is_initialized: return - self.goal_pose_visualizer.visualize(self.pose_command_w[:, :3], self.pose_command_w[:, 3:]) + # update the markers + # -- goal pose + pose_command_w = self._debug_pose_command_w() + self.goal_pose_visualizer.visualize(pose_command_w[:, :3], pose_command_w[:, 3:]) + # -- current body pose body_link_pose_w = self.robot.data.body_link_pose_w.torch[:, self.body_idx] self.current_pose_visualizer.visualize(body_link_pose_w[:, :3], body_link_pose_w[:, 3:7]) + + def _debug_pose_command_w(self) -> torch.Tensor: + """Return the world-frame pose command consumed by the goal marker. + + The value is the desired pose in the simulation world frame [m, m, m, quaternion xyzw], + shape ``(num_envs, 7)``. Warp-native terms override this to expose the Torch alias of + their pointer-stable Warp buffer. + """ + return self.pose_command_w diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py index d98a6b5e7f81..efecc3e2a1d0 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py @@ -14,17 +14,18 @@ from isaaclab.assets import Articulation from isaaclab.managers import CommandTerm -from isaaclab.markers import VisualizationMarkers from isaaclab.utils.leapp import POSE7_ELEMENT_NAMES from isaaclab.utils.math import combine_frame_transforms, compute_pose_error, quat_from_euler_xyz, quat_unique +from ._debug_vis import _PoseCommandDebugVis + if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv from .commands_cfg import UniformPoseCommandCfg -class UniformPoseCommand(CommandTerm): +class UniformPoseCommand(_PoseCommandDebugVis, CommandTerm): """Command generator for generating pose commands uniformly. The command generator generates poses by sampling positions uniformly within specified @@ -151,31 +152,3 @@ def _resample_command(self, env_ids: Sequence[int]): def _update_command(self): pass - - def _set_debug_vis_impl(self, debug_vis: bool): - # create markers if necessary for the first time - if debug_vis: - if not hasattr(self, "goal_pose_visualizer"): - # -- goal pose - self.goal_pose_visualizer = VisualizationMarkers(self.cfg.goal_pose_visualizer_cfg) - # -- current body pose - self.current_pose_visualizer = VisualizationMarkers(self.cfg.current_pose_visualizer_cfg) - # set their visibility to true - self.goal_pose_visualizer.set_visibility(True) - self.current_pose_visualizer.set_visibility(True) - else: - if hasattr(self, "goal_pose_visualizer"): - self.goal_pose_visualizer.set_visibility(False) - self.current_pose_visualizer.set_visibility(False) - - def _debug_vis_callback(self, event): - # check if robot is initialized - # note: this is needed in-case the robot is de-initialized. we can't access the data - if not self.robot.is_initialized: - return - # update the markers - # -- goal pose - self.goal_pose_visualizer.visualize(self.pose_command_w[:, :3], self.pose_command_w[:, 3:]) - # -- current body pose - body_link_pose_w = self.robot.data.body_link_pose_w.torch[:, self.body_idx] - self.current_pose_visualizer.visualize(body_link_pose_w[:, :3], body_link_pose_w[:, 3:7]) diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py index 930f663f65d3..31fad8296bc4 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py @@ -16,7 +16,8 @@ import isaaclab.utils.math as math_utils from isaaclab.assets import Articulation from isaaclab.managers import CommandTerm -from isaaclab.markers import VisualizationMarkers + +from ._debug_vis import _VelocityCommandDebugVis if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -27,7 +28,7 @@ logger = logging.getLogger(__name__) -class UniformVelocityCommand(CommandTerm): +class UniformVelocityCommand(_VelocityCommandDebugVis, CommandTerm): r"""Command generator that generates a velocity command in SE(2) from uniform distribution. The command comprises of a linear velocity in x and y direction and an angular velocity around @@ -197,63 +198,6 @@ def _update_command(self): standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten() self.vel_command_b[standing_env_ids, :] = 0.0 - def _set_debug_vis_impl(self, debug_vis: bool): - # set visibility of markers - # note: parent only deals with callbacks. not their visibility - if debug_vis: - # create markers if necessary for the first time - if not hasattr(self, "goal_vel_visualizer"): - # -- goal - self.goal_vel_visualizer = VisualizationMarkers(self.cfg.goal_vel_visualizer_cfg) - # -- current - self.current_vel_visualizer = VisualizationMarkers(self.cfg.current_vel_visualizer_cfg) - # set their visibility to true - self.goal_vel_visualizer.set_visibility(True) - self.current_vel_visualizer.set_visibility(True) - else: - if hasattr(self, "goal_vel_visualizer"): - self.goal_vel_visualizer.set_visibility(False) - self.current_vel_visualizer.set_visibility(False) - - def _debug_vis_callback(self, event): - # check if robot is initialized - # note: this is needed in-case the robot is de-initialized. we can't access the data - if not self.robot.is_initialized: - return - # get marker location - # -- base state - base_pos_w = self.robot.data.root_pos_w.torch.clone() - base_pos_w[:, 2] += 0.5 - # -- resolve the scales and quaternions - vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2]) - vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow( - self.robot.data.root_lin_vel_b.torch[:, :2] - ) - # display markers - self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) - self.current_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale) - - """ - Internal helpers. - """ - - def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Converts the XY base velocity command to arrow direction rotation.""" - # obtain default scale of the marker - default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale - # arrow-scale - arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1) - arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0 - # arrow-direction - heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0]) - zeros = torch.zeros_like(heading_angle) - arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) - # convert everything back from base to world frame - base_quat_w = self.robot.data.root_quat_w.torch - arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) - - return arrow_scale, arrow_quat - class NormalVelocityCommand(UniformVelocityCommand): """Command generator that generates a velocity command in SE(2) from a normal distribution. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst index 1bc8f2464ff5..f85cde838393 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -8,6 +8,9 @@ Added environment IDs only for legacy host consumers. * Added Warp-native terrain curricula backed by cached zero-copy Warp views of terrain and scene origins. +* Added a boolean ``env_mask`` option to :meth:`ManagerBasedEnvWarp.reset_to`, the + ``ISAACLAB_WARP_CAPTURE`` environment variable to force eager execution, and + :attr:`WarpGraphCache.captured_stages` introspection for capture-coverage checks. Changed ^^^^^^^ diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py index 928e2ebcbd78..565002bbcb03 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py @@ -225,12 +225,13 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.torch_reset_time_outs: torch.Tensor = None self.torch_episode_length_buf: torch.Tensor = None - # Direct-task stages stay eager. The end-pre stage resets through - # scene.reset(), whose legacy Torch actuator boundary materializes - # compact IDs on the host, and CUDA graph replay would silently skip - # that Python-side reset work. The follow-up launch-replay execution - # path restores capture for validated stages. - self._warp_graph_cache = WarpGraphCache(enabled=False, device=self.device) + # Pure-kernel task stages capture as on develop. Only the reset stage is + # registered eager: it dispatches scene.reset(), whose legacy Torch + # actuator boundary materializes compact IDs on the host, and CUDA graph + # replay would silently skip that Python-side reset work. The follow-up + # launch-replay execution path restores its speed correctly. + self._warp_graph_cache = WarpGraphCache(device=self.device) + self._warp_graph_cache.register_capturability("DirectReset", False) # setup the action and observation spaces for Gym self._configure_gym_env_spaces() @@ -413,7 +414,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: # set actions into buffers # simulate with Timer(name="apply_action", msg="Action processing step took:", enable=DEBUG_TIMERS): - self._warp_graph_cache.call("action", self.step_warp_action) + self._warp_graph_cache.call("DirectAction_step", self.step_warp_action) # Keep scene writes outside the task graph until scene, sensor, and # actuator capturability have been validated as one backend boundary. @@ -433,12 +434,16 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self.common_step_counter += 1 # total step (common for all envs) with Timer(name="end_pre_graph", msg="End pre-graph took:", enable=DEBUG_TIMERS): - self._warp_graph_cache.call("end_pre", self._step_warp_end_pre) + self._warp_graph_cache.call("DirectEndPre_compute", self._step_warp_end_pre) + # Reset stage: registered eager (see __init__) because its scene dispatch + # crosses the legacy actuator host boundary. + with Timer(name="reset_stage", msg="Reset stage took:", enable=DEBUG_TIMERS): + self._warp_graph_cache.call("DirectReset_apply", self._step_warp_reset) # Keep the post-reset scene write at the explicit backend boundary. with Timer(name="write_data_to_sim_post", msg="Write data to sim (post-reset) took:", enable=DEBUG_TIMERS): self.scene.write_data_to_sim() with Timer(name="end_post_graph", msg="End post-graph took:", enable=DEBUG_TIMERS): - self._warp_graph_cache.call("end_post", self._step_warp_end_post) + self._warp_graph_cache.call("DirectEndPost_step", self._step_warp_end_post) # Visualization hook — runs after CUDA graph scope. Override in subclass # to update markers or other non-graphable visual elements. @@ -469,7 +474,7 @@ def step_warp_action(self) -> None: # task path correct without making capture support a prerequisite. def _step_warp_end_pre(self) -> None: - """Capturable portion before write_data_to_sim (pure warp kernels).""" + """Capturable portion before the reset stage (pure warp kernels).""" wp.launch( increment_all_int32, dim=self.num_envs, @@ -481,7 +486,8 @@ def _step_warp_end_pre(self) -> None: self._get_dones() self._get_rewards() - # -- reset envs that terminated/timed-out and log the episode information + def _step_warp_reset(self) -> None: + """Reset envs that terminated/timed-out (eager: crosses host boundaries).""" self._reset_idx(mask=self.reset_buf) def _step_warp_end_post(self) -> None: diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py index 5e77a8ee8ad8..d14283016a1f 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py @@ -366,9 +366,14 @@ def load_managers(self): print("[INFO] Recorder Manager: ", self.recorder_manager) # -- action manager self.action_manager = ActionManager(self.cfg.actions, self) - self._action_in_wp = wp.zeros( - (self.num_envs, self.action_manager.total_action_dim), dtype=wp.float32, device=self.device - ) + # Manager-dependent persistent buffer: reuse the storage across manager + # reloads when the action dimension is unchanged so captured graphs keep + # reading the same pointer. + action_shape = (self.num_envs, self.action_manager.total_action_dim) + if getattr(self, "_action_in_wp", None) is None or tuple(self._action_in_wp.shape) != action_shape: + self._action_in_wp = wp.zeros(action_shape, dtype=wp.float32, device=self.device) + else: + self._action_in_wp.zero_() print("[INFO] Action Manager: ", self.action_manager) # -- observation manager self.observation_manager = ObservationManager(self.cfg.observations, self) @@ -401,7 +406,7 @@ def reset( ) -> tuple[VecEnvObs, dict]: """Resets the specified environments and returns observations. - This function calls the :meth:`_reset_idx` function to reset the specified environments. + This function calls the :meth:`_reset_mask` function to reset the specified environments. However, certain operations, such as procedural terrain generation, that happened during initialization are not repeated. @@ -441,7 +446,7 @@ def reset( device=self.device, ) - self._reset_idx(env_mask=reset_mask) + self._reset_mask(env_mask=reset_mask) if self._has_recorders: self.extras["log"].update(self.recorder_manager.reset(recorder_env_ids)) @@ -474,6 +479,7 @@ def reset_to( env_ids: Sequence[int] | None, seed: int | None = None, is_relative: bool = False, + env_mask: wp.array | None = None, ): """Resets specified environments to provided states. @@ -491,9 +497,14 @@ def reset_to( seed: The seed to use for randomization. Defaults to None, in which case the seed is not set. is_relative: If set to True, the state is considered relative to the environment origins. Defaults to False. + env_mask: Boolean Warp mask selecting environments. Takes precedence over + :paramref:`env_ids`. State setting and recorders are host consumers, so + compact IDs are still materialized once inside this method. """ # reset all envs in the scene if env_ids is None - if env_ids is None: + if env_mask is not None: + env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) # mask-boundary: host state dict + elif env_ids is None: env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device) if self._has_recorders: @@ -503,8 +514,8 @@ def reset_to( if seed is not None: self.seed(seed) - env_mask = self.resolve_env_mask(env_ids=env_ids) - self._reset_idx(env_mask=env_mask) + env_mask = self.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + self._reset_mask(env_mask=env_mask) if self._has_recorders: self.extras["log"].update(self.recorder_manager.reset(env_ids)) @@ -649,7 +660,7 @@ def close(self): Helper functions. """ - def _reset_idx( + def _reset_mask( self, *, env_mask: wp.array(dtype=wp.bool), diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index 789435f3309c..a9480e34823e 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -295,7 +295,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self.reward_buf = self._warp_graph_cache.call( "RewardManager_compute", self.reward_manager.compute, - dt=float(self.step_dt), + dt=self.step_dt, timer=DEBUG_TIMER_STEP, ) @@ -315,7 +315,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self._warp_graph_cache.call( "CommandManager_compute", self.command_manager.compute, - dt=float(self.step_dt), + dt=self.step_dt, timer=DEBUG_TIMER_STEP, ) @@ -325,7 +325,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: "EventManager_apply_interval", self.event_manager.apply, mode="interval", - dt=float(self.step_dt), + dt=self.step_dt, timer=DEBUG_TIMER_STEP, ) @@ -449,7 +449,7 @@ def _configure_gym_env_spaces(self): self.observation_space = gym.vector.utils.batch_space(self.single_observation_space, self.num_envs) self.action_space = gym.vector.utils.batch_space(self.single_action_space, self.num_envs) - def _reset_idx( + def _reset_mask( self, *, env_mask: wp.array(dtype=wp.bool), @@ -591,12 +591,12 @@ def _reset_terminated_envs(self) -> None: self.recorder_manager.record_pre_reset(recorder_env_ids) with Timer( - name="reset_idx", - msg="Reset idx took:", + name="reset_mask", + msg="Reset mask took:", enable=DEBUG_TIMER_STEP, time_unit="us", ): - self._reset_idx(env_mask=reset_mask, env_ids=recorder_env_ids) + self._reset_mask(env_mask=reset_mask, env_ids=recorder_env_ids) if self._has_recorders: self.extras["log"].update(self.recorder_manager.reset(recorder_env_ids)) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py index deaa5a5f569d..a0d4b90ff7de 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py @@ -13,12 +13,11 @@ import warp as wp from isaaclab.assets import Articulation +from isaaclab.envs.mdp.commands._debug_vis import _PoseCommandDebugVis from isaaclab.utils.leapp import POSE7_ELEMENT_NAMES from isaaclab_experimental.managers import CommandTerm -from ._debug_vis import _PoseCommandDebugVis - if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -121,25 +120,31 @@ def __init__(self, cfg: UniformPoseCommandCfg, env: ManagerBasedEnv): """ super().__init__(cfg, env) + # -- robot / body resolution self.robot: Articulation = env.scene[cfg.asset_name] self.body_idx = self.robot.find_bodies(cfg.body_name)[0][0] - self._pose_command_b_wp = wp.zeros((self.num_envs, 7), dtype=wp.float32, device=self.device) - self.pose_command_b = wp.to_torch(self._pose_command_b_wp) - self.pose_command_b[:, 6] = 1.0 - self._pose_command_w_wp = wp.zeros((self.num_envs, 7), dtype=wp.float32, device=self.device) - self.pose_command_w = wp.to_torch(self._pose_command_w_wp) + # -- command buffers (Warp-native, pointer-stable) + self.pose_command_b = wp.zeros((self.num_envs, 7), dtype=wp.float32, device=self.device) + self.pose_command_w = wp.zeros((self.num_envs, 7), dtype=wp.float32, device=self.device) + # -- metrics buffers self.metrics["position_error"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) self.metrics["orientation_error"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) - self._success_rate_wp = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._success_rate = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) self._track_success = cfg.position_success_threshold is not None if self._track_success: - self.metrics["success_rate"] = self._success_rate_wp + self.metrics["success_rate"] = self._success_rate + # adds (optional) cmd kind and element names for leapp export self.cfg.cmd_kind = self.cfg.cmd_kind or "command/body/pose" self.cfg.element_names = self.cfg.element_names or POSE7_ELEMENT_NAMES + # -- Torch views (zero-copy aliases for stable consumers) + self.pose_command_b_torch = wp.to_torch(self.pose_command_b) + self.pose_command_w_torch = wp.to_torch(self.pose_command_w) + self.pose_command_b_torch[:, 6] = 1.0 + def __str__(self) -> str: """Return a string representation of the command generator.""" msg = "UniformPoseCommand:\n" @@ -150,28 +155,28 @@ def __str__(self) -> str: @property def command(self) -> torch.Tensor: """Desired pose command [m, m, m, quaternion xyzw], shape ``(num_envs, 7)``.""" - return self.pose_command_b + return self.pose_command_b_torch @property def command_wp(self) -> wp.array(dtype=wp.float32, ndim=2): """Pointer-stable desired pose command [m, m, m, quaternion xyzw].""" - return self._pose_command_b_wp + return self.pose_command_b def _update_metrics(self): wp.launch( kernel=_update_pose_metrics, dim=self.num_envs, inputs=[ - self._pose_command_b_wp, + self.pose_command_b, self.robot.data.root_pos_w.warp, self.robot.data.root_quat_w.warp, self.robot.data.body_pos_w.warp, self.robot.data.body_quat_w.warp, self.body_idx, - self._pose_command_w_wp, + self.pose_command_w, self.metrics["position_error"], self.metrics["orientation_error"], - self._success_rate_wp, + self._success_rate, self._track_success, self.cfg.position_success_threshold or 0.0, ], @@ -185,7 +190,7 @@ def _resample_command(self, env_mask: wp.array): inputs=[ env_mask, self._env.rng_state_wp, - self._pose_command_b_wp, + self.pose_command_b, self.cfg.ranges.pos_x[0], self.cfg.ranges.pos_x[1], self.cfg.ranges.pos_y[0], @@ -205,3 +210,7 @@ def _resample_command(self, env_mask: wp.array): def _update_command(self): pass + + def _debug_pose_command_w(self) -> torch.Tensor: + """Return the Torch alias of the pointer-stable world-frame pose command buffer.""" + return self.pose_command_w_torch diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py index 24fbcefcdf16..7eb952fc6f5f 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py @@ -16,12 +16,11 @@ from isaaclab_newton.kernels.state_kernels import body_ang_vel_from_root, body_lin_vel_from_root from isaaclab.assets import Articulation +from isaaclab.envs.mdp.commands._debug_vis import _VelocityCommandDebugVis from isaaclab_experimental.managers import CommandTerm from isaaclab_experimental.utils.warp import wrap_to_pi -from ._debug_vis import _VelocityCommandDebugVis - if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -158,6 +157,7 @@ def __init__(self, cfg: UniformVelocityCommandCfg, env: ManagerBasedEnv): """ super().__init__(cfg, env) + # -- config validation if cfg.heading_command and cfg.ranges.heading is None: raise ValueError( "The velocity command has heading commands active (heading_command=True) but the `ranges.heading`" @@ -169,27 +169,33 @@ def __init__(self, cfg: UniformVelocityCommandCfg, env: ManagerBasedEnv): " but the heading command is not active. Consider setting the flag for the heading command to True." ) + # -- robot resolution self.robot: Articulation = env.scene[cfg.asset_name] - self._vel_command_b_wp = wp.zeros((self.num_envs, 3), dtype=wp.float32, device=self.device) - self.vel_command_b = wp.to_torch(self._vel_command_b_wp) - self._heading_target_wp = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) - self.heading_target = wp.to_torch(self._heading_target_wp) - self._is_heading_env_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) - self.is_heading_env = wp.to_torch(self._is_heading_env_wp) - self._is_standing_env_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) - self.is_standing_env = wp.to_torch(self._is_standing_env_wp) + # -- command buffers (Warp-native, pointer-stable) + self.vel_command_b = wp.zeros((self.num_envs, 3), dtype=wp.float32, device=self.device) + self.heading_target = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self.is_heading_env = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + self.is_standing_env = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + # -- metrics buffers self.metrics["error_vel_xy"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) self.metrics["error_vel_yaw"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) self.metrics["success_rate"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) - self._error_xy_sum_wp = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) - self._error_yaw_sum_wp = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) - self._step_count_wp = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._error_xy_sum = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._error_yaw_sum = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._step_count = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + # adds (optional) cmd kind and element names for leapp export self.cfg.cmd_kind = self.cfg.cmd_kind or "command/body/velocity" self.cfg.element_names = self.cfg.element_names or ["lin_vel_x", "lin_vel_y", "ang_vel_z"] + # -- Torch views (zero-copy aliases for stable consumers) + self.vel_command_b_torch = wp.to_torch(self.vel_command_b) + self.heading_target_torch = wp.to_torch(self.heading_target) + self.is_heading_env_torch = wp.to_torch(self.is_heading_env) + self.is_standing_env_torch = wp.to_torch(self.is_standing_env) + def __str__(self) -> str: """Return a string representation of the command generator.""" msg = "UniformVelocityCommand:\n" @@ -204,12 +210,12 @@ def __str__(self) -> str: @property def command(self) -> torch.Tensor: """Desired base velocity command [m/s, m/s, rad/s], shape ``(num_envs, 3)``.""" - return self.vel_command_b + return self.vel_command_b_torch @property def command_wp(self) -> wp.array(dtype=wp.float32, ndim=2): """Pointer-stable desired base velocity command [m/s, m/s, rad/s].""" - return self._vel_command_b_wp + return self.vel_command_b def reset( self, @@ -232,9 +238,9 @@ def reset( dim=self.num_envs, inputs=[ env_mask, - self._error_xy_sum_wp, - self._error_yaw_sum_wp, - self._step_count_wp, + self._error_xy_sum, + self._error_yaw_sum, + self._step_count, self.metrics["error_vel_xy"], self.metrics["error_vel_yaw"], self.metrics["success_rate"], @@ -250,12 +256,12 @@ def _update_metrics(self): kernel=_accumulate_velocity_metrics, dim=self.num_envs, inputs=[ - self._vel_command_b_wp, + self.vel_command_b, self.robot.data.root_pose_w.warp, self.robot.data.root_vel_w.warp, - self._error_xy_sum_wp, - self._error_yaw_sum_wp, - self._step_count_wp, + self._error_xy_sum, + self._error_yaw_sum, + self._step_count, ], device=self.device, ) @@ -268,10 +274,10 @@ def _resample_command(self, env_mask: wp.array): inputs=[ env_mask, self._env.rng_state_wp, - self._vel_command_b_wp, - self._heading_target_wp, - self._is_heading_env_wp, - self._is_standing_env_wp, + self.vel_command_b, + self.heading_target, + self.is_heading_env, + self.is_standing_env, self.cfg.ranges.lin_vel_x[0], self.cfg.ranges.lin_vel_x[1], self.cfg.ranges.lin_vel_y[0], @@ -292,10 +298,10 @@ def _update_command(self): kernel=_update_velocity_command, dim=self.num_envs, inputs=[ - self._vel_command_b_wp, - self._heading_target_wp, - self._is_heading_env_wp, - self._is_standing_env_wp, + self.vel_command_b, + self.heading_target, + self.is_heading_env, + self.is_standing_env, self.robot.data.root_pose_w.warp, self.cfg.heading_command, self.cfg.heading_control_stiffness, diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py index 9d1868407ccd..950496451c4f 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py @@ -104,6 +104,9 @@ def reset( Persistent scalar curriculum states keyed by their logging paths. """ env_mask = self._resolve_reset_mask(None, env_mask) + # Dispatch by term kind: Warp class terms consume the mask directly; stable + # class terms receive compact IDs (materialized at most once) or a global + # selection, depending on their declared mode. compact_env_ids = env_ids for term_cfg, mode in zip(self._term_cfgs, self._term_modes): if isinstance(term_cfg.func, ManagerTermBase): @@ -141,8 +144,13 @@ def compute( When omitted, legacy terms materialize IDs once at this boundary. """ env_mask = self._resolve_reset_mask(None, env_mask) + # Logging states are recomputed every call; mask-native terms write their + # scalar slot on-device through the pointer-stable ``term_cfg.out`` view. self._term_states_wp.zero_() compact_env_ids = env_ids + # Term modes: "mask" = Warp-native ``(env, env_mask, out)``; + # "legacy_ids" = stable ``(env, env_ids)``; "legacy_global" = stable term + # that ignores its ID argument (``requires_host_ids=False``). for term_idx, (term_cfg, mode) in enumerate(zip(self._term_cfgs, self._term_modes)): if mode == "mask": term_cfg.func(self._env, env_mask, term_cfg.out, **term_cfg.params) diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py index 437bd2fe1031..ee54c0e59121 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py @@ -5,6 +5,7 @@ """Warp CUDA graph capture-or-replay utility.""" +import os from collections.abc import Callable from typing import Any @@ -12,6 +13,9 @@ from isaaclab.utils.timer import Timer +CAPTURE_ENV_VAR = "ISAACLAB_WARP_CAPTURE" +"""Set ``ISAACLAB_WARP_CAPTURE=0`` to force every stage eager (debug / A-B validation).""" + class WarpGraphCache: """Execute Warp stages eagerly or through cached CUDA graphs. @@ -38,16 +42,27 @@ def __init__(self, *, enabled: bool = True, device: wp.DeviceLike = None): """Initialize the cache. Args: - enabled: Whether eligible stages may use CUDA graph capture. + enabled: Whether eligible stages may use CUDA graph capture. The + :data:`CAPTURE_ENV_VAR` environment variable can force this off. device: Warp device used by cached stages. CPU devices always run eagerly. """ - self._enabled = enabled + self._enabled = enabled and os.environ.get(CAPTURE_ENV_VAR, "1") != "0" self._device = wp.get_device(device) if device is not None else None self._graphs: dict[str, Any] = {} self._results: dict[str, Any] = {} self._capturable: dict[str, bool] = {} self._warmed: set[str] = set() + @property + def captured_stages(self) -> tuple[str, ...]: + """Stages currently backed by a captured CUDA graph, sorted by name.""" + return tuple(sorted(self._graphs)) + + @property + def eager_groups(self) -> tuple[str, ...]: + """Stage groups registered as non-capturable, sorted by name.""" + return tuple(sorted(group for group, capturable in self._capturable.items() if not capturable)) + def call( self, stage: str, @@ -140,6 +155,8 @@ def _capture( self._graphs[stage] = capture.graph self._results[stage] = result self._warmed.add(stage) + # One grep-able line per stage so runs can assert capture coverage. + print(f"[INFO] WarpGraphCache: captured stage '{stage}'") return result def register_capturability(self, key: str, capturable: bool) -> None: diff --git a/source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py b/source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py index c07980f0b6c3..9d5bef7edfa4 100644 --- a/source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py +++ b/source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py @@ -124,8 +124,8 @@ def test_configs_exports_and_command_accessors_preserve_shared_contracts(): assert experimental_mdp.UniformPoseCommandCfg is UniformPoseCommandCfg assert str(_velocity_cfg().class_type).endswith(".commands.velocity_command:UniformVelocityCommand") assert str(_pose_cfg().class_type).endswith(".commands.pose_command:UniformPoseCommand") - assert UniformVelocityCommand._set_debug_vis_impl.__module__.endswith(".commands._debug_vis") - assert UniformPoseCommand._set_debug_vis_impl.__module__.endswith(".commands._debug_vis") + assert UniformVelocityCommand._set_debug_vis_impl.__module__ == "isaaclab.envs.mdp.commands._debug_vis" + assert UniformPoseCommand._set_debug_vis_impl.__module__ == "isaaclab.envs.mdp.commands._debug_vis" env = _Env() velocity_term = UniformVelocityCommand(_velocity_cfg(), env) @@ -180,15 +180,15 @@ def test_uniform_velocity_command_updates_and_resets_selected_rows(): term._update_metrics() wp.synchronize() torch.testing.assert_close( - wp.to_torch(term._error_xy_sum_wp), + wp.to_torch(term._error_xy_sum), torch.tensor([np.sqrt(5.0), 1.0, 0.5, 1.0], dtype=torch.float32), ) - torch.testing.assert_close(wp.to_torch(term._error_yaw_sum_wp), torch.tensor([0.2, 0.1, 0.3, 0.0])) + torch.testing.assert_close(wp.to_torch(term._error_yaw_sum), torch.tensor([0.2, 0.1, 0.3, 0.0])) term.command.fill_(-9.0) - wp.to_torch(term._error_xy_sum_wp)[:] = torch.tensor([0.2, 9.0, 2.0, 8.0]) - wp.to_torch(term._error_yaw_sum_wp)[:] = torch.tensor([0.1, 9.0, 2.0, 8.0]) - wp.to_torch(term._step_count_wp)[:] = torch.tensor([1.0, 1.0, 2.0, 1.0]) + wp.to_torch(term._error_xy_sum)[:] = torch.tensor([0.2, 9.0, 2.0, 8.0]) + wp.to_torch(term._error_yaw_sum)[:] = torch.tensor([0.1, 9.0, 2.0, 8.0]) + wp.to_torch(term._step_count)[:] = torch.tensor([1.0, 1.0, 2.0, 1.0]) wp.to_torch(term.time_left_wp)[:] = torch.tensor([-1.0, 6.0, -1.0, 8.0]) wp.to_torch(term.command_counter_wp)[:] = torch.tensor([5, 6, 7, 8], dtype=torch.int32) env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") @@ -205,13 +205,13 @@ def test_uniform_velocity_command_updates_and_resets_selected_rows(): torch.testing.assert_close(extras["error_vel_xy"], torch.tensor(0.6)) torch.testing.assert_close(extras["error_vel_yaw"], torch.tensor(0.55)) torch.testing.assert_close(extras["success_rate"], torch.tensor(0.5)) - torch.testing.assert_close(wp.to_torch(term._error_xy_sum_wp), torch.tensor([0.0, 9.0, 0.0, 8.0])) - torch.testing.assert_close(wp.to_torch(term._step_count_wp), torch.tensor([0.0, 1.0, 0.0, 1.0])) + torch.testing.assert_close(wp.to_torch(term._error_xy_sum), torch.tensor([0.0, 9.0, 0.0, 8.0])) + torch.testing.assert_close(wp.to_torch(term._step_count), torch.tensor([0.0, 1.0, 0.0, 1.0])) term.command[:] = torch.tensor([[1.0, 2.0, 0.4]]).repeat(4, 1) - term.heading_target[:] = torch.tensor([np.pi / 2, 0.0, np.pi / 2, 0.0]) - term.is_heading_env[:] = torch.tensor([True, False, True, False]) - term.is_standing_env[:] = torch.tensor([False, False, True, True]) + term.heading_target_torch[:] = torch.tensor([np.pi / 2, 0.0, np.pi / 2, 0.0]) + term.is_heading_env_torch[:] = torch.tensor([True, False, True, False]) + term.is_standing_env_torch[:] = torch.tensor([False, False, True, True]) term._update_command() wp.synchronize() torch.testing.assert_close( @@ -230,9 +230,9 @@ def test_uniform_velocity_command_replay_reads_current_root_state(): term.command.zero_() term.command[:, 0] = 1.0 - term.heading_target.fill_(np.pi / 2) - term.is_heading_env.fill_(True) - term.is_standing_env.fill_(False) + term.heading_target_torch.fill_(np.pi / 2) + term.is_heading_env_torch.fill_(True) + term.is_standing_env_torch.fill_(False) with wp.ScopedCapture(device=env.device) as metrics_capture: term._update_metrics() @@ -243,7 +243,7 @@ def test_uniform_velocity_command_replay_reads_current_root_state(): robot_data.root_vel_w.torch[:, 0] = 1.0 wp.capture_launch(metrics_capture.graph) wp.synchronize() - torch.testing.assert_close(wp.to_torch(term._error_xy_sum_wp), torch.zeros(env.num_envs)) + torch.testing.assert_close(wp.to_torch(term._error_xy_sum), torch.zeros(env.num_envs)) # A 90-degree root yaw exactly matches the heading target, so replay should command zero yaw rate. half_sqrt = np.sqrt(0.5) @@ -309,8 +309,8 @@ def test_uniform_pose_command_matches_stable_math_and_tracks_sticky_success(): env.scene["robot"].data.body_pos_w.torch[:, 0], env.scene["robot"].data.body_quat_w.torch[:, 0], ) - torch.testing.assert_close(term.pose_command_w[:, :3], expected_position_w, atol=1.0e-5, rtol=1.0e-5) - torch.testing.assert_close(term.pose_command_w[:, 3:], expected_orientation_w, atol=1.0e-5, rtol=1.0e-5) + torch.testing.assert_close(term.pose_command_w_torch[:, :3], expected_position_w, atol=1.0e-5, rtol=1.0e-5) + torch.testing.assert_close(term.pose_command_w_torch[:, 3:], expected_orientation_w, atol=1.0e-5, rtol=1.0e-5) torch.testing.assert_close( wp.to_torch(term.metrics["position_error"]), torch.linalg.norm(expected_position_delta, dim=-1) ) @@ -320,15 +320,15 @@ def test_uniform_pose_command_matches_stable_math_and_tracks_sticky_success(): atol=1.0e-5, rtol=1.0e-5, ) - torch.testing.assert_close(wp.to_torch(term._success_rate_wp), torch.tensor([1.0, 0.0, 1.0, 0.0])) + torch.testing.assert_close(wp.to_torch(term._success_rate), torch.tensor([1.0, 0.0, 1.0, 0.0])) env.scene["robot"].data.body_pos_w.torch[0, 0] += 2.0 term._update_metrics() wp.synchronize() - torch.testing.assert_close(wp.to_torch(term._success_rate_wp), torch.tensor([1.0, 0.0, 1.0, 0.0])) + torch.testing.assert_close(wp.to_torch(term._success_rate), torch.tensor([1.0, 0.0, 1.0, 0.0])) reset_mask = wp.array([True, False, False, True], dtype=wp.bool, device="cpu") extras = term.reset(env_mask=reset_mask) wp.synchronize() torch.testing.assert_close(extras["success_rate"], torch.tensor(0.5)) - torch.testing.assert_close(wp.to_torch(term._success_rate_wp), torch.tensor([0.0, 0.0, 1.0, 0.0])) + torch.testing.assert_close(wp.to_torch(term._success_rate), torch.tensor([0.0, 0.0, 1.0, 0.0])) diff --git a/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py index e3855a981cc5..93ca80750ec4 100644 --- a/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py @@ -24,7 +24,7 @@ def test_reset_without_host_consumers_does_not_compact_ids() -> None: env.reset_buf = Mock() env.reset_buf.any.return_value.item.return_value = True env.reset_buf.nonzero.side_effect = AssertionError("reset IDs should not be materialized") - env._reset_idx = Mock() + env._reset_mask = Mock() env._has_recorders = False env.has_rtx_sensors = False env.cfg = SimpleNamespace(compute_final_obs=False, num_rerenders_on_reset=0) @@ -33,7 +33,7 @@ def test_reset_without_host_consumers_does_not_compact_ids() -> None: env._reset_terminated_envs() env.reset_buf.nonzero.assert_not_called() - env._reset_idx.assert_called_once_with(env_mask=reset_mask, env_ids=None) + env._reset_mask.assert_called_once_with(env_mask=reset_mask, env_ids=None) def test_curriculum_uses_mask_before_scene_reset() -> None: @@ -64,7 +64,7 @@ def call(stage, fn, /, *args, **kwargs): env._episode_length_buf_wp = wp.ones(3, dtype=wp.int64, device="cpu") env.extras = {} reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") - env._reset_idx(env_mask=reset_mask) + env._reset_mask(env_mask=reset_mask) assert [stage for stage, _ in stages[:2]] == ["CurriculumManager_compute", "Scene_reset"] assert stages[0][1]["env_mask"] is env.reset_mask_wp @@ -99,8 +99,8 @@ def call(stage, fn, /, *args, **kwargs): env._episode_length_buf_wp = wp.ones(3, dtype=wp.int64, device="cpu") env.extras = {} - env._reset_idx(env_mask=wp.array([False, True, False], dtype=wp.bool, device="cpu")) - env._reset_idx(env_mask=wp.zeros(3, dtype=wp.bool, device="cpu")) + env._reset_mask(env_mask=wp.array([False, True, False], dtype=wp.bool, device="cpu")) + env._reset_mask(env_mask=wp.zeros(3, dtype=wp.bool, device="cpu")) assert seen[0][0] is env.reset_mask_wp assert seen[1][0] is env.reset_mask_wp @@ -132,7 +132,7 @@ def call(stage, fn, /, *args, **kwargs): env._episode_length_buf_wp = wp.ones(3, dtype=wp.int64, device="cpu") env.extras = {} - env._reset_idx(env_mask=wp.array([False, True, True], dtype=wp.bool, device="cpu")) + env._reset_mask(env_mask=wp.array([False, True, True], dtype=wp.bool, device="cpu")) compute_ids = stages["CurriculumManager_compute"]["env_ids"] torch.testing.assert_close(compute_ids, torch.tensor([1, 2])) @@ -145,7 +145,7 @@ def test_reset_compacts_ids_only_for_active_recorders() -> None: reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") env.termination_manager = SimpleNamespace(dones_wp=reset_mask) env.reset_buf = torch.tensor([False, True, False]) - env._reset_idx = Mock() + env._reset_mask = Mock() env.recorder_manager = Mock() env.recorder_manager.reset.return_value = {"recorder": 1.0} env._has_recorders = True @@ -157,7 +157,7 @@ def test_reset_compacts_ids_only_for_active_recorders() -> None: reset_ids = env.recorder_manager.record_pre_reset.call_args.args[0] torch.testing.assert_close(reset_ids, torch.tensor([1])) - env._reset_idx.assert_called_once_with(env_mask=reset_mask, env_ids=reset_ids) + env._reset_mask.assert_called_once_with(env_mask=reset_mask, env_ids=reset_ids) env.recorder_manager.reset.assert_called_once_with(reset_ids) env.recorder_manager.record_post_reset.assert_called_once_with(reset_ids) assert env.extras["log"]["recorder"] == 1.0 @@ -169,7 +169,7 @@ def test_rtx_rerender_does_not_require_compact_reset_ids() -> None: reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") env.termination_manager = SimpleNamespace(dones_wp=reset_mask) env.reset_buf = torch.tensor([False, True, False]) - env._reset_idx = Mock() + env._reset_mask = Mock() env._has_recorders = False env.has_rtx_sensors = True env.cfg = SimpleNamespace(compute_final_obs=False, num_rerenders_on_reset=2) @@ -178,7 +178,7 @@ def test_rtx_rerender_does_not_require_compact_reset_ids() -> None: env._reset_terminated_envs() - env._reset_idx.assert_called_once_with(env_mask=reset_mask, env_ids=None) + env._reset_mask.assert_called_once_with(env_mask=reset_mask, env_ids=None) assert env.sim.render.call_count == 2 @@ -188,13 +188,13 @@ def test_empty_reset_skips_warp_pipeline_without_compacting_ids() -> None: reset_mask = wp.zeros(3, dtype=wp.bool, device="cpu") env.termination_manager = SimpleNamespace(dones_wp=reset_mask) env.reset_buf = torch.zeros(3, dtype=torch.bool) - env._reset_idx = Mock() + env._reset_mask = Mock() env._has_recorders = False env.extras = {"log": {"previous": 1.0}} env._reset_terminated_envs() - env._reset_idx.assert_not_called() + env._reset_mask.assert_not_called() assert env.extras["log"] == {"previous": 1.0} @@ -206,7 +206,7 @@ def test_reset_captures_final_observation_before_reset() -> None: env.termination_manager = SimpleNamespace(dones_wp=reset_mask) env.reset_buf = torch.tensor([False, True, False]) env.observation_manager = SimpleNamespace(compute=Mock(return_value=final_obs)) - env._reset_idx = Mock() + env._reset_mask = Mock() env._has_recorders = False env.has_rtx_sensors = False env.cfg = SimpleNamespace(compute_final_obs=True, num_rerenders_on_reset=0) @@ -216,4 +216,4 @@ def test_reset_captures_final_observation_before_reset() -> None: assert env.extras["final_obs"] is final_obs env.observation_manager.compute.assert_called_once_with() - env._reset_idx.assert_called_once_with(env_mask=reset_mask, env_ids=None) + env._reset_mask.assert_called_once_with(env_mask=reset_mask, env_ids=None) diff --git a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py index 62585d30d62f..d49aa5b498cf 100644 --- a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py +++ b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py @@ -7,10 +7,12 @@ from __future__ import annotations +import os import unittest +import unittest.mock import warp as wp -from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache +from isaaclab_experimental.utils.warp_graph_cache import CAPTURE_ENV_VAR, WarpGraphCache @wp.kernel @@ -37,6 +39,37 @@ def setUp(self): # Capture policy # ------------------------------------------------------------------ + def test_env_var_forces_eager(self): + """Setting the capture environment variable to 0 should force every stage eager.""" + with unittest.mock.patch.dict(os.environ, {CAPTURE_ENV_VAR: "0"}): + cache = WarpGraphCache(device=self.device) + call_count = 0 + + def counted_call(): + nonlocal call_count + call_count += 1 + + for _ in range(3): + cache.call("RewardManager_compute", counted_call) + + self.assertEqual(call_count, 3) + self.assertEqual(cache.captured_stages, ()) + + def test_captured_stages_and_eager_groups_reflect_state(self): + """Introspection should report captured stages and non-capturable groups.""" + self.cache.register_capturability("Scene", False) + src = wp.full(4, value=1.0, dtype=wp.float32, device=self.device) + dst = wp.zeros(4, dtype=wp.float32, device=self.device) + + def launch(): + wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device) + + self.cache.call("RewardManager_compute", launch) + self.cache.call("RewardManager_compute", launch) + + self.assertEqual(self.cache.captured_stages, ("RewardManager_compute",)) + self.assertEqual(self.cache.eager_groups, ("Scene",)) + def test_non_capturable_group_stays_eager(self): """Every stage in a non-capturable group should execute eagerly.""" call_count = 0 From c1bf3941d41cd0ba3fae282e1c4c98aaab08c4c1 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 21 Jul 2026 11:46:33 -0700 Subject: [PATCH 46/58] Add warp-native height_scan and NullCommand twins Rough-terrain tasks need a warp twin for the height_scan observation and command-less tasks need a NullCommand twin once command class_types swap to warp-native implementations. height_scan stays eager (warp_capturable False) until the sensor lazy-update path is audited for graph capture; its output width resolves through a new out_dim='sensor:rays' sentinel. --- .../jichuanh-warp-frontend-cleanup.minor.rst | 5 +- .../envs/mdp/__init__.pyi | 2 + .../envs/mdp/commands/__init__.pyi | 5 +- .../envs/mdp/commands/commands_cfg.py | 9 +++ .../envs/mdp/commands/null_command.py | 62 +++++++++++++++++++ .../envs/mdp/observations.py | 39 ++++++++++++ .../managers/observation_manager.py | 8 +++ 7 files changed, 128 insertions(+), 2 deletions(-) create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/null_command.py diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst index f85cde838393..98b4d7c8ea20 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -2,7 +2,10 @@ Added ^^^^^ * Added Warp-native uniform pose and velocity command terms with pointer-stable - command arrays. + command arrays, and a Warp-native :class:`NullCommand` term for environments + without commands. +* Added a Warp-native ``height_scan`` observation term (eager until sensor + capture-readiness lands) with ray-count output-dimension inference. * Added :class:`~isaaclab_experimental.managers.CurriculumManager` and a boolean-mask reset path for Warp manager-based environments, retaining compact environment IDs only for legacy host consumers. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi index 804dbb689c7a..9459feae7a18 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi @@ -20,6 +20,8 @@ from .actions import ( # noqa: F401 JointPositionActionCfg, ) from .commands import ( # noqa: F401 + NullCommand, + NullCommandCfg, UniformPoseCommand, UniformPoseCommandCfg, UniformVelocityCommand, diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi index e30d38de19df..dcc7b0b99b82 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi @@ -4,12 +4,15 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "NullCommandCfg", "UniformPoseCommandCfg", "UniformVelocityCommandCfg", + "NullCommand", "UniformPoseCommand", "UniformVelocityCommand", ] -from .commands_cfg import UniformPoseCommandCfg, UniformVelocityCommandCfg +from .commands_cfg import NullCommandCfg, UniformPoseCommandCfg, UniformVelocityCommandCfg +from .null_command import NullCommand from .pose_command import UniformPoseCommand from .velocity_command import UniformVelocityCommand diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/commands_cfg.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/commands_cfg.py index 04d09e2a9948..483f62cbd65c 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/commands_cfg.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/commands_cfg.py @@ -9,15 +9,24 @@ from typing import TYPE_CHECKING +from isaaclab.envs.mdp.commands.commands_cfg import NullCommandCfg as _NullCommandCfg from isaaclab.envs.mdp.commands.commands_cfg import UniformPoseCommandCfg as _UniformPoseCommandCfg from isaaclab.envs.mdp.commands.commands_cfg import UniformVelocityCommandCfg as _UniformVelocityCommandCfg from isaaclab.utils.configclass import configclass if TYPE_CHECKING: + from .null_command import NullCommand from .pose_command import UniformPoseCommand from .velocity_command import UniformVelocityCommand +@configclass +class NullCommandCfg(_NullCommandCfg): + """Configuration for the Warp-native null command generator.""" + + class_type: type[NullCommand] | str = "{DIR}.null_command:NullCommand" + + @configclass class UniformVelocityCommandCfg(_UniformVelocityCommandCfg): """Configuration for the Warp-native uniform velocity command generator.""" diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/null_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/null_command.py new file mode 100644 index 000000000000..b210daeefb55 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/null_command.py @@ -0,0 +1,62 @@ +# 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 + +"""Warp-native command term that does nothing.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import warp as wp + +from isaaclab_experimental.managers import CommandTerm + +if TYPE_CHECKING: + from .commands_cfg import NullCommandCfg + + +class NullCommand(CommandTerm): + """Command generator that does nothing. + + Warp-native twin of :class:`isaaclab.envs.mdp.commands.NullCommand`. It does not + generate any commands and is used for environments that do not require one. All + inherited bookkeeping runs as device-side kernels on empty metrics with an + infinite resampling time, so the term is trivially graph-capturable. + """ + + cfg: NullCommandCfg + """Configuration for the command generator.""" + + def __str__(self) -> str: + msg = "NullCommand:\n" + msg += "\tCommand dimension: N/A\n" + msg += f"\tResampling time range: {self.cfg.resampling_time_range}" + return msg + + """ + Properties + """ + + @property + def command(self): + """Null command. + + Raises: + RuntimeError: No command is generated. Always raises this error. + """ + raise RuntimeError("NullCommandTerm does not generate any commands.") + + """ + Implementation specific functions. + """ + + def _update_metrics(self): + pass + + def _resample_command(self, env_mask: wp.array): + pass + + def _update_command(self): + pass diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py index fb0320251a9c..9616e96ef29b 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py @@ -36,10 +36,12 @@ record_shape, ) from isaaclab_experimental.managers import SceneEntityCfg +from isaaclab_experimental.utils.warp import warp_capturable if TYPE_CHECKING: from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedEnv + from isaaclab.sensors import RayCaster # --------------------------------------------------------------------------- @@ -356,3 +358,40 @@ def generated_commands(env: ManagerBasedEnv, out, command_name: str) -> None: Reads the command manager's pointer-stable Warp storage directly. """ wp.copy(out, env.command_manager.get_command_wp(command_name)) + + +""" +Sensors. +""" + + +@wp.kernel +def _height_scan_kernel( + pos_w: wp.array(dtype=wp.vec3f), + ray_hits_w: wp.array2d(dtype=wp.vec3f), + offset: wp.float32, + out: wp.array2d(dtype=wp.float32), +): + env_id, ray_id = wp.tid() + out[env_id, ray_id] = pos_w[env_id][2] - ray_hits_w[env_id, ray_id][2] - offset + + +# Sensor reads go through the lazy-update path, whose host-side timestamp bookkeeping +# decides when rays are re-cast and has not been audited for graph capture. +@warp_capturable(False) +@generic_io_descriptor_warp(units="m", out_dim="sensor:rays", observation_type="SensorState", on_inspect=[record_shape]) +def height_scan(env: ManagerBasedEnv, out, sensor_cfg: SceneEntityCfg, offset: float = 0.5) -> None: + """Height scan from the given sensor w.r.t. the sensor's frame [m]. Writes into ``out``. + + Warp-first override of :func:`isaaclab.envs.mdp.observations.height_scan`. + The provided offset (Defaults to 0.5) is subtracted from the returned values. + Missed rays carry ``inf`` hit positions, matching the stable term's semantics + (the resulting ``-inf`` heights are handled by the term's ``clip`` setting). + """ + sensor: RayCaster = env.scene.sensors[sensor_cfg.name] + wp.launch( + kernel=_height_scan_kernel, + dim=(env.num_envs, sensor.num_rays), + inputs=[sensor.data.pos_w.warp, sensor.data.ray_hits_w.warp, float(offset), out], + device=env.device, + ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index 0825e917e2f2..23468ed50c45 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -897,10 +897,18 @@ def _resolve_out_dim(self, out_dim: int | str, term_cfg: ObservationTermCfg) -> - ``"body:N"``: ``N`` components per selected body from ``asset_cfg``. - ``"command"``: query ``command_manager.get_command(name).shape[-1]``. - ``"action"``: query ``action_manager.action.shape[-1]``. + - ``"sensor:rays"``: number of rays of the ray-caster sensor from ``sensor_cfg``. """ if isinstance(out_dim, int): return out_dim + if out_dim == "sensor:rays": + sensor_cfg = term_cfg.params.get("sensor_cfg") + if sensor_cfg is None: + raise ValueError("out_dim='sensor:rays' requires a 'sensor_cfg' entry in the term's params.") + sensor = self._env.scene.sensors[sensor_cfg.name] + return int(sensor.num_rays) + if out_dim == "joint": asset_cfg = term_cfg.params.get("asset_cfg") asset = self._env.scene[asset_cfg.name] From 54a3c5f5e77286e5fa56d83a47c776b0d5e76446 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 21 Jul 2026 11:05:59 -0700 Subject: [PATCH 47/58] Swap command term classes to warp twins in the bridge The warp env runs commands on the experimental captured CommandManager, so stable command term class_types must be swapped to their warp twins like every other warp-managed group. Stable CommandTermCfg is a standalone base (like ActionTermCfg), so the term walker must match it explicitly; without this the swap silently skipped command terms and env construction failed with a CommandType TypeError. --- .../isaaclab_experimental/envs/frontend.py | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index 04d9d1886969..d918016ff387 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -85,13 +85,15 @@ class WarpFrontend: # are adapted (SceneEntityCfg promotion + MDP twin swap). The event manager # is warp-first too — it invokes term funcs with a Warp env-mask, so a stable # event func (which expects torch ``env_ids``) breaks at runtime; its funcs - # must be swapped to warp twins. The curriculum, recorder and command - # managers run on the stable (torch) implementation, so their terms are left - # untouched. A stable term left on a warp manager would break, so a missing - # twin in these groups is a hard error; a stable term on a stable manager is - # correct, so those groups are skipped. + # must be swapped to warp twins. The command manager is the warp-native + # ``CommandManager`` (captured), so command term ``class_type``s must swap + # to their warp twins as well. The curriculum and recorder managers run on + # the stable (torch) implementation, so their terms are left untouched. A + # stable term left on a warp manager would break, so a missing twin in these + # groups is a hard error; a stable term on a stable manager is correct, so + # those groups are skipped. WARP_MANAGED_GROUPS: ClassVar[frozenset[str]] = frozenset( - {"observations", "rewards", "terminations", "actions", "events"} + {"observations", "rewards", "terminations", "actions", "events", "commands"} ) # ------------------------------------------------------------------ @@ -483,10 +485,10 @@ def _walk_terms(node: Any, path: tuple[str, ...] = ()) -> Iterator[tuple[tuple[s """Yield ``(path, term)`` for every MDP term cfg in the cfg tree. A "term" is a :class:`ManagerTermBaseCfg` (observation/reward/ - termination/event/curriculum) *or* an :class:`ActionTermCfg` — the - latter is a separate base that is **not** a ``ManagerTermBaseCfg`` - subclass, yet carries a swappable ``class_type``, so it must be matched - explicitly. + termination/event/curriculum) *or* an :class:`ActionTermCfg` / + :class:`CommandTermCfg` — the latter two are separate bases that are + **not** ``ManagerTermBaseCfg`` subclasses, yet carry a swappable + ``class_type``, so they must be matched explicitly. Behavior at each node: @@ -510,9 +512,9 @@ def _walk_terms(node: Any, path: tuple[str, ...] = ()) -> Iterator[tuple[tuple[s cfg layouts (extra observation groups, new nesting, etc.) are picked up automatically as long as their terms subclass one of the term base cfgs. """ - from isaaclab.managers.manager_term_cfg import ActionTermCfg, ManagerTermBaseCfg + from isaaclab.managers.manager_term_cfg import ActionTermCfg, CommandTermCfg, ManagerTermBaseCfg - if isinstance(node, (ManagerTermBaseCfg, ActionTermCfg)): + if isinstance(node, (ManagerTermBaseCfg, ActionTermCfg, CommandTermCfg)): yield path, node # a term: yield and stop; never descend into params/func return if not hasattr(node, "__dataclass_fields__"): From bb4e29a3f666aa5928fb49d6f0cd7cba19c49dfd Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 21 Jul 2026 12:12:40 -0700 Subject: [PATCH 48/58] Reroute test imports to the restructured task packages The bridge merge moves experimental velocity/reach MDP modules under isaaclab_tasks_experimental.core and retires reach_env_cfg; the curriculum-manager and terrain-out-of-bounds tests import from the new locations. --- .../test/managers/test_curriculum_manager.py | 20 ++++++++++++++----- .../test/test_terrain_out_of_bounds.py | 2 +- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/source/isaaclab_experimental/test/managers/test_curriculum_manager.py b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py index 532a0fadeb85..11910d139907 100644 --- a/source/isaaclab_experimental/test/managers/test_curriculum_manager.py +++ b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py @@ -14,10 +14,8 @@ import torch import warp as wp from isaaclab_experimental.managers import CurriculumManager, CurriculumTermCfg -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp import TerrainLevelsVel, terrain_levels_vel -from isaaclab_tasks_experimental.manager_based.manipulation.reach.reach_env_cfg import ( - CurriculumCfg as ReachCurriculumCfg, -) +from isaaclab_tasks_experimental.core.reach.mdp import ModifyRewardWeight +from isaaclab_tasks_experimental.core.velocity.mdp import TerrainLevelsVel, terrain_levels_vel from isaaclab.managers import ManagerTermBase as StableManagerTermBase from isaaclab.terrains import TerrainImporter @@ -419,7 +417,19 @@ def test_registered_reach_curricula_update_weights_without_a_host_boundary(): ) env.reward_manager = _RewardManager() env.common_step_counter = 0 - manager = CurriculumManager(ReachCurriculumCfg(), env) + manager = CurriculumManager( + { + "action_rate": CurriculumTermCfg( + func=ModifyRewardWeight, + params={"term_name": "action_rate", "weight": -0.005, "num_steps": 4500}, + ), + "joint_vel": CurriculumTermCfg( + func=ModifyRewardWeight, + params={"term_name": "joint_vel", "weight": -0.001, "num_steps": 4500}, + ), + }, + env, + ) env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") assert not manager.requires_host_ids diff --git a/source/isaaclab_tasks_experimental/test/test_terrain_out_of_bounds.py b/source/isaaclab_tasks_experimental/test/test_terrain_out_of_bounds.py index c306608323ec..f04d225433a8 100644 --- a/source/isaaclab_tasks_experimental/test/test_terrain_out_of_bounds.py +++ b/source/isaaclab_tasks_experimental/test/test_terrain_out_of_bounds.py @@ -9,7 +9,7 @@ import numpy as np import warp as wp -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp.terminations import terrain_out_of_bounds +from isaaclab_tasks_experimental.core.velocity.mdp.terminations import terrain_out_of_bounds class _Scene(dict): From 466decc9d02bb9f25bbf35d5c0e5236b960844ec Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 21 Jul 2026 12:56:07 -0700 Subject: [PATCH 49/58] Document all Warp kernels in the experimental manager stack Every @wp.kernel carries a purpose docstring (domain quantity, units, launch dimensionality where non-trivial) so kernels self-explain instead of reading as bare index math. --- .../envs/manager_based_env_warp.py | 4 +++ .../envs/mdp/commands/pose_command.py | 12 +++++++ .../envs/mdp/commands/velocity_command.py | 22 +++++++++++++ .../isaaclab_experimental/envs/mdp/events.py | 22 +++++++++++++ .../envs/mdp/observations.py | 32 +++++++++++++++++++ .../isaaclab_experimental/envs/mdp/rewards.py | 29 +++++++++++++++++ .../managers/command_manager.py | 12 +++++++ .../managers/event_manager.py | 22 +++++++++++++ .../managers/observation_manager.py | 2 ++ .../managers/reward_manager.py | 14 ++++++++ .../test/managers/test_reward_manager.py | 1 + 11 files changed, 172 insertions(+) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py index d14283016a1f..55afdb1df06b 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py @@ -53,6 +53,10 @@ def initialize_rng_state( # output state: wp.array(dtype=wp.uint32), ): + """Initialize each env's persistent RNG state as an independent stream of the + + global ``seed`` (one ``wp.rand_init`` stream per env id). + """ env_id = wp.tid() state[env_id] = wp.rand_init(seed, wp.int32(env_id)) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py index a0d4b90ff7de..bb3944ad07ea 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py @@ -43,6 +43,12 @@ def _resample_pose_command( yaw_max: float, make_quat_unique: bool, ): + """Draw a new body-frame pose command for selected envs: uniform position [m] + + and roll/pitch/yaw [rad] composed into a quaternion (optionally + sign-canonicalized to the w >= 0 hemisphere). RNG state is written back so + the sequence advances across replays. + """ env_id = wp.tid() if env_mask[env_id]: state = rng_state[env_id] @@ -81,6 +87,12 @@ def _update_pose_metrics( track_success: bool, position_success_threshold: float, ): + """Refresh the world-frame goal pose (root pose composed with the body-frame + + command) and the tracking metrics: end-effector position [m] / orientation + [rad] errors vs the goal, and a position-threshold success flag when success + tracking is enabled. + """ env_id = wp.tid() position_b = wp.vec3f(command_b[env_id, 0], command_b[env_id, 1], command_b[env_id, 2]) orientation_b = wp.quatf(command_b[env_id, 3], command_b[env_id, 4], command_b[env_id, 5], command_b[env_id, 6]) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py index 7eb952fc6f5f..87707d64f103 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py @@ -38,6 +38,12 @@ def _accumulate_velocity_metrics( error_yaw_sum: wp.array(dtype=wp.float32), step_count: wp.array(dtype=wp.float32), ): + """Accumulate per-step command-tracking error sums for logging: planar (xy) + + linear-velocity error [m/s] and yaw-rate error [rad/s] vs the command, plus + the step count used later for averaging. Body-frame velocities are derived + inline from the root pose/velocity (capture-safe, no lazy buffers). + """ env_id = wp.tid() root_lin_vel_b = body_lin_vel_from_root(root_pose_w[env_id], root_vel_w[env_id]) root_ang_vel_b = body_ang_vel_from_root(root_pose_w[env_id], root_vel_w[env_id]) @@ -60,6 +66,11 @@ def _finalize_velocity_metrics( error_xy_threshold: float, error_yaw_threshold: float, ): + """Turn selected envs' accumulated error sums into episode means and a binary + + success flag (both mean errors under their thresholds), then zero the + accumulators for the next episode. + """ env_id = wp.tid() if env_mask[env_id]: denominator = wp.max(step_count[env_id], 1.0) @@ -95,6 +106,12 @@ def _resample_velocity_command( rel_heading_envs: float, rel_standing_envs: float, ): + """Draw a new SE(2) velocity command (vx, vy [m/s], wz [rad/s]) for selected + + envs, plus their heading target [rad] and heading/standing role flags. + The per-env device RNG state is written back so the random sequence advances + across CUDA-graph replays. + """ env_id = wp.tid() if env_mask[env_id]: state = rng_state[env_id] @@ -122,6 +139,11 @@ def _update_velocity_command( ang_vel_z_min: float, ang_vel_z_max: float, ): + """Per-step command shaping: heading-controlled envs convert heading error into + + a clamped yaw-rate command (proportional control); standing envs zero their + command entirely. + """ env_id = wp.tid() if heading_command and is_heading_env[env_id]: root_quat_w = wp.transform_get_rotation(root_pose_w[env_id]) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py index a07a69e4a811..c3349d485efb 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py @@ -173,6 +173,10 @@ def _apply_external_force_torque_kernel( torque_lo: float, torque_hi: float, ): + """Sample uniform random external force [N] and torque [N·m] vectors for + + selected envs' target bodies into the asset's wrench composer buffers. + """ env_id = wp.tid() if not env_mask[env_id]: return @@ -281,6 +285,10 @@ def _push_by_setting_velocity_kernel( ang_lo: wp.vec3f, ang_hi: wp.vec3f, ): + """Add a uniform random velocity kick [m/s, rad/s] to selected envs' current + + root velocity, writing the result for the masked sim write. + """ env_id = wp.tid() if not env_mask[env_id]: return @@ -379,6 +387,11 @@ def _reset_root_state_uniform_kernel( vel_ang_lo: wp.vec3f, vel_ang_hi: wp.vec3f, ): + """Compose selected envs' reset root state: default pose offset by the env + + origin plus uniform position [m] / roll-pitch-yaw [rad] noise, and default + velocity plus uniform noise [m/s, rad/s]. + """ env_id = wp.tid() if not env_mask[env_id]: return @@ -523,6 +536,11 @@ def _reset_joints_by_offset_kernel( vel_lo: float, vel_hi: float, ): + """Reset selected envs' selected joints to defaults plus uniform offsets + + [rad or m, depending on joint type], clamped to the soft position and + velocity limits. + """ env_id = wp.tid() if not env_mask[env_id]: return @@ -621,6 +639,10 @@ def _reset_joints_by_scale_kernel( vel_lo: float, vel_hi: float, ): + """Reset selected envs' selected joints to defaults scaled by uniform random + + factors, clamped to the soft position and velocity limits. + """ env_id = wp.tid() if not env_mask[env_id]: return diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py index 497fe6699aff..8aa4fedfd9b5 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py @@ -54,6 +54,7 @@ def _vec3_to_out3_kernel( src: wp.array(dtype=wp.vec3f), out: wp.array(dtype=wp.float32, ndim=2), ): + """Scatter a vec3 array into three float32 columns of a term's ``out`` block.""" env_id = wp.tid() v = src[env_id] out[env_id, 0] = v[0] @@ -67,6 +68,10 @@ def _joint_gather_kernel( joint_ids: wp.array(dtype=wp.int32), out: wp.array(dtype=wp.float32, ndim=2), ): + """Gather selected joints' values into contiguous term columns. + + Launch with dim=(num_envs, num_selected_joints). + """ env_id, k = wp.tid() j = joint_ids[k] out[env_id, k] = src[env_id, j] @@ -82,6 +87,7 @@ def _base_pos_z_kernel( root_pos_w: wp.array(dtype=wp.vec3f), out: wp.array(dtype=wp.float32, ndim=2), ): + """Write the root height above the world origin [m] into the term's single column.""" env_id = wp.tid() out[env_id, 0] = root_pos_w[env_id][2] @@ -114,6 +120,10 @@ def _base_lin_vel_kernel( root_vel_w: wp.array(dtype=wp.spatial_vectorf), out: wp.array(dtype=wp.float32, ndim=2), ): + """Root linear velocity [m/s] in the body frame, derived inline from the root + + pose and CoM velocity (Tier-1 access; avoids non-capturable lazy buffers). + """ i = wp.tid() v = body_lin_vel_from_root(root_pose_w[i], root_vel_w[i]) out[i, 0] = v[0] @@ -141,6 +151,10 @@ def _base_ang_vel_kernel( root_vel_w: wp.array(dtype=wp.spatial_vectorf), out: wp.array(dtype=wp.float32, ndim=2), ): + """Root angular velocity [rad/s] in the body frame, derived inline from the root + + pose and CoM velocity (Tier-1 access; avoids non-capturable lazy buffers). + """ i = wp.tid() v = body_ang_vel_from_root(root_pose_w[i], root_vel_w[i]) out[i, 0] = v[0] @@ -168,6 +182,7 @@ def _projected_gravity_kernel( gravity_w: wp.array(dtype=wp.vec3f), out: wp.array(dtype=wp.float32, ndim=2), ): + """Unit gravity direction rotated into the body frame (flat pose -> (0, 0, -1)).""" i = wp.tid() g = rotate_vec_to_body_frame(wp.normalize(gravity_w[i]), root_pose_w[i]) out[i, 0] = g[0] @@ -221,6 +236,10 @@ def _joint_rel_gather_kernel( joint_ids: wp.array(dtype=wp.int32), out: wp.array(dtype=wp.float32, ndim=2), ): + """Gather selected joints' values relative to their per-env defaults into + + contiguous term columns. Launch with dim=(num_envs, num_selected_joints). + """ env_id, k = wp.tid() j = joint_ids[k] out[env_id, k] = values[env_id, j] - defaults[env_id, j] @@ -257,6 +276,10 @@ def _joint_pos_limit_normalized_kernel( joint_ids: wp.array(dtype=wp.int32), out: wp.array(dtype=wp.float32, ndim=2), ): + """Selected joints' positions normalized to [-1, 1] within their soft limits. + + Launch with dim=(num_envs, num_selected_joints). + """ env_id, k = wp.tid() j = joint_ids[k] pos = joint_pos[env_id, j] @@ -372,6 +395,10 @@ def _body_incoming_wrench_kernel( body_ids: wp.array(dtype=wp.int32), out: wp.array(dtype=wp.float32, ndim=2), ): + """Gather selected bodies' incoming force [N] and torque [N·m] into per-body + + ``[fx, fy, fz, tx, ty, tz]`` column blocks. + """ env_id = wp.tid() for k in range(body_ids.shape[0]): b = body_ids[k] @@ -411,6 +438,11 @@ def _height_scan_kernel( offset: wp.float32, out: wp.array2d(dtype=wp.float32), ): + """Height of the scan frame above each ray hit [m]: sensor_z - hit_z - offset. + + Launch with dim=(num_envs, num_rays). Missed hits carry ``inf`` and produce + ``-inf`` heights, matching the stable term (handled by the term's clip). + """ env_id, ray_id = wp.tid() out[env_id, ray_id] = pos_w[env_id][2] - ray_hits_w[env_id, ray_id][2] - offset diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py index c4bde2e7b5a4..833a656a9a45 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py @@ -38,6 +38,7 @@ @wp.kernel def _is_alive_kernel(terminated: wp.array(dtype=wp.bool), out: wp.array(dtype=wp.float32)): + """1.0 for envs still running, 0.0 for terminated ones.""" i = wp.tid() out[i] = wp.where(terminated[i], 0.0, 1.0) @@ -54,6 +55,7 @@ def is_alive(env: ManagerBasedRLEnv, out: wp.array(dtype=wp.float32)) -> None: @wp.kernel def _is_terminated_kernel(terminated: wp.array(dtype=wp.bool), out: wp.array(dtype=wp.float32)): + """1.0 for envs terminated by a non-timeout cause, else 0.0.""" i = wp.tid() out[i] = wp.where(terminated[i], 1.0, 0.0) @@ -87,6 +89,10 @@ def _lin_vel_z_l2_kernel( root_vel_w: wp.array(dtype=wp.spatial_vectorf), out: wp.array(dtype=wp.float32), ): + """Squared body-frame vertical velocity [(m/s)^2] per env, derived inline from + + the root state. + """ i = wp.tid() vz = body_lin_vel_from_root(root_pose_w[i], root_vel_w[i])[2] out[i] = vz * vz @@ -109,6 +115,7 @@ def _ang_vel_xy_l2_kernel( root_vel_w: wp.array(dtype=wp.spatial_vectorf), out: wp.array(dtype=wp.float32), ): + """Sum of squared body-frame roll/pitch rates [(rad/s)^2] per env.""" i = wp.tid() v = body_ang_vel_from_root(root_pose_w[i], root_vel_w[i]) out[i] = v[0] * v[0] + v[1] * v[1] @@ -131,6 +138,10 @@ def _flat_orientation_l2_kernel( gravity_w: wp.array(dtype=wp.vec3f), out: wp.array(dtype=wp.float32), ): + """Sum of squared xy components of body-frame projected gravity (zero when the + + base is perfectly flat). + """ # ``gravity_w`` is per-env and may carry magnitude (Newton, m/s^2), so index # per env and normalize before projecting. i = wp.tid() @@ -160,6 +171,7 @@ def flat_orientation_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = def _sum_sq_masked_kernel( x: wp.array(dtype=wp.float32, ndim=2), joint_mask: wp.array(dtype=wp.bool), out: wp.array(dtype=wp.float32) ): + """Per-env sum of squares over mask-selected joint columns.""" i = wp.tid() s = float(0.0) for j in range(x.shape[1]): @@ -184,6 +196,7 @@ def joint_torques_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = Sc def _sum_abs_masked_kernel( x: wp.array(dtype=wp.float32, ndim=2), joint_mask: wp.array(dtype=wp.bool), out: wp.array(dtype=wp.float32) ): + """Per-env L1 sum over mask-selected joint columns.""" i = wp.tid() s = float(0.0) for j in range(x.shape[1]): @@ -233,6 +246,7 @@ def _sum_abs_diff_masked_kernel( joint_mask: wp.array(dtype=wp.bool), out: wp.array(dtype=wp.float32), ): + """Per-env L1 sum of ``a - b`` over mask-selected joint columns.""" i = wp.tid() s = float(0.0) for j in range(a.shape[1]): @@ -260,6 +274,11 @@ def _joint_pos_limits_kernel( joint_mask: wp.array(dtype=wp.bool), out: wp.array(dtype=wp.float32), ): + """Per-env sum of soft-limit violations [rad or m, depending on joint type]: + + distance each selected joint sits below its lower or above its upper soft + limit (zero inside the limits). + """ i = wp.tid() s = float(0.0) for j in range(joint_pos.shape[1]): @@ -302,6 +321,7 @@ def _sum_sq_diff_2d_kernel( b: wp.array(dtype=wp.float32, ndim=2), out: wp.array(dtype=wp.float32), ): + """Per-env sum of squared differences across all columns (action-rate penalty).""" i = wp.tid() s = float(0.0) for j in range(a.shape[1]): @@ -323,6 +343,7 @@ def action_rate_l2(env: ManagerBasedRLEnv, out) -> None: # TODO(warp-migration): Revisit 2D kernel + wp.atomic_add vs 1D inner loop. @wp.kernel def _sum_sq_2d_kernel(x: wp.array(dtype=wp.float32, ndim=2), out: wp.array(dtype=wp.float32)): + """Per-env sum of squares across all columns (action-magnitude penalty).""" i = wp.tid() s = float(0.0) for j in range(x.shape[1]): @@ -395,6 +416,10 @@ def _track_lin_vel_xy_exp_kernel( std_sq_inv: float, out: wp.array(dtype=wp.float32), ): + """Exponential tracking shaping ``exp(-err^2 / std^2)`` of planar body-frame + + linear velocity vs the command's (vx, vy). + """ i = wp.tid() v = body_lin_vel_from_root(root_pose_w[i], root_vel_w[i]) dx = command[i, 0] - v[0] @@ -438,6 +463,10 @@ def _track_ang_vel_z_exp_kernel( std_sq_inv: float, out: wp.array(dtype=wp.float32), ): + """Exponential tracking shaping ``exp(-err^2 / std^2)`` of body-frame yaw rate + + vs the command's wz. + """ i = wp.tid() dz = command[i, cmd_col] - body_ang_vel_from_root(root_pose_w[i], root_vel_w[i])[2] out[i] = wp.exp(-dz * dz * std_sq_inv) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py index 16d6758716f5..6a459dcf9b66 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py @@ -36,6 +36,11 @@ def _sum_and_zero_masked( metric: wp.array(dtype=wp.float32), out_mean: wp.array(dtype=wp.float32), ): + """Fold selected envs' episode metric into a running mean and zero it for the next episode. + + ``scale[0]`` carries the 1/count normalization; the atomic add reduces across all + selected envs into ``out_mean[0]``. + """ env_id = wp.tid() if mask[env_id]: wp.atomic_add(out_mean, 0, metric[env_id] * scale[0]) @@ -44,6 +49,7 @@ def _sum_and_zero_masked( @wp.kernel def _zero_counter_masked(mask: wp.array(dtype=wp.bool), counter: wp.array(dtype=wp.int32)): + """Zero the per-env resample counter for envs selected by ``mask``.""" env_id = wp.tid() if mask[env_id]: counter[env_id] = 0 @@ -55,6 +61,7 @@ def _step_time_left_and_build_resample_mask( dt: wp.float32, out_mask: wp.array(dtype=wp.bool), ): + """Advance each env's resample countdown by ``dt`` [s] and flag expired envs in ``out_mask``.""" env_id = wp.tid() t = time_left[env_id] - dt time_left[env_id] = t @@ -70,6 +77,11 @@ def _resample_time_left_and_increment_counter( lower: wp.float32, upper: wp.float32, ): + """Draw a new resampling interval [s] for flagged envs and count the resample. + + Uses the per-env device RNG state with write-back so the random sequence + advances across CUDA-graph replays. + """ env_id = wp.tid() if mask[env_id]: s = rng_state[env_id] diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py index 0a7ad5ae74f9..b00ffc6daf92 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py @@ -43,6 +43,7 @@ def _interval_init_per_env( lower: wp.float32, upper: wp.float32, ): + """Seed each env's interval-event countdown with a uniform draw from [``lower``, ``upper``] [s].""" env_id = wp.tid() s = rng_state[env_id] time_left[env_id] = wp.randf(s, lower, upper) @@ -56,6 +57,7 @@ def _interval_init_global( lower: wp.float32, upper: wp.float32, ): + """Seed the single shared countdown of a ``is_global_time`` interval event.""" # single element s = rng_state[0] time_left[0] = wp.randf(s, lower, upper) @@ -71,6 +73,10 @@ def _interval_step_per_env( lower: wp.float32, upper: wp.float32, ): + """Tick each env's interval countdown by ``dt`` [s]; on expiry, flag the env in + + ``trigger_mask`` and immediately draw its next interval (RNG state written back). + """ env_id = wp.tid() t = time_left[env_id] - dt if t < wp.float32(1.0e-6): @@ -92,6 +98,10 @@ def _interval_step_global( lower: wp.float32, upper: wp.float32, ): + """Tick the shared countdown of a ``is_global_time`` interval event; on expiry set + + ``trigger_flag[0]`` and draw the next interval. + """ t = time_left[0] - dt if t < wp.float32(1.0e-6): trigger_flag[0] = True @@ -111,6 +121,7 @@ def _interval_reset_selected( lower: wp.float32, upper: wp.float32, ): + """Re-draw interval countdowns for envs selected by ``env_mask`` (episode-reset path).""" env_id = wp.tid() if env_mask[env_id]: s = rng_state[env_id] @@ -123,6 +134,10 @@ def _seed_global_rng_from_env_rng( env_rng_state: wp.array(dtype=wp.uint32), global_rng_state: wp.array(dtype=wp.uint32), ): + """Derive the global-event RNG stream from env 0's per-env state so global draws + + stay seed-reproducible without a host round-trip. + """ global_rng_state[0] = wp.rand_init(wp.int32(env_rng_state[0]), wp.int32(0)) @@ -135,6 +150,13 @@ def _reset_compute_valid_mask( global_step_count_buf: wp.array(dtype=wp.int32), min_step_count: wp.int32, ): + """Gate a reset-event mask by trigger spacing entirely on device. + + An env stays selected only if at least ``min_step_count`` steps passed since it + last triggered, or it has never triggered (stable-parity: first reset always + fires). Passing envs update ``last_triggered_step``/``triggered_once`` in the + same launch, keeping the spacing bookkeeping replay-safe. + """ env_id = wp.tid() if not in_mask[env_id]: out_mask[env_id] = False diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index ac14165434e7..12fb601acbae 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -70,6 +70,7 @@ @wp.kernel def _apply_clip(out: wp.array(dtype=wp.float32, ndim=2), clip_lo: wp.float32, clip_hi: wp.float32): + """Clamp every column of a term's ``out`` block to [``clip_lo``, ``clip_hi``] in place.""" env_id = wp.tid() for j in range(out.shape[1]): out[env_id, j] = wp.clamp(out[env_id, j], clip_lo, clip_hi) @@ -77,6 +78,7 @@ def _apply_clip(out: wp.array(dtype=wp.float32, ndim=2), clip_lo: wp.float32, cl @wp.kernel def _apply_scale(out: wp.array(dtype=wp.float32, ndim=2), scale: wp.array(dtype=wp.float32)): + """Multiply each column of a term's ``out`` block by its per-column ``scale`` in place.""" env_id = wp.tid() for j in range(out.shape[1]): out[env_id, j] = out[env_id, j] * scale[j] diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py index 3cf7c5d2ec43..9f8490ed1c31 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py @@ -38,6 +38,13 @@ def _sum_and_zero_masked( # output out_avg: wp.array(dtype=wp.float32), ): + """Fold selected envs' per-term episode reward sums into per-term means and zero + + them for the next episode. + Launch with dim=(num_terms, num_envs); ``scale[0]`` carries the 1/count + normalization and the atomic add reduces across selected envs into + ``out_avg[term]``. + """ term_idx, env_id = wp.tid() if mask[env_id]: wp.atomic_add(out_avg, term_idx, episode_sums[term_idx, env_id] * scale[0]) @@ -55,6 +62,13 @@ def _reward_finalize( episode_sums: wp.array(dtype=wp.float32, ndim=2), step_reward: wp.array(dtype=wp.float32, ndim=2), ): + """Combine per-term outputs into the env step reward and episode bookkeeping. + + Per env: ``step_reward`` records each term's weighted reward rate, + ``reward_buf`` accumulates ``weight * out * dt`` across terms, and + ``episode_sums`` grows by the same amount for episode metrics. Zero-weight + terms are skipped so disabled terms cost nothing. + """ env_id = wp.tid() total = wp.float32(0.0) diff --git a/source/isaaclab_experimental/test/managers/test_reward_manager.py b/source/isaaclab_experimental/test/managers/test_reward_manager.py index 569465822377..3d056470540f 100644 --- a/source/isaaclab_experimental/test/managers/test_reward_manager.py +++ b/source/isaaclab_experimental/test/managers/test_reward_manager.py @@ -14,6 +14,7 @@ @wp.kernel def _fill_unit_reward(out: wp.array(dtype=wp.float32)): + """Test helper: write 1.0 into every env's term output.""" out[wp.tid()] = 1.0 From 3333e2aeae72aeb90b5f6b6e9415628c9e3ff549 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 22 Jul 2026 00:02:22 -0700 Subject: [PATCH 50/58] Route play entrypoints through the frontend factory Play backends constructed environments with a direct gym.make, so warp-trained policies could not be played through the official CLI: the shared parser accepted --frontend but play ignored it. All four dispatched play backends now build their environment via create_isaaclab_env, mirroring their train twins' MARL-conversion semantics, and a unit test pins the routing. The rlinf integration constructs environments inside the external framework and remains frontend-unroutable (documented in the changelog). --- .../rl_games/play_rl_games.py | 18 +++++++++--------- .../rsl_rl/play_rsl_rl.py | 18 +++++++++--------- scripts/reinforcement_learning/sb3/play_sb3.py | 16 ++++++++-------- .../reinforcement_learning/skrl/play_skrl.py | 16 ++++++++-------- .../test/test_reinforcement_learning_common.py | 14 ++++++++++++++ .../changelog.d/warp-manager-bridge.rst | 6 ++++-- 6 files changed, 52 insertions(+), 36 deletions(-) diff --git a/scripts/reinforcement_learning/rl_games/play_rl_games.py b/scripts/reinforcement_learning/rl_games/play_rl_games.py index e9fb02af9a12..d08126f08981 100644 --- a/scripts/reinforcement_learning/rl_games/play_rl_games.py +++ b/scripts/reinforcement_learning/rl_games/play_rl_games.py @@ -16,7 +16,7 @@ import gymnasium as gym import torch -from common import CHECKPOINT_SELECTORS, resolve_checkpoint_selector +from common import CHECKPOINT_SELECTORS, add_frontend_args, create_isaaclab_env, resolve_checkpoint_selector from rl_games.common import env_configurations, vecenv from rl_games.common.player import BasePlayer from rl_games.torch_runner import Runner @@ -49,6 +49,7 @@ ) 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.") +add_frontend_args(parser) parser.add_argument( "--agent", type=str, default="rl_games_cfg_entry_point", help="Name of the RL agent configuration entry point." ) @@ -137,14 +138,13 @@ def main(): obs_groups = agent_cfg["params"]["env"].get("obs_groups") concate_obs_groups = agent_cfg["params"]["env"].get("concate_obs_groups", True) - # create isaac environment - env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) - - # convert to single-agent instance if required by the RL algorithm - if isinstance(env.unwrapped.cfg, DirectMARLEnvCfg): - from isaaclab.envs import multi_agent_to_single_agent - - env = multi_agent_to_single_agent(env) + # create isaac environment on the selected frontend + env = create_isaaclab_env( + args_cli.task, + env_cfg, + args_cli, + convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg), + ) # wrap for video recording if args_cli.video: diff --git a/scripts/reinforcement_learning/rsl_rl/play_rsl_rl.py b/scripts/reinforcement_learning/rsl_rl/play_rsl_rl.py index 6f04485f71ad..2bcf37d95c7b 100644 --- a/scripts/reinforcement_learning/rsl_rl/play_rsl_rl.py +++ b/scripts/reinforcement_learning/rsl_rl/play_rsl_rl.py @@ -14,7 +14,7 @@ import gymnasium as gym import torch -from common import CHECKPOINT_SELECTORS, resolve_checkpoint_selector +from common import CHECKPOINT_SELECTORS, add_frontend_args, create_isaaclab_env, resolve_checkpoint_selector from packaging import version from rsl_rl.runners import DistillationRunner, OnPolicyRunner @@ -56,6 +56,7 @@ ) 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.") +add_frontend_args(parser) parser.add_argument( "--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point." ) @@ -140,14 +141,13 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen # set the log directory for the environment env_cfg.log_dir = log_dir - # create isaac environment - env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) - - # convert to single-agent instance if required by the RL algorithm - if isinstance(env.unwrapped.cfg, DirectMARLEnvCfg): - from isaaclab.envs import multi_agent_to_single_agent - - env = multi_agent_to_single_agent(env) + # create isaac environment on the selected frontend + env = create_isaaclab_env( + args_cli.task, + env_cfg, + args_cli, + convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg), + ) # wrap for video recording if args_cli.video: diff --git a/scripts/reinforcement_learning/sb3/play_sb3.py b/scripts/reinforcement_learning/sb3/play_sb3.py index 87af5950d24a..fb5faebaf8fe 100644 --- a/scripts/reinforcement_learning/sb3/play_sb3.py +++ b/scripts/reinforcement_learning/sb3/play_sb3.py @@ -15,7 +15,7 @@ import gymnasium as gym import torch -from common import CHECKPOINT_SELECTORS, resolve_checkpoint_selector +from common import CHECKPOINT_SELECTORS, add_frontend_args, create_isaaclab_env, resolve_checkpoint_selector from stable_baselines3 import PPO from stable_baselines3.common.vec_env import VecNormalize @@ -46,6 +46,7 @@ ) 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.") +add_frontend_args(parser) parser.add_argument( "--agent", type=str, default="sb3_cfg_entry_point", help="Name of the RL agent configuration entry point." ) @@ -122,17 +123,16 @@ def main(): env_cfg.log_dir = log_dir # create isaac environment - env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + env = create_isaaclab_env( + args_cli.task, + env_cfg, + args_cli, + convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg), + ) # post-process agent configuration agent_cfg = process_sb3_cfg(agent_cfg, env.unwrapped.num_envs) - # convert to single-agent instance if required by the RL algorithm - if isinstance(env.unwrapped.cfg, DirectMARLEnvCfg): - from isaaclab.envs import multi_agent_to_single_agent - - env = multi_agent_to_single_agent(env) - # wrap for video recording if args_cli.video: video_kwargs = { diff --git a/scripts/reinforcement_learning/skrl/play_skrl.py b/scripts/reinforcement_learning/skrl/play_skrl.py index 49659285dca6..0544adb0312d 100644 --- a/scripts/reinforcement_learning/skrl/play_skrl.py +++ b/scripts/reinforcement_learning/skrl/play_skrl.py @@ -20,7 +20,7 @@ import gymnasium as gym import skrl import torch -from common import CHECKPOINT_SELECTORS, resolve_checkpoint_selector +from common import CHECKPOINT_SELECTORS, add_frontend_args, create_isaaclab_env, resolve_checkpoint_selector from packaging import version from isaaclab.app import add_launcher_args, launch_simulation @@ -51,6 +51,7 @@ ) 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.") +add_frontend_args(parser) parser.add_argument( "--agent", type=str, @@ -174,13 +175,12 @@ def main(): env_cfg.log_dir = log_dir # create isaac environment - env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) - - # convert to single-agent instance if required by the RL algorithm - if isinstance(env.unwrapped.cfg, DirectMARLEnvCfg) and algorithm in ["ppo"]: - from isaaclab.envs import multi_agent_to_single_agent - - env = multi_agent_to_single_agent(env) + env = create_isaaclab_env( + args_cli.task, + env_cfg, + args_cli, + convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg) and algorithm in ["ppo"], + ) # get environment (step) dt for real-time evaluation try: diff --git a/source/isaaclab/test/test_reinforcement_learning_common.py b/source/isaaclab/test/test_reinforcement_learning_common.py index 9e2b4c24e701..5085f156c43b 100644 --- a/source/isaaclab/test/test_reinforcement_learning_common.py +++ b/source/isaaclab/test/test_reinforcement_learning_common.py @@ -232,6 +232,20 @@ def test_common_train_args_register_frontend_with_torch_default() -> None: parser.parse_args(["--frontend", "tensorflow"]) +def test_play_entrypoints_route_through_frontend_factory() -> None: + """Every dispatched play backend constructs its env via the frontend-aware factory.""" + play_scripts = sorted((_repo_root() / "scripts" / "reinforcement_learning").glob("*/play_*.py")) + # rlinf constructs environments inside the external framework; the frontend cannot + # reach it (documented limitation). + play_scripts = [path for path in play_scripts if path.name != "play_rlinf.py"] + assert len(play_scripts) == 4, sorted(path.name for path in play_scripts) + for script in play_scripts: + source = script.read_text() + assert "create_isaaclab_env(" in source, f"{script.name} bypasses the frontend factory" + assert "gym.make(args_cli.task" not in source, f"{script.name} constructs directly via gym.make" + assert "add_frontend_args(parser)" in source, f"{script.name} does not expose --frontend" + + def test_create_isaaclab_env_uses_registered_torch_env_by_default(monkeypatch: pytest.MonkeyPatch) -> None: """The shared factory preserves the existing Gym path when no frontend is selected.""" expected_env = object() diff --git a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst index de5082942259..c5a43cbde658 100644 --- a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst +++ b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst @@ -2,8 +2,10 @@ Added ^^^^^ * Added ``--frontend {torch,warp}`` to the shared reinforcement-learning - training CLI (all supported RL libraries) for selecting the environment - runtime; default ``torch`` is unchanged. + training and play CLIs (all supported RL libraries) for selecting the + environment runtime; default ``torch`` is unchanged. The ``rlinf`` + integration constructs environments inside the external framework and is + not frontend-routable. * Added :mod:`isaaclab_experimental.envs.frontend` runtime selector and :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable` used by ``--frontend=warp`` to adapt stable cfgs onto the warp runtime. From 923b88cbcb2c81e1c4bd8cba90b7a5a09e33efa2 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 22 Jul 2026 01:07:27 -0700 Subject: [PATCH 51/58] Address second review: fixes, structural sync gate, capturability API Fix the action-parity mock for origin-less tests, guard the clone-origin warp view against reallocation, accept non-contiguous terrain origins, and pin the reset quaternion composition against the stable math reference. Replace the mask-boundary marker comments with a structural gate: host syncs live in sanctioned functions listed exactly in the gate test, ISAACLAB_SYNC_DEBUG runs warp stages under torch's sync-debug trap as the runtime net, and every @WarpCapturable(False) opt-out is pinned by an inventory test. Consolidate capturability to the class-aware WarpCapturable decorator. The new gate immediately caught a per-step host sync in CircularBuffer.max_length (read on the append path), fixed by caching the host int. Classic-PhysX env_mask behavior is documented on InteractiveScene.reset instead of guarded. --- AGENTS.md | 34 ++- .../isaaclab/scene/interactive_scene.py | 8 +- .../isaaclab/terrains/terrain_importer.py | 2 +- .../envs/manager_based_env_warp.py | 5 +- .../envs/manager_based_rl_env_warp.py | 8 +- .../isaaclab_experimental/envs/mdp/events.py | 7 +- .../envs/mdp/observations.py | 4 +- .../managers/curriculum_manager.py | 2 +- .../managers/manager_base.py | 8 +- .../utils/buffers/circular_buffer.py | 5 +- .../utils/warp/__init__.py | 5 +- .../isaaclab_experimental/utils/warp/utils.py | 88 ++++---- .../utils/warp_graph_cache.py | 33 ++- .../test/envs/mdp/parity_helpers.py | 6 +- .../test/envs/mdp/test_events_warp_parity.py | 41 ++++ .../test/managers/test_manager_base.py | 4 +- .../test/utils/test_mask_boundary_gate.py | 211 +++++++++++++++--- .../assets/articulation/articulation.py | 6 +- .../rigid_object_collection.py | 4 +- .../cloner/newton_clone_utils.py | 4 +- .../envs/mdp/actions/newton_ik_actions.py | 2 +- .../isaaclab_newton/sensors/pva/pva.py | 2 +- 22 files changed, 378 insertions(+), 111 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 497448ae4026..854ac7b8b290 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -330,14 +330,30 @@ mask-to-ID compaction (`nonzero()` on device data) is a GPU-to-host synchronizat with a data-dependent shape, which both stalls the pipeline and permanently disqualifies the containing code path from CUDA graph capture. -1. **Compaction only at reviewed boundaries.** Any production `nonzero(` in the - gated source trees must carry a same-line `# mask-boundary: ` marker; - the static gate `source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py` - fails otherwise. Legitimate boundaries are host consumers: recorders, legacy - compatibility fallbacks, init-time tooling, diagnostics, and debug visualization. -2. **Mask-native Torch means elementwise masked ops.** Use `masked_fill_` or +1. **Syncs only inside named, sanctioned boundary helpers.** Host syncs + (`nonzero`, `argwhere`, `item`, `cpu`, `tolist`, `numpy`) on the mask-native + path live in small single-purpose functions listed in + `SANCTIONED_BOUNDARIES` of + `source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py` — no + inline marker comments. The gate scans the experimental production tree in + full and `*mask*`-named functions in core/Newton; the list is exact, so + adding or removing a boundary is a conscious, reviewable test change. + Legitimate boundaries are host consumers: recorders, legacy compatibility + fallbacks, checkpoint/restore, and empty-reset predicates. +2. **Runtime backstop: `ISAACLAB_SYNC_DEBUG=1`.** Warp stages then execute under + `torch.cuda.set_sync_debug_mode("error")` (see `WarpGraphCache`), so a hidden + sync in any term — including future ones — raises at the exact call site. + Zero cost when the variable is unset. +3. **Capturability is annotated with `@WarpCapturable(False, reason=...)`** (the + single decorator form; works on term functions and class terms). Every + opt-out is pinned by the inventory test next to the gate. Unannotated terms + default to capturable and are covered by the runtime backstop. +4. **Mask-native Torch means elementwise masked ops.** Use `masked_fill_` or `torch.where`; boolean advanced indexing (`buf[:, mask] = 0`) gathers indices with a data-dependent allocation and is not capture-safe. -3. **Dual reset APIs.** Keep the ID-based `reset(env_ids)` signature untouched and - add a sibling `reset_mask(env_mask)`; the base-class fallback compacts (and is - the marked boundary), mask-native implementations override it. +5. **Dual reset APIs.** Keep the ID-based `reset(env_ids)` signature untouched and + add a sibling `reset_mask(env_mask)`; the base-class fallback compacts (a + sanctioned boundary), mask-native implementations override it. Backends whose + assets accept but ignore `env_mask` (classic PhysX, a develop-era condition) + are a documented hazard on `InteractiveScene.reset`, not guarded machinery — + the warp env only drives Newton scenes. diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index 538cbc19d33a..b8ca77aa1b31 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -387,7 +387,7 @@ def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): """ if self._terrain is not None: return self._terrain.env_origins_wp - if self._clone_origins_wp is None: + if self._clone_origins_wp is None or self._clone_origins_wp.ptr != self.env_origins.data_ptr(): self._clone_origins_wp = wp.from_torch(self.env_origins, dtype=wp.vec3f) return self._clone_origins_wp @@ -483,6 +483,12 @@ def reset( env_mask: Boolean Warp mask selecting environments. When provided, it takes precedence over :paramref:`env_ids`. + .. caution:: + ``env_mask`` is honored by mask-native assets (the Newton backend). + Classic PhysX assets currently accept and ignore it, so passing a + mask in such scenes resets unintended environments — use + :paramref:`env_ids` there. + """ reset_kwargs = {"env_mask": env_mask} if env_mask is not None else {"env_ids": env_ids} # -- assets diff --git a/source/isaaclab/isaaclab/terrains/terrain_importer.py b/source/isaaclab/isaaclab/terrains/terrain_importer.py index 5f9fee1d4d3a..7ebfa58d4e5a 100644 --- a/source/isaaclab/isaaclab/terrains/terrain_importer.py +++ b/source/isaaclab/isaaclab/terrains/terrain_importer.py @@ -349,7 +349,7 @@ def configure_env_origins(self, origins: np.ndarray | torch.Tensor | None = None origins = torch.from_numpy(origins) # store the origins self.terrain_origins = self._assign_pointer_stable( - self.terrain_origins, origins.to(self.device, dtype=torch.float) + self.terrain_origins, origins.to(self.device, dtype=torch.float).contiguous() ) # compute environment origins self.env_origins = self._assign_pointer_stable( diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py index 55afdb1df06b..38c82ddf0fe0 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py @@ -433,8 +433,7 @@ def reset( if self._has_recorders: recorder_env_ids = env_ids if recorder_env_ids is None or isinstance(recorder_env_ids, wp.array): - torch_mask = wp.to_torch(reset_mask) - recorder_env_ids = torch_mask.nonzero(as_tuple=False).squeeze(-1) # mask-boundary: recorders + recorder_env_ids = wp.to_torch(reset_mask).nonzero(as_tuple=False).squeeze(-1) self.recorder_manager.record_pre_reset(recorder_env_ids) # set the seed @@ -507,7 +506,7 @@ def reset_to( """ # reset all envs in the scene if env_ids is None if env_mask is not None: - env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) # mask-boundary: host state dict + env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) elif env_ids is None: env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index 252fcb86948d..85390a074577 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -29,7 +29,7 @@ from isaaclab_experimental.managers import CommandManager, CurriculumManager, RewardManager, TerminationManager from isaaclab_experimental.utils.torch_utils import clone_obs_buffer -from isaaclab_experimental.utils.warp import increment_all_int64, zero_masked_int64 +from isaaclab_experimental.utils.warp import any_env_set, increment_all_int64, zero_masked_int64 from .manager_based_env_warp import ManagerBasedEnvWarp @@ -477,7 +477,7 @@ def _reset_mask( curriculum_kwargs: dict[str, Any] = {"env_mask": env_mask} if self.curriculum_manager.requires_host_ids: if env_ids is None: - env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) # mask-boundary: legacy terms + env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) curriculum_kwargs["env_ids"] = env_ids self._warp_graph_cache.call( @@ -577,7 +577,7 @@ def _reset_terminated_envs(self) -> None: reset_mask = self.termination_manager.dones_wp # Keep the mask as the canonical selection, but use one host predicate # to avoid dispatching the complete reset pipeline when it is empty. - if not self.reset_buf.any().item(): + if not any_env_set(self.reset_buf): return # Same-step autoreset exposes terminal observations before any selected @@ -593,7 +593,7 @@ def _reset_terminated_envs(self) -> None: enable=DEBUG_TIMER_STEP, time_unit="us", ): - recorder_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) # mask-boundary: recorders + recorder_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) self.recorder_manager.record_pre_reset(recorder_env_ids) with Timer( diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py index c3349d485efb..faadee5f1770 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py @@ -37,7 +37,7 @@ from isaaclab_experimental.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg from isaaclab_experimental.managers import ManagerTermBase as _WarpManagerTermBase -from isaaclab_experimental.utils.warp import WarpCapturable, warp_capturable +from isaaclab_experimental.utils.warp import WarpCapturable __all__ = [ "apply_external_force_torque", @@ -95,7 +95,7 @@ def _randomize_com_kernel( rng_state[env_id] = state -@warp_capturable(False) +@WarpCapturable(False, reason="set_coms_mask calls SimulationManager.add_model_change") class randomize_rigid_body_com(ManagerTermBase): """Randomize rigid-body centers of mass from a persistent default baseline. @@ -122,7 +122,6 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv) -> None: self._com_lo = wp.vec3f(ranges[0][0], ranges[1][0], ranges[2][0]) self._com_hi = wp.vec3f(ranges[0][1], ranges[1][1], ranges[2][1]) - @WarpCapturable(False, reason="set_coms_mask calls SimulationManager.add_model_change") def __call__( self, env: ManagerBasedEnv, @@ -732,7 +731,7 @@ def reset_joints_by_scale( def _mask_to_env_ids(env_mask: wp.array) -> torch.Tensor: """Convert a Warp boolean env-mask to the torch index tensor the stable terms expect.""" - return torch.nonzero(wp.to_torch(env_mask), as_tuple=False).squeeze(-1) # mask-boundary: startup event twins + return torch.nonzero(wp.to_torch(env_mask), as_tuple=False).squeeze(-1) class randomize_rigid_body_material(_StableRandomizeRigidBodyMaterial, _WarpManagerTermBase): diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py index 8aa4fedfd9b5..d534d70dbcec 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py @@ -36,7 +36,7 @@ record_shape, ) from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_experimental.utils.warp import warp_capturable +from isaaclab_experimental.utils.warp import WarpCapturable if TYPE_CHECKING: from isaaclab.assets import Articulation @@ -449,7 +449,7 @@ def _height_scan_kernel( # Sensor reads go through the lazy-update path, whose host-side timestamp bookkeeping # decides when rays are re-cast and has not been audited for graph capture. -@warp_capturable(False) +@WarpCapturable(False, reason="sensor lazy-update host bookkeeping is not capture-audited (Part 3 scope)") @generic_io_descriptor_warp(units="m", out_dim="sensor:rays", observation_type="SensorState", on_inspect=[record_shape]) def height_scan(env: ManagerBasedEnv, out, sensor_cfg: SceneEntityCfg, offset: float = 0.5) -> None: """Height scan from the given sensor w.r.t. the sensor's frame [m]. Writes into ``out``. diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py index 950496451c4f..3e5aea6cc4b0 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py @@ -254,4 +254,4 @@ def _resolve_legacy_term_cfg(self, term_name: str, term_cfg: StableCurriculumTer @staticmethod def _compact_legacy_env_ids(env_mask: wp.array(dtype=wp.bool)) -> torch.Tensor: """Materialize compact Torch IDs at the legacy curriculum boundary.""" - return wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) # mask-boundary: legacy curriculum terms + return wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py index 1e13f77377f2..bb2decbd77e8 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py @@ -29,7 +29,7 @@ from isaaclab.managers.manager_term_cfg import ManagerTermBaseCfg from isaaclab.utils import class_to_dict, string_to_callable -from isaaclab_experimental.utils.warp import is_warp_capturable +from isaaclab_experimental.utils.warp import WarpCapturable from .scene_entity_cfg import SceneEntityCfg @@ -463,7 +463,11 @@ def _resolve_common_term_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg, def _register_term_capturability(self, term: Callable) -> None: """Keep the complete manager eager when a configured term is unsafe.""" - if not is_warp_capturable(term): + # TODO(#6611): granularity is whole-manager — one non-capturable term forces the + # entire manager eager. Per-term record/replay could keep the capturable terms + # on graphs while only the unsafe term runs eagerly; that is execution-layer + # scope and deliberately not part of the mask-first PR. + if not WarpCapturable.is_capturable(term): graph_cache = getattr(self._env, "_warp_graph_cache", None) if graph_cache is not None: graph_cache.register_capturability(type(self).__name__, False) diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/buffers/circular_buffer.py b/source/isaaclab_experimental/isaaclab_experimental/utils/buffers/circular_buffer.py index f0dcf57b2a30..96bcf21aaf10 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/buffers/circular_buffer.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/buffers/circular_buffer.py @@ -42,6 +42,9 @@ def __init__(self, max_len: int, batch_size: int, device: str): self._device = device self._ALL_INDICES = torch.arange(batch_size, device=device) + # max length as a host int: the property is read on the per-step append path, + # so it must not synchronize with the device tensor below. + self._max_len_int = int(max_len) # max length tensor for comparisons self._max_len = torch.full((batch_size,), max_len, dtype=torch.int, device=device) # number of data pushes passed since the last call to :meth:`reset` @@ -69,7 +72,7 @@ def device(self) -> str: @property def max_length(self) -> int: """The maximum length of the ring buffer.""" - return int(self._max_len[0].item()) + return self._max_len_int @property def current_length(self) -> torch.Tensor: diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py index 27e117d43478..c7e33acf47bf 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py @@ -16,9 +16,10 @@ zero_masked_int64, ) from .utils import ( + SYNC_DEBUG_ENV_VAR, WarpCapturable, - is_warp_capturable, - warp_capturable, + any_env_set, + sync_debug_enabled, wrap_to_pi, zero_masked_2d, ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py index 6b8aafc415f8..87c4833f89e9 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py @@ -5,38 +5,36 @@ from __future__ import annotations -import warp as wp - - -def warp_capturable(capturable: bool): - """Annotate an MDP term's CUDA-graph capturability. +import os - No-wrapper decorator: sets ``_warp_capturable`` directly on the function - and returns it unchanged. Safe to stack with any other decorator in any order. +import torch +import warp as wp - By default all MDP terms are assumed capturable (True). Use - ``@warp_capturable(False)`` on terms that call non-capturable external APIs. - """ +SYNC_DEBUG_ENV_VAR = "ISAACLAB_SYNC_DEBUG" +"""Set ``ISAACLAB_SYNC_DEBUG=1`` to trap hidden GPU->host syncs in warp stages (CI/debug).""" - def decorator(func): - func._warp_capturable = capturable - return func - return decorator +def sync_debug_enabled() -> bool: + """Whether the sync-debug trap is requested for this process.""" + return os.environ.get(SYNC_DEBUG_ENV_VAR, "0") == "1" -def is_warp_capturable(func) -> bool: - """Check if a term function is CUDA-graph-capturable. +def any_env_set(mask: torch.Tensor) -> bool: + """Host predicate: does the boolean mask select any environment? - Checks ``_warp_capturable`` on the function and its ``__wrapped__`` target. - Returns True (capturable) by default if no annotation is found. + This is the single sanctioned per-step host sync of the mask-native reset + pipeline (one predicate instead of dispatching an empty reset). It suspends + the sync-debug trap around itself so the exemption is code, not annotation. """ - for f in (func, getattr(func, "__wrapped__", None)): - if f is not None: - val = getattr(f, "_warp_capturable", None) - if val is not None: - return val - return True + if mask.is_cuda: + prev = torch.cuda.get_sync_debug_mode() + if prev != 0: + torch.cuda.set_sync_debug_mode(0) + try: + return bool(mask.any().item()) + finally: + torch.cuda.set_sync_debug_mode(prev) + return bool(mask.any().item()) @wp.func @@ -52,32 +50,40 @@ def wrap_to_pi(angle: float) -> float: class WarpCapturable: """CUDA graph capture safety: decorator, annotation checker, and runtime guard. - Decorator usage:: - - @WarpCapturable(False) - def reset_root_state_uniform(env, env_mask, ...): - ... + The single capturability annotation (see ``is_capturable``). Decorator usage:: @WarpCapturable(False, reason="calls write_root_pose_to_sim") def push_by_setting_velocity(env, env_mask, ...): ... + @WarpCapturable(False, reason="notifies the solver of model changes") + class randomize_rigid_body_com(ManagerTermBase): + ... + - ``@WarpCapturable(True)`` or no decorator: capturable, returned unwrapped. - - ``@WarpCapturable(False)``: sets ``func._warp_capturable = False``, wraps with - runtime guard that raises if ``wp.get_device().is_capturing`` is ``True``. + - ``@WarpCapturable(False)`` on a function: sets ``_warp_capturable = False`` and + wraps it with a runtime guard that raises if called during CUDA graph capture. + - ``@WarpCapturable(False)`` on a class term: annotates the class (so manager + registration sees it) and guards its ``__call__``. """ def __init__(self, capturable: bool, *, reason: str | None = None): self._capturable = capturable self._reason = reason - def __call__(self, func): - """Decorate *func* with capture safety annotation and optional runtime guard.""" - import functools - - func._warp_capturable = self._capturable + def __call__(self, target): + """Decorate a term function or class with the annotation and optional guard.""" + target._warp_capturable = self._capturable if self._capturable: - return func + return target + if isinstance(target, type): + target.__call__ = self._guarded(target.__call__) + return target + return self._guarded(target) + + def _guarded(self, func): + """Wrap *func* so calling it during CUDA graph capture raises.""" + import functools reason = self._reason @@ -94,12 +100,14 @@ def wrapper(*args, **kwargs): return wrapper @staticmethod - def is_capturable(func) -> bool: - """Check capturability annotation. Default: ``True``. + def is_capturable(term) -> bool: + """Check the capturability annotation on a term function, class, or instance. Checks ``__wrapped__`` for decorated functions to handle stacked decorators. + Unannotated terms default to capturable; hidden syncs in them are trapped by + the ``ISAACLAB_SYNC_DEBUG`` stage sweep and by CUDA graph capture itself. """ - for f in (func, getattr(func, "__wrapped__", None)): + for f in (term, getattr(term, "__wrapped__", None)): if f is not None: val = getattr(f, "_warp_capturable", None) if val is not None: diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py index ee54c0e59121..0bdb00d02c45 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py @@ -9,10 +9,35 @@ from collections.abc import Callable from typing import Any +import torch import warp as wp from isaaclab.utils.timer import Timer +SYNC_DEBUG_ENV_VAR = "ISAACLAB_SYNC_DEBUG" +"""Set ``ISAACLAB_SYNC_DEBUG=1`` to run eager and warm-up stage invocations under +``torch.cuda.set_sync_debug_mode("error")`` so hidden GPU->host syncs raise (CI/debug).""" + +_SYNC_DEBUG = os.environ.get(SYNC_DEBUG_ENV_VAR, "0") == "1" + + +def _invoke_stage(fn: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: + """Run one stage callable, optionally under the sync-debug trap. + + The trap is the choke point that guards every current and future term inside a + warp stage without per-function annotations; opt-outs suspend it explicitly + (see :func:`isaaclab_experimental.utils.warp.any_env_set`). + """ + if not _SYNC_DEBUG: + return fn(*args, **kwargs) + prev = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode(2) + try: + return fn(*args, **kwargs) + finally: + torch.cuda.set_sync_debug_mode(prev) + + CAPTURE_ENV_VAR = "ISAACLAB_WARP_CAPTURE" """Set ``ISAACLAB_WARP_CAPTURE=0`` to force every stage eager (debug / A-B validation).""" @@ -94,14 +119,14 @@ def call( group = stage.partition("_")[0] with Timer(name=stage, msg=f"{stage} took:", enable=timer, time_unit="us"): if not self._capture_enabled or not self.is_capturable(group): - result = fn(*args, **kwargs) + result = _invoke_stage(fn, args, kwargs) elif stage in self._graphs: wp.capture_launch(self._graphs[stage]) result = self._results[stage] elif stage not in self._warmed: # The first real call doubles as warm-up. Stateful manager stages # must execute exactly once per environment step. - result = fn(*args, **kwargs) + result = _invoke_stage(fn, args, kwargs) self._warmed.add(stage) else: result = self._capture(stage, fn, args, kwargs) @@ -133,13 +158,13 @@ def capture_or_replay( if kwargs is None: kwargs = {} if not self._capture_enabled: - return fn(*args, **kwargs) + return _invoke_stage(fn, args, kwargs) graph = self._graphs.get(stage) if graph is not None: wp.capture_launch(graph) return self._results[stage] # Warm-up: run eagerly to flush first-call allocations / hasattr guards. - fn(*args, **kwargs) + _invoke_stage(fn, args, kwargs) return self._capture(stage, fn, args, kwargs) def _capture( diff --git a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py index 7cb1e8b073be..86a352958a15 100644 --- a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py +++ b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py @@ -383,7 +383,11 @@ class MockScene: def __init__(self, assets: dict, env_origins, sensors=None): self._assets = assets - if isinstance(env_origins, wp.array): + if env_origins is None: + # Action-term tests don't need origins; keep the attributes present but empty. + self.env_origins = None + self.env_origins_wp = None + elif isinstance(env_origins, wp.array): self.env_origins_wp = env_origins self.env_origins = wp.to_torch(env_origins) else: diff --git a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py index 480ec829cd97..1b21c97d06d0 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py @@ -345,6 +345,47 @@ def test_apply_external_force_torque(self, warp_env, art_data, all_joints_cfg): # -- reset_root_state_uniform ----------------------------------------------- + def test_reset_root_state_uniform_composes_orientation_like_stable(self): + """Pin the quaternion path: ``final = default ∘ euler_xyz(roll, pitch, yaw)``. + + Zero-width ranges at non-zero angles make the uniform draw deterministic, + so the composition order, Euler convention, and xyzw storage layout are + all pinned against the stable math reference. + """ + import isaaclab.utils.math as math_utils + + env, data, asset = _make_event_env(1234) + # Non-identity default orientation (90 deg about x), xyzw layout. + default_pose = np.zeros((NUM_ENVS, 7), dtype=np.float32) + default_pose[:, 3] = np.sin(np.pi / 4.0) + default_pose[:, 6] = np.cos(np.pi / 4.0) + copy_np_to_wp(data.default_root_pose, default_pose) + copy_np_to_wp(data.default_root_vel, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + captured = {} + asset.write_root_pose_to_sim_mask = lambda **kwargs: captured.update(kwargs) + asset.write_root_velocity_to_sim_mask = lambda **kwargs: None + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + roll, pitch, yaw = 0.3, -0.4, 1.1 + term = _make_event_term( + warp_evt.reset_root_state_uniform, + env, + pose_range={"roll": (roll, roll), "pitch": (pitch, pitch), "yaw": (yaw, yaw)}, + velocity_range={}, + ) + + term(env, env_mask, **term.cfg.params) + wp.synchronize() + + result_xyzw = wp.to_torch(captured["root_pose"])[:, 3:7] + angles = torch.tensor([[roll], [pitch], [yaw]], device=DEVICE) + # Core math utils use the (x, y, z, w) layout, matching the warp buffers. + delta_xyzw = math_utils.quat_from_euler_xyz(angles[0], angles[1], angles[2]) + default_xyzw = torch.tensor([[np.sin(np.pi / 4.0), 0.0, 0.0, np.cos(np.pi / 4.0)]], device=DEVICE) + expected_xyzw = math_utils.quat_mul(default_xyzw, delta_xyzw).expand(NUM_ENVS, 4) + # Quaternions double-cover rotations: align signs before comparing. + sign = torch.sign((result_xyzw * expected_xyzw).sum(dim=1, keepdim=True)) + assert_close(result_xyzw * sign, expected_xyzw) + # -- env_mask selectivity --------------------------------------------------- def test_reset_joints_mask_selectivity(self, warp_env, art_data, all_joints_cfg): diff --git a/source/isaaclab_experimental/test/managers/test_manager_base.py b/source/isaaclab_experimental/test/managers/test_manager_base.py index 5450a49835ed..bef4ad902216 100644 --- a/source/isaaclab_experimental/test/managers/test_manager_base.py +++ b/source/isaaclab_experimental/test/managers/test_manager_base.py @@ -12,7 +12,7 @@ import warp as wp from isaaclab_experimental.managers.action_manager import ActionManager from isaaclab_experimental.managers.manager_base import _resolve_reset_mask -from isaaclab_experimental.utils.warp import warp_capturable +from isaaclab_experimental.utils.warp import WarpCapturable def test_reset_mask_rejects_wrong_device(monkeypatch: pytest.MonkeyPatch) -> None: @@ -29,7 +29,7 @@ def test_reset_mask_rejects_wrong_device(monkeypatch: pytest.MonkeyPatch) -> Non def test_class_term_capturability_is_registered() -> None: """Class-based manager terms should honor explicit capture metadata.""" - @warp_capturable(False) + @WarpCapturable(False) class HostTerm: pass diff --git a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py index cc62a94e0a33..caea47f38cb2 100644 --- a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py +++ b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py @@ -3,37 +3,198 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Static gate for the mask-first contract. +"""Static gate for the mask-native execution path. -Host-side mask-to-ID compaction is a synchronization point and must only happen -at explicitly reviewed boundaries (recorders, legacy term/actuator APIs, -init-time tooling, diagnostics). Every such call site carries a same-line -``# mask-boundary: `` marker; this gate fails when a new unmarked site -appears in the scanned production packages. +The warp frontend's contract is that mask-native code performs no data-dependent +GPU->host synchronization outside sanctioned boundaries. Instead of inline marker +comments, the gate is structural: + +* **Scope** — every function in the experimental production tree that sits on the + step/reset pipeline, plus every ``*mask*``-named function in the core and Newton + trees (the functions that *are* the warp path there). Torch-first core code is + deliberately out of scope. +* **Sanctioned boundaries** — small, named helper functions listed in + ``SANCTIONED_BOUNDARIES``; the list is exact (a stale entry fails the gate). +* **Runtime backstop** — ``ISAACLAB_SYNC_DEBUG=1`` runs every warp stage under + ``torch.cuda.set_sync_debug_mode("error")`` (see ``WarpGraphCache``), trapping + syncs this static scan cannot see. + +The companion inventory test pins every ``@WarpCapturable(False)`` opt-out so +non-capturable terms are documented in exactly one reviewable place. """ +from __future__ import annotations + +import ast from pathlib import Path _REPO_ROOT = Path(__file__).resolve().parents[4] -_SCAN_ROOTS = ( - "source/isaaclab_experimental/isaaclab_experimental", - "source/isaaclab_newton/isaaclab_newton", + +# Attribute calls that synchronize (or copy) device memory to the host. +SYNC_ATTRS = {"nonzero", "argwhere", "item", "cpu", "tolist", "numpy"} + +# Trees scanned in full (production only; every function is in scope unless excluded). +FULL_SCAN_ROOTS = [ + "source/isaaclab_experimental/isaaclab_experimental/envs", + "source/isaaclab_experimental/isaaclab_experimental/managers", + "source/isaaclab_experimental/isaaclab_experimental/utils", +] + +# Trees where only ``*mask*``-named functions are in scope (the warp path inside +# otherwise torch-first packages). +MASK_SCAN_ROOTS = [ + "source/isaaclab/isaaclab/actuators", "source/isaaclab/isaaclab/scene", + "source/isaaclab/isaaclab/sensors", "source/isaaclab/isaaclab/terrains", -) -_COMPACTION_PATTERN = "nonzero(" -_BOUNDARY_MARKER = "mask-boundary:" - - -def test_host_compaction_sites_carry_mask_boundary_markers(): - """Every production mask-to-ID compaction must be a marked, reviewed boundary.""" - violations = [] - for scan_root in _SCAN_ROOTS: - for path in sorted((_REPO_ROOT / scan_root).rglob("*.py")): - for line_number, line in enumerate(path.read_text().splitlines(), start=1): - if _COMPACTION_PATTERN in line and _BOUNDARY_MARKER not in line: - violations.append(f"{path.relative_to(_REPO_ROOT)}:{line_number}: {line.strip()}") - assert not violations, ( - "Host mask-to-ID compaction found outside marked boundaries. Make the call mask-native, or" - " append '# mask-boundary: ' after review:\n" + "\n".join(violations) + "source/isaaclab_newton/isaaclab_newton", +] + +# Functions that never run on the per-step path: construction, reporting, and +# debug surfaces. Host syncs there are init-time or human-facing. +NON_STEP_FUNCTIONS = { + "__init__", + "__post_init__", + "__str__", + "__repr__", + "__del__", + "_prepare_terms", + "serialize", + "get_active_iterable_terms", + "_debug_vis_callback", + "_set_debug_vis_impl", + # reporting / config-access surfaces (export- or user-facing, not per-step) + "IO_descriptor", + "get_IO_descriptors", + "get_term_cfg", +} + +# The sanctioned mask->ID / host-sync boundaries, by (path suffix, function name). +# Every entry must exist and contain a sync call; anything else that syncs fails. +SANCTIONED_BOUNDARIES = { + ("managers/curriculum_manager.py", "_compact_legacy_env_ids"), + ("envs/mdp/events.py", "_mask_to_env_ids"), + ("envs/manager_based_env_warp.py", "reset_to"), + # Coarse-grained by review preference: the recorder / legacy-term compactions + # stay inline in their reset functions rather than one-line helpers; the + # ISAACLAB_SYNC_DEBUG runtime trap is the fine-grained net inside them. + ("envs/manager_based_env_warp.py", "reset"), + ("envs/manager_based_rl_env_warp.py", "_reset_terminated_envs"), + ("envs/manager_based_rl_env_warp.py", "_reset_mask"), + ("utils/warp/utils.py", "any_env_set"), + # camera's empty-reset predicate (the sensor-side analogue of any_env_set) + ("isaaclab/sensors/camera/camera.py", "_env_mask_has_any"), + # joint-limit writes mutate the solver model (host-side by nature, event-driven) + ("isaaclab_newton/assets/articulation/articulation.py", "write_joint_position_limit_to_sim_mask"), + # develop-era compaction helpers; deleted by Part 3 (#6646) — its rebase must + # drop these entries (the stale-entry assertion enforces that). + ("isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py", "_resolve_env_mask"), + ("isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py", "_resolve_body_mask"), +} + + +def _sync_calls(func_node: ast.AST) -> list[str]: + """Names of synchronizing attribute calls inside a function body.""" + hits = [] + for node in ast.walk(func_node): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr in SYNC_ATTRS: + hits.append(node.func.attr) + return hits + + +def _iter_functions(tree: ast.Module): + """Yield (function node, name) for every def in the module.""" + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + yield node, node.name + + +def _scan(root: Path, mask_only: bool) -> dict[tuple[str, str], list[str]]: + """Map (relative path, function) -> sync hits for in-scope functions under root.""" + findings: dict[tuple[str, str], list[str]] = {} + for path in sorted(root.rglob("*.py")): + rel = path.relative_to(_REPO_ROOT).as_posix() + if "/test/" in rel: + continue + tree = ast.parse(path.read_text(), filename=rel) + for func, name in _iter_functions(tree): + if name in NON_STEP_FUNCTIONS: + continue + if mask_only and "mask" not in name: + continue + hits = _sync_calls(func) + if hits: + findings[(rel, name)] = hits + return findings + + +def _is_sanctioned(rel: str, name: str) -> bool: + return any(rel.endswith(suffix) and name == func for suffix, func in SANCTIONED_BOUNDARIES) + + +def test_mask_native_code_has_no_unsanctioned_host_syncs(): + """Every sync in mask-native scope is one of the named, sanctioned boundaries.""" + findings: dict[tuple[str, str], list[str]] = {} + for root in FULL_SCAN_ROOTS: + findings.update(_scan(_REPO_ROOT / root, mask_only=False)) + for root in MASK_SCAN_ROOTS: + findings.update(_scan(_REPO_ROOT / root, mask_only=True)) + + violations = {k: v for k, v in findings.items() if not _is_sanctioned(*k)} + assert not violations, "Unsanctioned host syncs on the mask-native path:\n" + "\n".join( + f" {rel}::{name} -> {hits}" for (rel, name), hits in sorted(violations.items()) + ) + + # The sanctioned list is exact: every entry must still exist and still sync. + matched = {k for k in findings if _is_sanctioned(*k)} + stale = { + (suffix, func) + for suffix, func in SANCTIONED_BOUNDARIES + if not any(rel.endswith(suffix) and name == func for rel, name in matched) + } + assert not stale, f"Stale sanctioned-boundary entries (function gone or no longer syncs): {sorted(stale)}" + + +# Expected ``@WarpCapturable(False)`` opt-outs: (path suffix, decorated name). +EXPECTED_NON_CAPTURABLE = { + ("envs/mdp/events.py", "randomize_rigid_body_com"), + ("envs/mdp/observations.py", "height_scan"), +} + + +def _warp_capturable_false_targets(tree: ast.Module) -> list[str]: + """Names decorated with ``@WarpCapturable(False, ...)`` in a module.""" + out = [] + for node in ast.walk(tree): + if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + continue + for deco in node.decorator_list: + if ( + isinstance(deco, ast.Call) + and isinstance(deco.func, ast.Name) + and deco.func.id == "WarpCapturable" + and deco.args + and isinstance(deco.args[0], ast.Constant) + and deco.args[0].value is False + ): + out.append(node.name) + return out + + +def test_non_capturable_terms_are_inventoried(): + """Every ``@WarpCapturable(False)`` opt-out is documented here, and only these.""" + found = set() + root = _REPO_ROOT / "source/isaaclab_experimental/isaaclab_experimental" + for path in sorted(root.rglob("*.py")): + rel = path.relative_to(_REPO_ROOT).as_posix() + for name in _warp_capturable_false_targets(ast.parse(path.read_text(), filename=rel)): + found.add((rel, name)) + expected = { + (f"source/isaaclab_experimental/isaaclab_experimental/{suffix}", name) + for suffix, name in EXPECTED_NON_CAPTURABLE + } + assert found == expected, ( + "Non-capturable inventory drift.\n" + f" unexpected: {sorted(found - expected)}\n" + f" missing: {sorted(expected - found)}" ) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 76246171593a..0bf09ce4a210 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -279,7 +279,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None actuator_env_ids = env_ids if self.actuators and env_mask is not None: torch_mask = wp.to_torch(env_mask) - actuator_env_ids = torch_mask.nonzero(as_tuple=False).squeeze(-1) # mask-boundary: legacy actuators + actuator_env_ids = torch_mask.nonzero(as_tuple=False).squeeze(-1) for actuator in self.actuators.values(): actuator.reset(actuator_env_ids) # reset the global Newton actuator adapter (its ``_states_a/_b`` buffers @@ -3966,7 +3966,7 @@ def _validate_cfg(self): default_joint_pos = self._data.default_joint_pos.torch[0] out_of_range = default_joint_pos < joint_pos_limits_lower out_of_range |= default_joint_pos > joint_pos_limits_upper - violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) # mask-boundary: diagnostics + violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) # throw error if any of the default joint positions are out of the limits if len(violated_indices) > 0: # prepare message for violated joints @@ -3983,7 +3983,7 @@ def _validate_cfg(self): joint_max_vel = self._data.joint_vel_limits.torch[0] default_joint_vel = self._data.default_joint_vel.torch[0] out_of_range = torch.abs(default_joint_vel) > joint_max_vel - violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) # mask-boundary: diagnostics + violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) if len(violated_indices) > 0: # prepare message for violated joints msg = "The following joints have default velocities out of the limits: \n" diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py index 7f6c523eb391..e0eb6276d5ec 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py @@ -1284,7 +1284,7 @@ def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array | torch.Tenso if env_mask is not None: if isinstance(env_mask, wp.array): env_mask = wp.to_torch(env_mask) - env_ids = torch.nonzero(env_mask)[:, 0].to(torch.int32) # mask-boundary: legacy index write path + env_ids = torch.nonzero(env_mask)[:, 0].to(torch.int32) else: env_ids = self._ALL_ENV_INDICES return env_ids @@ -1294,7 +1294,7 @@ def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array | torch.Ten if body_mask is not None: if isinstance(body_mask, wp.array): body_mask = wp.to_torch(body_mask) - body_ids = torch.nonzero(body_mask)[:, 0].to(torch.int32) # mask-boundary: legacy index write path + body_ids = torch.nonzero(body_mask)[:, 0].to(torch.int32) else: body_ids = self._ALL_BODY_INDICES return body_ids diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py index e411f7550eef..586746744398 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py @@ -177,7 +177,7 @@ def replicate_builder_mapping( site_idx = builder.add_site(body=-1, xform=wp.transform_multiply(world_xform, xform), label=label) local_site_map.setdefault(label, [[] for _ in range(num_worlds)])[col].append(site_idx) - for row in torch.nonzero(mapping[:, col], as_tuple=True)[0].tolist(): # mask-boundary: init-time cloning + for row in torch.nonzero(mapping[:, col], as_tuple=True)[0].tolist(): source_builder = source_builders[sources[int(row)]] offset = builder.shape_count source_col = int(source_world_indices[int(row)]) @@ -222,7 +222,7 @@ def rename_builder_labels( for source_index, source in enumerate(sources): source_root = source.rstrip("/") - world_cols = torch.nonzero(mapping[source_index], as_tuple=True)[0].tolist() # mask-boundary: init-time cloning + world_cols = torch.nonzero(mapping[source_index], as_tuple=True)[0].tolist() world_roots = {int(env_ids[col]): destinations[source_index].format(int(env_ids[col])) for col in world_cols} def _rename_pair(values, worlds, *, collect_body_bindings: bool = False): diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py index 715a5197b4ec..5a91f113111a 100644 --- a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py @@ -335,7 +335,7 @@ def _validate_matching_root_orientations(self) -> None: # q and -q represent the same orientation, so compare absolute dot products. same_orientation = torch.abs(torch.sum(root_quat_w * root_quat_w[0:1], dim=-1)) > 1.0 - 1e-5 if not torch.all(same_orientation): - bad_env_ids = torch.nonzero(~same_orientation, as_tuple=False).flatten().tolist() # mask-boundary: errors + bad_env_ids = torch.nonzero(~same_orientation, as_tuple=False).flatten().tolist() raise RuntimeError( "NewtonInverseKinematicsAction solves against the env 0 prototype root orientation, but " f"root orientations differ in env ids {bad_env_ids}. Use identical fixed-base root orientations " diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py b/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py index 3a82e8540738..2be89399d03a 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py @@ -220,7 +220,7 @@ def _debug_vis_callback(self, event): pos_w_torch = self._data.pos_w.torch accel_w = math_utils.quat_apply(self._data.quat_w.torch, self._data.lin_acc_b.torch) accel_valid = torch.linalg.norm(accel_w, dim=-1) > 1e-5 - valid_indices = accel_valid.nonzero(as_tuple=True)[0] # mask-boundary: debug vis + valid_indices = accel_valid.nonzero(as_tuple=True)[0] if valid_indices.numel() == 0: return pos_filtered = pos_w_torch.index_select(0, valid_indices) From e8a21985443c1c5e80963d174ce18cb460b8a776 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 22 Jul 2026 02:08:22 -0700 Subject: [PATCH 52/58] Convert observation noise cfgs to warp twins in the bridge MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Warp observation manager applies noise through the experimental noise cfg family, so stable noise cfgs on bridged tasks were silently skipped and training ran without the configured corruption. The frontend now swaps stable constant/uniform/gaussian noise cfgs to their warp-native twins during cfg adaptation (data fields copied, warp func default); cfgs it cannot represent — class-based noise models, customized funcs, tensor-valued parameters — are an aggregated hard error instead of a silent MDP change. The observation manager rejects non-warp noise cfgs at construction and plumbs the shared per-env RNG state to function-style noise kernels, wiring the contract the noise module documented. Audit: every function-style noise cfg in the repo has a swappable warp twin; class-based models are loud unsupported. --- .../changelog.d/warp-manager-bridge.rst | 6 ++ .../isaaclab_experimental/envs/frontend.py | 57 +++++++++++++++++++ .../managers/observation_manager.py | 14 +++++ .../test/envs/test_frontend_cfg_conversion.py | 50 ++++++++++++++++ 4 files changed, 127 insertions(+) diff --git a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst index c5a43cbde658..5f47328da3b0 100644 --- a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst +++ b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst @@ -12,6 +12,12 @@ Added * Added ``warp_entry_point`` registration support: a stable direct task may declare its warp environment class, which ``--frontend=warp`` constructs with the stable configuration. +* Added observation-noise conversion to the warp frontend: stable noise cfgs + (constant/uniform/gaussian) swap to their warp-native twins during cfg + adaptation; cfgs without a twin (class-based noise models, customized or + tensor-valued parameters) are a hard error instead of being silently + ignored. The Warp observation manager rejects non-warp noise cfgs and + plumbs the shared per-env RNG state to function-style noise kernels. * Added warp adapters for the stable :class:`~isaaclab.envs.mdp.events.randomize_rigid_body_material` and :class:`~isaaclab.envs.mdp.events.randomize_rigid_body_mass` startup event diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index 04d9d1886969..cb3ecf101ae7 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -37,6 +37,7 @@ from typing import Any, ClassVar import gymnasium as gym +import torch from isaaclab.envs import DirectMARLEnvCfg, DirectRLEnvCfg, ManagerBasedRLEnvCfg from isaaclab.managers.scene_entity_cfg import SceneEntityCfg as _StableSceneEntityCfg @@ -179,6 +180,7 @@ def adapt_cfg(cls, cfg: Any) -> None: cls._require_newton_physics(cfg, label) cls._promote_scene_entity_cfgs(cfg) cls._swap_mdp(cfg, label) + cls._swap_noise_cfgs(cfg, label) @staticmethod def _require_newton_physics(cfg: Any, label: str) -> None: @@ -401,6 +403,61 @@ def _nearest_mdp_module(mirrored: str) -> str | None: continue return None + @classmethod + def _swap_noise_cfgs(cls, cfg: Any, label: str) -> None: + """Replace stable observation-noise cfgs with their warp-native twins. + + The warp observation manager applies noise through the experimental noise + cfg family (in-place Warp kernels); a stable noise cfg would otherwise be + silently ignored, training against a different MDP. Twins resolve by class + name; data fields are copied while ``func`` keeps the twin's warp default. + Noise cfgs without a warp twin (e.g. class-based ``NoiseModelCfg``) are a + hard error — silently changing the MDP is never acceptable. + """ + import dataclasses + + from isaaclab.utils import noise as stable_noise + + from isaaclab_experimental.utils import noise as warp_noise + + unsupported: list[str] = [] + swapped = 0 + for path, term in cls._walk_terms(cfg): + if not path or path[0] != "observations": + continue + noise_cfg = getattr(term, "noise", None) + if noise_cfg is None: + continue + if type(noise_cfg).__module__.startswith(cls.WARP_ROOT_PREFIXES): + continue # already warp-native + twin_cls = getattr(warp_noise, type(noise_cfg).__name__, None) + if ( + twin_cls is None + or not twin_cls.__module__.startswith(cls.WARP_ROOT_PREFIXES) + or not isinstance(noise_cfg, stable_noise.NoiseCfg) + ): + unsupported.append(f"{'.'.join(path)}: {type(noise_cfg).__name__}") + continue + if noise_cfg.func != type(noise_cfg)().func: + # A user-customized noise func has no warp twin; replacing it with the + # twin's default would silently change the MDP. + unsupported.append(f"{'.'.join(path)}: {type(noise_cfg).__name__} with custom func {noise_cfg.func!r}") + continue + fields = {f.name: getattr(noise_cfg, f.name) for f in dataclasses.fields(noise_cfg) if f.name != "func"} + if any(isinstance(value, torch.Tensor) for value in fields.values()): + # The warp noise kernels take scalar parameters; tensor-valued ranges + # (per-element noise) have no warp twin yet. + unsupported.append(f"{'.'.join(path)}: {type(noise_cfg).__name__} with tensor-valued parameters") + continue + term.noise = twin_cls(**fields) + swapped += 1 + if unsupported: + raise FrontendIncompatibleError( + f"warp env {label!r}: observation noise cfgs without warp twins:\n " + "\n ".join(unsupported) + ) + if swapped: + logger.info("frontend.warp: swapped %d observation noise cfg(s) to warp twins", swapped) + @staticmethod def _import_twin_module(target: str) -> ModuleType: """Import a resolved twin module; a module that exists but breaks is a hard error.""" diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index 46c93f891fa5..8adb74ff8da3 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -679,6 +679,16 @@ def _prepare_terms(self): # noqa: C901 # check noise settings if not group_cfg.enable_corruption: term_cfg.noise = None + elif term_cfg.noise is not None and not type(term_cfg.noise).__module__.startswith( + "isaaclab_experimental." + ): + raise TypeError( + f"Observation term '{term_name}' uses noise cfg" + f" '{type(term_cfg.noise).__module__}.{type(term_cfg.noise).__name__}', which the Warp" + " observation manager would silently ignore. Use the warp-native noise cfgs from" + " isaaclab_experimental.utils.noise (the warp frontend converts stable cfgs" + " automatically)." + ) # check group history params and override terms if group_cfg.history_length is not None: term_cfg.history_length = group_cfg.history_length @@ -760,6 +770,10 @@ def _prepare_terms(self): # noqa: C901 f" and optional parameters: {args_with_defaults}, but received: {term_params}." ) + # plumb the shared per-env RNG state so Warp noise kernels can consume it + # (function-style NoiseCfg kernels read it off the cfg at call time) + if term_cfg.noise is not None and isinstance(term_cfg.noise, noise.NoiseCfg): + term_cfg.noise.rng_state_wp = self._env.rng_state_wp # prepare noise model classes if term_cfg.noise is not None and isinstance(term_cfg.noise, noise.NoiseModelCfg): # plumb the shared per-env RNG state so Warp noise kernels can consume it diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index ea534f3a4b33..f833d032964e 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -169,3 +169,53 @@ def test_stable_cartpole_cfg_adapts_to_current_warp_module_layout(): assert cfg.rewards.pole_pos.func.__module__.startswith("isaaclab_tasks_experimental.core.cartpole.mdp") assert cfg.rewards.success_rate.func.__module__.startswith("isaaclab_tasks_experimental.core.cartpole.mdp") assert issubclass(cfg.actions.joint_effort.class_type, ActionTerm) + + +def _iter_obs_terms(cfg): + """Yield (name, term cfg) for every observation term that carries a noise slot.""" + from isaaclab.managers.manager_term_cfg import ObservationGroupCfg + + for group in vars(cfg.observations).values(): + if not isinstance(group, ObservationGroupCfg): + continue + for name, term in vars(group).items(): + if hasattr(term, "noise"): + yield name, term + + +def test_stable_observation_noise_converts_to_warp_twins(): + """Stable noise cfgs (e.g. Unoise) become warp-native twins with the same parameters.""" + from isaaclab.utils import noise as stable_noise + + entry = _cfg_entry_point("Isaac-Velocity-Flat-UnitreeGo2") + module_path, class_name = entry.split(":") + raw = getattr(importlib.import_module(module_path), class_name)() + stable_params = { + name: (term.noise.n_min, term.noise.n_max) + for name, term in _iter_obs_terms(raw) + if isinstance(term.noise, stable_noise.UniformNoiseCfg) + } + assert stable_params, "expected uniform noise on the stable velocity observations" + + cfg = _load_adapted_cfg(entry) + converted = dict(_iter_obs_terms(cfg)) + for name, (n_min, n_max) in stable_params.items(): + twin = converted[name].noise + assert type(twin).__module__.startswith("isaaclab_experimental."), name + assert (twin.n_min, twin.n_max) == (n_min, n_max), name + + +def test_noise_cfg_without_warp_twin_is_a_hard_error(): + """Class-based noise models have no warp twin yet: adapting must fail loudly.""" + from isaaclab_experimental.envs.frontend import FrontendIncompatibleError + + from isaaclab.utils.noise import NoiseModelCfg, UniformNoiseCfg + + entry = _cfg_entry_point("Isaac-Velocity-Flat-UnitreeGo2") + module_path, class_name = entry.split(":") + cfg = getattr(importlib.import_module(module_path), class_name)() + cfg = resolve_presets(cfg, selected=("newton_mjwarp",)) + name, term = next(iter(_iter_obs_terms(cfg))) + term.noise = NoiseModelCfg(noise_cfg=UniformNoiseCfg()) + with pytest.raises(FrontendIncompatibleError, match="noise"): + WarpFrontend.adapt_cfg(cfg) From daa14b38b8abe2e7d9bb39fede85915cc1c64f9f Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 22 Jul 2026 02:09:03 -0700 Subject: [PATCH 53/58] Adapt curriculum terms opportunistically in the warp frontend Curriculum terms swap to warp twins when one exists, keeping the legacy ID-based fallback for stable terms without twins; the mask-native task classes take their stable functions' names (terrain_levels_vel, modify_reward_weight) so plain name resolution finds them, and the velocity legacy shim is removed with them. The curriculum manager's legacy modes collapse to one (compact IDs, materialized at most once); the unused requires_host_ids cfg knob is removed. Rough-terrain curriculum now runs mask-native and captured; reach's weight curriculum runs mask-native but stays eager by its own declaration (host step counter). The non-capturable inventory detects attribute-style annotations and scans the tasks package. --- .../jichuanh-warp-frontend-cleanup.minor.rst | 3 + .../isaaclab_experimental/envs/frontend.py | 82 +++++++++++++++-- .../managers/curriculum_manager.py | 30 +++---- .../managers/manager_term_cfg.py | 9 +- .../managers/observation_manager.py | 14 +++ .../test/envs/test_frontend.py | 8 +- .../test/envs/test_frontend_cfg_conversion.py | 90 +++++++++++++++++++ .../test/managers/test_curriculum_manager.py | 30 +++---- .../test/utils/test_mask_boundary_gate.py | 43 ++++++--- .../core/reach/mdp/__init__.pyi | 4 +- .../core/reach/mdp/curriculums.py | 2 +- .../core/velocity/mdp/__init__.pyi | 5 +- .../core/velocity/mdp/curriculums.py | 19 +--- 13 files changed, 245 insertions(+), 94 deletions(-) diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst index 98b4d7c8ea20..03bcd096bfc8 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -6,6 +6,9 @@ Added without commands. * Added a Warp-native ``height_scan`` observation term (eager until sensor capture-readiness lands) with ray-count output-dimension inference. +* Added opportunistic curriculum adaptation: curriculum terms swap to warp + twins when one exists (e.g. terrain levels), keeping the legacy ID-based + fallback for stable terms without twins. * Added :class:`~isaaclab_experimental.managers.CurriculumManager` and a boolean-mask reset path for Warp manager-based environments, retaining compact environment IDs only for legacy host consumers. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index d918016ff387..91f65daef0d0 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -37,6 +37,7 @@ from typing import Any, ClassVar import gymnasium as gym +import torch from isaaclab.envs import DirectMARLEnvCfg, DirectRLEnvCfg, ManagerBasedRLEnvCfg from isaaclab.managers.scene_entity_cfg import SceneEntityCfg as _StableSceneEntityCfg @@ -87,15 +88,20 @@ class WarpFrontend: # event func (which expects torch ``env_ids``) breaks at runtime; its funcs # must be swapped to warp twins. The command manager is the warp-native # ``CommandManager`` (captured), so command term ``class_type``s must swap - # to their warp twins as well. The curriculum and recorder managers run on - # the stable (torch) implementation, so their terms are left untouched. A - # stable term left on a warp manager would break, so a missing twin in these - # groups is a hard error; a stable term on a stable manager is correct, so - # those groups are skipped. + # to their warp twins as well. A stable term left on a warp manager would + # break, so a missing twin in these groups is a hard error. The recorder + # manager runs on the stable implementation and is skipped entirely; + # curriculum is adapted opportunistically (see ``OPTIONAL_WARP_GROUPS``). WARP_MANAGED_GROUPS: ClassVar[frozenset[str]] = frozenset( {"observations", "rewards", "terminations", "actions", "events", "commands"} ) + # Groups adapted opportunistically: a warp twin is swapped in when one exists, + # otherwise the stable term stays on the manager's legacy (host-ID) fallback. + # Curriculum is the only such group: its manager supports stable terms by + # design, at the cost of eager execution and one ID materialization per reset. + OPTIONAL_WARP_GROUPS: ClassVar[frozenset[str]] = frozenset({"curriculum"}) + # ------------------------------------------------------------------ # Env construction # ------------------------------------------------------------------ @@ -181,6 +187,7 @@ def adapt_cfg(cls, cfg: Any) -> None: cls._require_newton_physics(cfg, label) cls._promote_scene_entity_cfgs(cfg) cls._swap_mdp(cfg, label) + cls._swap_noise_cfgs(cfg, label) @staticmethod def _require_newton_physics(cfg: Any, label: str) -> None: @@ -218,7 +225,7 @@ def _promote_scene_entity_cfgs(cls, cfg: Any) -> None: promoted: list[str] = [] for path, term in cls._walk_terms(cfg): - if not path or path[0] not in cls.WARP_MANAGED_GROUPS: + if not path or path[0] not in (cls.WARP_MANAGED_GROUPS | cls.OPTIONAL_WARP_GROUPS): continue # term runs on a stable manager; keep the stable entity params = getattr(term, "params", None) if not isinstance(params, dict): @@ -267,8 +274,9 @@ def _swap_mdp(cls, cfg: Any, label: str) -> None: swapped = 0 missing: list[tuple[str, str, str]] = [] # (location, attr, symbol) for path, term in cls._walk_terms(cfg): - if not path or path[0] not in cls.WARP_MANAGED_GROUPS: + if not path or path[0] not in (cls.WARP_MANAGED_GROUPS | cls.OPTIONAL_WARP_GROUPS): continue # term runs on a stable manager; leave it stable + optional = path[0] in cls.OPTIONAL_WARP_GROUPS location = ".".join(path) for attr in ("func", "class_type"): # a term implements via either attr stable = getattr(term, attr, None) @@ -280,8 +288,9 @@ def _swap_mdp(cls, cfg: Any, label: str) -> None: searched.update(m.__name__ for m in module_cache[origin]) twin = cls._resolve_warp_twin(stable.__name__, module_cache[origin]) if twin is None: - missing.append((location, attr, stable.__name__)) # collect every miss; report once below - continue + if not optional: + missing.append((location, attr, stable.__name__)) # collect every miss; report once below + continue # optional group: the stable term stays on the legacy fallback setattr(term, attr, twin) swapped += 1 @@ -403,6 +412,61 @@ def _nearest_mdp_module(mirrored: str) -> str | None: continue return None + @classmethod + def _swap_noise_cfgs(cls, cfg: Any, label: str) -> None: + """Replace stable observation-noise cfgs with their warp-native twins. + + The warp observation manager applies noise through the experimental noise + cfg family (in-place Warp kernels); a stable noise cfg would otherwise be + silently ignored, training against a different MDP. Twins resolve by class + name; data fields are copied while ``func`` keeps the twin's warp default. + Noise cfgs without a warp twin (e.g. class-based ``NoiseModelCfg``) are a + hard error — silently changing the MDP is never acceptable. + """ + import dataclasses + + from isaaclab.utils import noise as stable_noise + + from isaaclab_experimental.utils import noise as warp_noise + + unsupported: list[str] = [] + swapped = 0 + for path, term in cls._walk_terms(cfg): + if not path or path[0] != "observations": + continue + noise_cfg = getattr(term, "noise", None) + if noise_cfg is None: + continue + if type(noise_cfg).__module__.startswith(cls.WARP_ROOT_PREFIXES): + continue # already warp-native + twin_cls = getattr(warp_noise, type(noise_cfg).__name__, None) + if ( + twin_cls is None + or not twin_cls.__module__.startswith(cls.WARP_ROOT_PREFIXES) + or not isinstance(noise_cfg, stable_noise.NoiseCfg) + ): + unsupported.append(f"{'.'.join(path)}: {type(noise_cfg).__name__}") + continue + if noise_cfg.func != type(noise_cfg)().func: + # A user-customized noise func has no warp twin; replacing it with the + # twin's default would silently change the MDP. + unsupported.append(f"{'.'.join(path)}: {type(noise_cfg).__name__} with custom func {noise_cfg.func!r}") + continue + fields = {f.name: getattr(noise_cfg, f.name) for f in dataclasses.fields(noise_cfg) if f.name != "func"} + if any(isinstance(value, torch.Tensor) for value in fields.values()): + # The warp noise kernels take scalar parameters; tensor-valued ranges + # (per-element noise) have no warp twin yet. + unsupported.append(f"{'.'.join(path)}: {type(noise_cfg).__name__} with tensor-valued parameters") + continue + term.noise = twin_cls(**fields) + swapped += 1 + if unsupported: + raise FrontendIncompatibleError( + f"warp env {label!r}: observation noise cfgs without warp twins:\n " + "\n ".join(unsupported) + ) + if swapped: + logger.info("frontend.warp: swapped %d observation noise cfg(s) to warp twins", swapped) + @staticmethod def _import_twin_module(target: str) -> ModuleType: """Import a resolved twin module; a module that exists but breaks is a hard error.""" diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py index 3e5aea6cc4b0..6c358e1c061e 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py @@ -105,26 +105,21 @@ def reset( """ env_mask = self._resolve_reset_mask(None, env_mask) # Dispatch by term kind: Warp class terms consume the mask directly; stable - # class terms receive compact IDs (materialized at most once) or a global - # selection, depending on their declared mode. + # class terms receive compact IDs, materialized at most once per call. compact_env_ids = env_ids for term_cfg, mode in zip(self._term_cfgs, self._term_modes): if isinstance(term_cfg.func, ManagerTermBase): term_cfg.func.reset(env_mask=env_mask) elif isinstance(term_cfg.func, StableManagerTermBase): - if mode == "legacy_ids": - if compact_env_ids is None: - compact_env_ids = self._compact_legacy_env_ids(env_mask) - term_env_ids = compact_env_ids - else: - term_env_ids = slice(None) - term_cfg.func.reset(env_ids=term_env_ids) + if compact_env_ids is None: + compact_env_ids = self._compact_legacy_env_ids(env_mask) + term_cfg.func.reset(env_ids=compact_env_ids) return self._reset_extras @property def requires_host_ids(self) -> bool: """Whether any active legacy term genuinely consumes compact environment IDs.""" - return "legacy_ids" in self._term_modes + return "legacy" in self._term_modes @property def requires_host_boundary(self) -> bool: @@ -149,19 +144,14 @@ def compute( self._term_states_wp.zero_() compact_env_ids = env_ids # Term modes: "mask" = Warp-native ``(env, env_mask, out)``; - # "legacy_ids" = stable ``(env, env_ids)``; "legacy_global" = stable term - # that ignores its ID argument (``requires_host_ids=False``). + # "legacy" = stable ``(env, env_ids)`` with compact IDs materialized at most once. for term_idx, (term_cfg, mode) in enumerate(zip(self._term_cfgs, self._term_modes)): if mode == "mask": term_cfg.func(self._env, env_mask, term_cfg.out, **term_cfg.params) continue - if mode == "legacy_ids": - if compact_env_ids is None: - compact_env_ids = self._compact_legacy_env_ids(env_mask) - term_env_ids = compact_env_ids - else: - term_env_ids = slice(None) - state = term_cfg.func(self._env, term_env_ids, **term_cfg.params) + if compact_env_ids is None: + compact_env_ids = self._compact_legacy_env_ids(env_mask) + state = term_cfg.func(self._env, compact_env_ids, **term_cfg.params) if state is not None: if isinstance(state, torch.Tensor) and state.numel() != 1: raise TypeError( @@ -204,7 +194,7 @@ def _prepare_terms(self): mode = "mask" self._resolve_common_term_cfg(term_name, term_cfg, min_argc=3) else: - mode = "legacy_global" if getattr(term_cfg, "requires_host_ids", None) is False else "legacy_ids" + mode = "legacy" self._resolve_legacy_term_cfg(term_name, term_cfg) self._term_names.append(term_name) self._term_cfgs.append(term_cfg) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py index bcdfa816ce8a..9dad78dc25c9 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py @@ -13,7 +13,7 @@ ``func(env, out, **params) -> None`` signature. :class:`CurriculumTermCfg` is the one override: it extends the stable class with -the ``requires_host_ids`` selector consumed by the Warp-first curriculum manager. +the Warp-first curriculum term configuration. """ from __future__ import annotations @@ -37,13 +37,6 @@ class CurriculumTermCfg(_CurriculumTermCfg): """Configuration for a Warp-mask or legacy curriculum term.""" - requires_host_ids: bool | None = None - """Whether a legacy term consumes compact environment IDs. - - ``None`` selects the manager default: mask-native terms do not require IDs, while legacy terms do. - Set this to ``False`` for global legacy terms that ignore their ``env_ids`` argument. - """ - __all__ = [ "ActionTermCfg", diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index 12fb601acbae..0c5e9a2ea1de 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -673,6 +673,16 @@ def _prepare_terms(self): # noqa: C901 # check noise settings if not group_cfg.enable_corruption: term_cfg.noise = None + elif term_cfg.noise is not None and not type(term_cfg.noise).__module__.startswith( + "isaaclab_experimental." + ): + raise TypeError( + f"Observation term '{term_name}' uses noise cfg" + f" '{type(term_cfg.noise).__module__}.{type(term_cfg.noise).__name__}', which the Warp" + " observation manager would silently ignore. Use the warp-native noise cfgs from" + " isaaclab_experimental.utils.noise (the warp frontend converts stable cfgs" + " automatically)." + ) # check group history params and override terms if group_cfg.history_length is not None: term_cfg.history_length = group_cfg.history_length @@ -754,6 +764,10 @@ def _prepare_terms(self): # noqa: C901 f" and optional parameters: {args_with_defaults}, but received: {term_params}." ) + # plumb the shared per-env RNG state so Warp noise kernels can consume it + # (function-style NoiseCfg kernels read it off the cfg at call time) + if term_cfg.noise is not None and isinstance(term_cfg.noise, noise.NoiseCfg): + term_cfg.noise.rng_state_wp = self._env.rng_state_wp # prepare noise model classes if term_cfg.noise is not None and isinstance(term_cfg.noise, noise.NoiseModelCfg): # plumb the shared per-env RNG state so Warp noise kernels can consume it diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py index a5bbedd02214..20c043180434 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend.py +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -473,10 +473,10 @@ def test_promote_scene_entity_cfgs_walks_all_term_groups(): assert isinstance(cfg.observations.perception.o3.params["asset_cfg"], WarpSceneEntityCfg) # The event manager is warp-first, so its terms are promoted too. assert isinstance(cfg.events.e1.params["asset_cfg"], WarpSceneEntityCfg) - # Curriculum runs on the stable (torch) manager, so its SceneEntityCfg is - # left untouched — promoting it would hand a warp variant to a torch manager. - assert isinstance(cfg.curriculum.c1.params["asset_cfg"], StableSceneEntityCfg) - assert not isinstance(cfg.curriculum.c1.params["asset_cfg"], WarpSceneEntityCfg) + # Curriculum is adapted opportunistically (terms swap to warp twins when one + # exists), so its entities are promoted as well; the warp SceneEntityCfg + # subclasses the stable one, keeping legacy fallback terms valid. + assert isinstance(cfg.curriculum.c1.params["asset_cfg"], WarpSceneEntityCfg) # ====================================================================== diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index ea534f3a4b33..6e36713afae5 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -169,3 +169,93 @@ def test_stable_cartpole_cfg_adapts_to_current_warp_module_layout(): assert cfg.rewards.pole_pos.func.__module__.startswith("isaaclab_tasks_experimental.core.cartpole.mdp") assert cfg.rewards.success_rate.func.__module__.startswith("isaaclab_tasks_experimental.core.cartpole.mdp") assert issubclass(cfg.actions.joint_effort.class_type, ActionTerm) + + +def _iter_obs_terms(cfg): + """Yield (name, term cfg) for every observation term that carries a noise slot.""" + from isaaclab.managers.manager_term_cfg import ObservationGroupCfg + + for group in vars(cfg.observations).values(): + if not isinstance(group, ObservationGroupCfg): + continue + for name, term in vars(group).items(): + if hasattr(term, "noise"): + yield name, term + + +def test_stable_observation_noise_converts_to_warp_twins(): + """Stable noise cfgs (e.g. Unoise) become warp-native twins with the same parameters.""" + from isaaclab.utils import noise as stable_noise + + entry = _cfg_entry_point("Isaac-Velocity-Flat-UnitreeGo2") + module_path, class_name = entry.split(":") + raw = getattr(importlib.import_module(module_path), class_name)() + stable_params = { + name: (term.noise.n_min, term.noise.n_max) + for name, term in _iter_obs_terms(raw) + if isinstance(term.noise, stable_noise.UniformNoiseCfg) + } + assert stable_params, "expected uniform noise on the stable velocity observations" + + cfg = _load_adapted_cfg(entry) + converted = dict(_iter_obs_terms(cfg)) + for name, (n_min, n_max) in stable_params.items(): + twin = converted[name].noise + assert type(twin).__module__.startswith("isaaclab_experimental."), name + assert (twin.n_min, twin.n_max) == (n_min, n_max), name + + +def test_noise_cfg_without_warp_twin_is_a_hard_error(): + """Class-based noise models have no warp twin yet: adapting must fail loudly.""" + from isaaclab_experimental.envs.frontend import FrontendIncompatibleError + + from isaaclab.utils.noise import NoiseModelCfg, UniformNoiseCfg + + entry = _cfg_entry_point("Isaac-Velocity-Flat-UnitreeGo2") + module_path, class_name = entry.split(":") + cfg = getattr(importlib.import_module(module_path), class_name)() + cfg = resolve_presets(cfg, selected=("newton_mjwarp",)) + name, term = next(iter(_iter_obs_terms(cfg))) + term.noise = NoiseModelCfg(noise_cfg=UniformNoiseCfg()) + with pytest.raises(FrontendIncompatibleError, match="noise"): + WarpFrontend.adapt_cfg(cfg) + + +def test_curriculum_terms_swap_to_warp_twins_when_available(): + """Rough-terrain curriculum swaps to its warp twin instead of the legacy host fallback.""" + cfg = _load_adapted_cfg(_cfg_entry_point("Isaac-Velocity-Rough-AnymalD")) + terms = {n: t for n, t in vars(cfg.curriculum).items() if t is not None and hasattr(t, "func")} + assert terms, "expected curriculum terms on the rough task" + for name, term in terms.items(): + assert not isinstance(term.func, str) and term.func.__module__.startswith(_WARP_ROOTS), ( + f"curriculum term '{name}' was not swapped to a warp twin (got {term.func!r})" + ) + assert getattr(term.func, "_curriculum_mask_native", False), ( + f"curriculum term '{name}' resolved to a legacy ID-based twin instead of the" + f" mask-native class (got {term.func!r})" + ) + + +def test_reach_curriculum_terms_swap_to_mask_native_class(): + """Reach's weight-modification curricula resolve to the mask-native class twin.""" + from isaaclab_tasks_experimental.core.reach.mdp import modify_reward_weight + + cfg = _load_adapted_cfg(_cfg_entry_point("Isaac-Reach-Franka")) + terms = {n: t for n, t in vars(cfg.curriculum).items() if t is not None and hasattr(t, "func")} + assert terms, "expected curriculum terms on the reach task" + for name, term in terms.items(): + assert term.func is modify_reward_weight, f"curriculum term '{name}' got {term.func!r}" + + +def test_curriculum_term_without_twin_stays_stable(): + """Optional group: a curriculum term lacking a twin stays on the legacy fallback, no error.""" + import isaaclab.utils.math as math_utils + + entry = _cfg_entry_point("Isaac-Velocity-Rough-AnymalD") + module_path, class_name = entry.split(":") + cfg = getattr(importlib.import_module(module_path), class_name)() + cfg = resolve_presets(cfg, selected=("newton_mjwarp",)) + term = next(t for t in vars(cfg.curriculum).values() if t is not None and hasattr(t, "func")) + term.func = math_utils.quat_mul # stable-rooted symbol with no warp mirror + WarpFrontend.adapt_cfg(cfg) + assert term.func is math_utils.quat_mul diff --git a/source/isaaclab_experimental/test/managers/test_curriculum_manager.py b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py index 11910d139907..bf9a30f299ac 100644 --- a/source/isaaclab_experimental/test/managers/test_curriculum_manager.py +++ b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py @@ -14,8 +14,8 @@ import torch import warp as wp from isaaclab_experimental.managers import CurriculumManager, CurriculumTermCfg -from isaaclab_tasks_experimental.core.reach.mdp import ModifyRewardWeight -from isaaclab_tasks_experimental.core.velocity.mdp import TerrainLevelsVel, terrain_levels_vel +from isaaclab_tasks_experimental.core.reach.mdp import modify_reward_weight +from isaaclab_tasks_experimental.core.velocity.mdp import terrain_levels_vel from isaaclab.managers import ManagerTermBase as StableManagerTermBase from isaaclab.terrains import TerrainImporter @@ -25,7 +25,7 @@ class _GlobalCurriculumTerm(StableManagerTermBase): """Legacy global term that intentionally ignores compact environment IDs.""" def __call__(self, env, env_ids, value: float) -> float: - assert env_ids == slice(None) + del env_ids # a stable term is free to ignore its selection return value @@ -299,7 +299,6 @@ def test_terrain_importer_mask_update_preserves_sparse_and_stable_semantics(): def test_curriculum_manager_updates_masked_levels_and_logs_without_host_compaction(monkeypatch: pytest.MonkeyPatch): """Manager compute/reset should remain mask-native and expose persistent scalar logging.""" - assert terrain_levels_vel is not TerrainLevelsVel terrain = _make_terrain([0, 1, 2, 2, 1, 0]) root_positions = terrain.env_origins.numpy().copy() root_positions[:, 0] += np.array([5.0, 1.0, 3.0, 5.0, 9.0, 0.0], dtype=np.float32) @@ -310,12 +309,7 @@ def test_curriculum_manager_updates_masked_levels_and_logs_without_host_compacti env = _Env(terrain, root_positions, commands) manager = CurriculumManager( { - "terrain_levels": CurriculumTermCfg(func=TerrainLevelsVel), - "global_state": CurriculumTermCfg( - func=_GlobalCurriculumTerm, - params={"value": 7.0}, - requires_host_ids=False, - ), + "terrain_levels": CurriculumTermCfg(func=terrain_levels_vel), }, env, ) @@ -326,7 +320,8 @@ def test_curriculum_manager_updates_masked_levels_and_logs_without_host_compacti state_ptr = manager._term_states_wp.ptr extras_ref = manager.reset_extras assert not manager.requires_host_ids - assert manager.requires_host_boundary + # a purely mask-native manager needs no host boundary at all + assert not manager.requires_host_boundary def _fail_host_compaction(*args, **kwargs): raise AssertionError("Torch host compaction/scalar extraction reached the Warp curriculum path") @@ -349,7 +344,6 @@ def _fail_host_compaction(*args, **kwargs): assert extras is extras_ref assert extras["Curriculum/terrain_levels"] is manager.reset_extras["Curriculum/terrain_levels"] torch.testing.assert_close(extras["Curriculum/terrain_levels"], terrain.terrain_levels.float().mean()) - torch.testing.assert_close(extras["Curriculum/global_state"], torch.tensor(7.0)) assert term._move_up_wp.ptr == move_up_ptr assert term._move_down_wp.ptr == move_down_ptr assert manager._term_states_wp.ptr == state_ptr @@ -363,7 +357,10 @@ def test_curriculum_manager_compacts_ids_inside_legacy_term_boundary(): sim=SimpleNamespace(is_playing=lambda: True), ) manager = CurriculumManager( - {"id_count": CurriculumTermCfg(func=_LegacyIdCurriculumTerm, params={"scale": 2.0})}, + { + "id_count": CurriculumTermCfg(func=_LegacyIdCurriculumTerm, params={"scale": 2.0}), + "global_state": CurriculumTermCfg(func=_GlobalCurriculumTerm, params={"value": 7.0}), + }, env, ) env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") @@ -377,6 +374,7 @@ def test_curriculum_manager_compacts_ids_inside_legacy_term_boundary(): torch.testing.assert_close(term.compute_env_ids, torch.tensor([0, 2])) torch.testing.assert_close(term.reset_env_ids, torch.tensor([0, 2])) torch.testing.assert_close(extras["Curriculum/id_count"], torch.tensor(4.0)) + torch.testing.assert_close(extras["Curriculum/global_state"], torch.tensor(7.0)) def test_curriculum_manager_consumes_threaded_ids_without_compaction(monkeypatch: pytest.MonkeyPatch): @@ -420,11 +418,11 @@ def test_registered_reach_curricula_update_weights_without_a_host_boundary(): manager = CurriculumManager( { "action_rate": CurriculumTermCfg( - func=ModifyRewardWeight, + func=modify_reward_weight, params={"term_name": "action_rate", "weight": -0.005, "num_steps": 4500}, ), "joint_vel": CurriculumTermCfg( - func=ModifyRewardWeight, + func=modify_reward_weight, params={"term_name": "joint_vel", "weight": -0.001, "num_steps": 4500}, ), }, @@ -456,7 +454,7 @@ def test_curriculum_manager_rejects_structured_legacy_state() -> None: terrain = _make_terrain([0, 1]) env = _Env(terrain, terrain.env_origins.numpy().copy(), np.zeros((2, 3), dtype=np.float32)) manager = CurriculumManager( - {"structured": CurriculumTermCfg(func=_structured_state, requires_host_ids=False)}, + {"structured": CurriculumTermCfg(func=_structured_state)}, env, ) env_mask = wp.array([True, False], dtype=wp.bool, device="cpu") diff --git a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py index caea47f38cb2..2bdc680d23a8 100644 --- a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py +++ b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py @@ -155,15 +155,24 @@ def test_mask_native_code_has_no_unsanctioned_host_syncs(): assert not stale, f"Stale sanctioned-boundary entries (function gone or no longer syncs): {sorted(stale)}" -# Expected ``@WarpCapturable(False)`` opt-outs: (path suffix, decorated name). +# Expected capturability opt-outs — ``@WarpCapturable(False)`` decorations or +# ``_warp_capturable = False`` class attributes: (path suffix, name). EXPECTED_NON_CAPTURABLE = { - ("envs/mdp/events.py", "randomize_rigid_body_com"), - ("envs/mdp/observations.py", "height_scan"), + ("isaaclab_experimental/envs/mdp/events.py", "randomize_rigid_body_com"), + ("isaaclab_experimental/envs/mdp/observations.py", "height_scan"), + # Reads the Python-owned common step counter; captured replay would bake it. + ("isaaclab_tasks_experimental/core/reach/mdp/curriculums.py", "modify_reward_weight"), } +NON_CAPTURABLE_SCAN_ROOTS = [ + "source/isaaclab_experimental/isaaclab_experimental", + "source/isaaclab_tasks_experimental/isaaclab_tasks_experimental", +] + def _warp_capturable_false_targets(tree: ast.Module) -> list[str]: - """Names decorated with ``@WarpCapturable(False, ...)`` in a module.""" + """Names annotated non-capturable: ``@WarpCapturable(False, ...)`` decorations + or ``_warp_capturable = False`` class attributes.""" out = [] for node in ast.walk(tree): if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): @@ -178,21 +187,29 @@ def _warp_capturable_false_targets(tree: ast.Module) -> list[str]: and deco.args[0].value is False ): out.append(node.name) + if isinstance(node, ast.ClassDef): + for stmt in node.body: + if ( + isinstance(stmt, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "_warp_capturable" for t in stmt.targets) + and isinstance(stmt.value, ast.Constant) + and stmt.value.value is False + ): + out.append(node.name) return out def test_non_capturable_terms_are_inventoried(): """Every ``@WarpCapturable(False)`` opt-out is documented here, and only these.""" found = set() - root = _REPO_ROOT / "source/isaaclab_experimental/isaaclab_experimental" - for path in sorted(root.rglob("*.py")): - rel = path.relative_to(_REPO_ROOT).as_posix() - for name in _warp_capturable_false_targets(ast.parse(path.read_text(), filename=rel)): - found.add((rel, name)) - expected = { - (f"source/isaaclab_experimental/isaaclab_experimental/{suffix}", name) - for suffix, name in EXPECTED_NON_CAPTURABLE - } + for scan_root in NON_CAPTURABLE_SCAN_ROOTS: + for path in sorted((_REPO_ROOT / scan_root).rglob("*.py")): + rel = path.relative_to(_REPO_ROOT).as_posix() + if "/test/" in rel: + continue + for name in _warp_capturable_false_targets(ast.parse(path.read_text(), filename=rel)): + found.add((rel, name)) + expected = {(f"source/{suffix.split('/')[0]}/{suffix}", name) for suffix, name in EXPECTED_NON_CAPTURABLE} assert found == expected, ( "Non-capturable inventory drift.\n" f" unexpected: {sorted(found - expected)}\n" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.pyi index 1384097609f8..0fe12a032cc9 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.pyi +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.pyi @@ -4,7 +4,7 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ - "ModifyRewardWeight", + "modify_reward_weight", "position_command_error", "position_command_error_tanh", "orientation_command_error", @@ -14,5 +14,5 @@ __all__ = [ # override with reach-specific terms below. from isaaclab_experimental.envs.mdp import * # noqa: F401, F403 -from .curriculums import ModifyRewardWeight +from .curriculums import modify_reward_weight from .rewards import orientation_command_error, position_command_error, position_command_error_tanh diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/curriculums.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/curriculums.py index 66bb0b0eabfb..a25adf0176bd 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/curriculums.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/curriculums.py @@ -39,7 +39,7 @@ def _update_reward_weight( reset_requested[0] = 0 -class ModifyRewardWeight(ManagerTermBase): +class modify_reward_weight(ManagerTermBase): """Update one reward weight after a reset reaches the configured step threshold.""" # The Python-owned common step counter is intentionally read in eager mode. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi index 463c665cd67a..e76bfc598bff 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi @@ -5,8 +5,7 @@ __all__ = [ "terrain_levels_vel", - "TerrainLevelsVel", - "feet_air_time", + "feet_air_time", "feet_air_time_positive_biped", "feet_slide", "stand_still_joint_deviation_l1", @@ -19,7 +18,7 @@ __all__ = [ # observations, ...) lazily, then override with locomotion-specific terms below. from isaaclab_experimental.envs.mdp import * # noqa: F401, F403 -from .curriculums import TerrainLevelsVel, terrain_levels_vel +from .curriculums import terrain_levels_vel from .rewards import ( feet_air_time, feet_air_time_positive_biped, diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py index f55af9254efa..1b96ad087c1c 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py @@ -7,10 +7,8 @@ from __future__ import annotations -from collections.abc import Sequence from typing import TYPE_CHECKING -import torch import warp as wp from isaaclab_experimental.managers import CurriculumTermCfg, ManagerTermBase, SceneEntityCfg @@ -58,7 +56,7 @@ def _compute_terrain_level_mean( wp.atomic_add(out, 0, wp.float32(terrain_levels[env_id]) * scale) -class TerrainLevelsVel(ManagerTermBase): +class terrain_levels_vel(ManagerTermBase): """Update terrain levels from commanded locomotion progress using a boolean environment mask.""" _curriculum_mask_native = True @@ -127,18 +125,3 @@ def __call__( inputs=[self._terrain_levels_wp, self._mean_scale, out], device=self.device, ) - - -def terrain_levels_vel( - env: ManagerBasedRLEnv, env_ids: Sequence[int], asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") -) -> torch.Tensor: - """Update terrain levels through the legacy ID-based public API.""" - asset: Articulation = env.scene[asset_cfg.name] - terrain: TerrainImporter = env.scene.terrain - command = env.command_manager.get_command("base_velocity") - distance = torch.norm(asset.data.root_pos_w.torch[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1) - move_up = distance > terrain.cfg.terrain_generator.size[0] / 2 - move_down = distance < torch.norm(command[env_ids, :2], dim=1) * env.max_episode_length_s * 0.5 - move_down *= ~move_up - terrain.update_env_origins(env_ids, move_up, move_down) - return torch.mean(terrain.terrain_levels.float()) From 35c652bff9288bdfa661112a8a396020ce92195d Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 22 Jul 2026 02:14:16 -0700 Subject: [PATCH 54/58] Merge the curriculum conversion pins into one parametrized test Both tasks exercise the same converter mechanism (name resolution to a mask-native class twin); one parametrized test covers both data points. --- .../test/envs/test_frontend_cfg_conversion.py | 29 +++++++------------ 1 file changed, 11 insertions(+), 18 deletions(-) diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index 6e36713afae5..82ec8489de9c 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -221,32 +221,25 @@ def test_noise_cfg_without_warp_twin_is_a_hard_error(): WarpFrontend.adapt_cfg(cfg) -def test_curriculum_terms_swap_to_warp_twins_when_available(): - """Rough-terrain curriculum swaps to its warp twin instead of the legacy host fallback.""" - cfg = _load_adapted_cfg(_cfg_entry_point("Isaac-Velocity-Rough-AnymalD")) +@pytest.mark.parametrize( + "task_id", + ["Isaac-Velocity-Rough-AnymalD", "Isaac-Reach-Franka"], + ids=["rough-terrain-levels", "reach-weight-modifiers"], +) +def test_curriculum_terms_swap_to_mask_native_twins(task_id: str): + """Curriculum terms resolve to the mask-native warp classes, not legacy ID paths.""" + cfg = _load_adapted_cfg(_cfg_entry_point(task_id)) terms = {n: t for n, t in vars(cfg.curriculum).items() if t is not None and hasattr(t, "func")} - assert terms, "expected curriculum terms on the rough task" + assert terms, f"expected curriculum terms on {task_id}" for name, term in terms.items(): assert not isinstance(term.func, str) and term.func.__module__.startswith(_WARP_ROOTS), ( f"curriculum term '{name}' was not swapped to a warp twin (got {term.func!r})" ) - assert getattr(term.func, "_curriculum_mask_native", False), ( - f"curriculum term '{name}' resolved to a legacy ID-based twin instead of the" - f" mask-native class (got {term.func!r})" + assert isinstance(term.func, type), ( + f"curriculum term '{name}' resolved to a legacy function instead of a mask-native class" ) -def test_reach_curriculum_terms_swap_to_mask_native_class(): - """Reach's weight-modification curricula resolve to the mask-native class twin.""" - from isaaclab_tasks_experimental.core.reach.mdp import modify_reward_weight - - cfg = _load_adapted_cfg(_cfg_entry_point("Isaac-Reach-Franka")) - terms = {n: t for n, t in vars(cfg.curriculum).items() if t is not None and hasattr(t, "func")} - assert terms, "expected curriculum terms on the reach task" - for name, term in terms.items(): - assert term.func is modify_reward_weight, f"curriculum term '{name}' got {term.func!r}" - - def test_curriculum_term_without_twin_stays_stable(): """Optional group: a curriculum term lacking a twin stays on the legacy fallback, no error.""" import isaaclab.utils.math as math_utils From 2c9b2af8c0b8fe4c78f1657f9a60346f40d53e0f Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 22 Jul 2026 10:33:38 -0700 Subject: [PATCH 55/58] Align warp environments doc with the merged catalog Use the presets= Hydra group name in the PhysX limitation note and drop performance rows for flat-velocity tasks removed by the catalog update (AnymalB, AnymalC, A1, Go1); rename the Go2 row to the UnitreeGo2 task name. --- .../newton/warp-environments.rst | 26 +++---------------- 1 file changed, 3 insertions(+), 23 deletions(-) diff --git a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst index 988fe7ac561a..610d2af98fa1 100644 --- a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst +++ b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst @@ -146,16 +146,6 @@ both running on the Newton physics backend. Measured over 300 iterations with 40 - 11,458 - 7,813 - -31.83% - * - AnymalB - - Manager - - 29,188 - - 21,781 - - -25.38% - * - AnymalC - - Manager - - 30,938 - - 22,228 - - -28.15% * - AnymalD - Manager - 32,294 @@ -176,17 +166,7 @@ both running on the Newton physics backend. Measured over 300 iterations with 40 - 22,202 - 15,864 - -28.55% - * - A1 - - Manager - - 15,257 - - 9,907 - - -35.07% - * - Go1 - - Manager - - 16,515 - - 11,869 - - -28.13% - * - Go2 + * - UnitreeGo2 - Manager - 15,221 - 9,966 @@ -203,7 +183,7 @@ overhead. Reading the table above: - **Manager-based classic RL** (Cartpole, Ant) — biggest gains (-52% to -54%). Many small reward / observation terms with low compute per term, so per-launch CPU overhead dominated the stable baseline. -- **Manager-based locomotion** (Anymal, G1, H1, Cassie, Unitree) — consistent -25% to -38% +- **Manager-based locomotion** (AnymalD, G1, H1, Cassie, UnitreeGo2) — consistent -21% to -38% range. The MDP has more terms but the underlying physics step is heavier, so the relative Python savings shrink. - **Direct workflow** — gains scale with how much the env's step body was Python (Ant -51%, @@ -231,7 +211,7 @@ specific to warp envs; for Newton physics limitations see :doc:`supported-featur ``class_type`` fields resolve to ``isaaclab_physx.*`` classes that depend on ``omni.physics.tensors`` (a Kit module the warp runtime does not initialise), and several warp APIs (env-mask reset, CUDA graph capture) require the Newton articulation. Configure - the cfg with a Newton physics block (or ``physics=newton_mjwarp``). + the cfg with a Newton physics block (or ``presets=newton_mjwarp``). **MDP coverage** From 74734302e062299e2eb1c0b25766b8750975b44e Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 22 Jul 2026 18:02:30 -0700 Subject: [PATCH 56/58] Back terrain and scene origin buffers with ProxyArray Store the terrain importer origin, level, and type buffers (and the scene clone-plan origin fallback) as ProxyArray so the Torch accessors and Warp views alias one owner. This drops the per-buffer view caches and pointer-drift guards: pointer stability holds by construction, and the warp path reaches the views through a single boundary accessor. --- .../jichuanh-warp-frontend-cleanup.minor.rst | 10 +- .../isaaclab/scene/interactive_scene.py | 23 +-- .../isaaclab/terrains/terrain_importer.py | 139 +++++++++--------- .../test/scene/test_interactive_scene.py | 10 +- .../test/terrains/test_terrain_importer.py | 8 +- .../envs/manager_based_env_warp.py | 8 +- .../test/envs/mdp/parity_helpers.py | 6 +- .../test/envs/mdp/test_events_warp_parity.py | 4 +- .../test/managers/test_curriculum_manager.py | 69 +++++---- .../core/locomotion/locomotion_env_warp.py | 2 +- .../reorient/inhand_manipulation_warp_env.py | 2 +- .../core/velocity/mdp/curriculums.py | 2 +- 12 files changed, 148 insertions(+), 135 deletions(-) diff --git a/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst index 6e203e14eb99..453998f2954b 100644 --- a/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst +++ b/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -1,11 +1,11 @@ Added ^^^^^ -* Added :attr:`~isaaclab.terrains.TerrainImporter.env_origins_wp`, - :attr:`~isaaclab.terrains.TerrainImporter.terrain_levels_wp`, and +* Added :attr:`~isaaclab.terrains.TerrainImporter.env_origins_pa`, + :attr:`~isaaclab.terrains.TerrainImporter.terrain_levels_pa`, and :meth:`~isaaclab.terrains.TerrainImporter.update_env_origins_mask` for zero-copy Warp terrain access and mask-based curriculum updates. -* Added :attr:`~isaaclab.scene.InteractiveScene.env_origins_wp` and a boolean ``env_mask`` +* Added :attr:`~isaaclab.scene.InteractiveScene.env_origins_pa` and a boolean ``env_mask`` option to :meth:`~isaaclab.scene.InteractiveScene.reset` for mask-based scene resets. Changed @@ -13,3 +13,7 @@ Changed * Changed the uniform pose and velocity command debug visualization to shared private mixins reused by the Warp-native command terms. +* Changed the :class:`~isaaclab.terrains.TerrainImporter` origin, level, and type buffers to + :class:`~isaaclab.utils.warp.ProxyArray` storage: the Torch accessors and Warp views alias + one owner, so consumer-held views stay pointer-stable across reconfiguration without + per-buffer view caches. diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index b8ca77aa1b31..000590c4fa0c 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -36,6 +36,7 @@ from isaaclab.sim import SimulationContext from isaaclab.sim.utils.stage import get_current_stage, get_current_stage_id from isaaclab.sim.views import FrameView +from isaaclab.utils.warp import ProxyArray # Note: This is a temporary import for the VisuoTactileSensorCfg class. # It will be removed once the VisuoTactileSensor class is added to the core Isaac Lab framework. @@ -191,10 +192,10 @@ def __init__(self, cfg: InteractiveSceneCfg): self._aggregate_scene_data_requirements(requested_viz_types) # The terrain importer owns its origins buffer together with its Warp view. - # For clone-plan origins, :attr:`env_origins_wp` caches a zero-copy view on + # For clone-plan origins, :attr:`env_origins_pa` caches a zero-copy proxy on # first access (manager/term initialization, outside capture), so scene # construction does not require a published clone plan. - self._clone_origins_wp = None + self._clone_origins_pa = None # Collision filtering is PhysX-only (matches both physx and ovphysx). if self.cfg.filter_collisions and "physx" in self.physics_backend and self._is_scene_setup_from_cfg(): @@ -378,18 +379,18 @@ def env_origins(self) -> torch.Tensor: return self.sim.get_clone_plan().positions @property - def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): - """Zero-copy Warp view of environment origins [m], shape ``(num_envs,)``. + def env_origins_pa(self) -> ProxyArray: + """Environment origins proxy [m], shape ``(num_envs,)`` vec3; use ``.warp`` / ``.torch`` views. - Terrain-backed scenes forward the terrain importer's cached view so both - accessors always share one storage owner. Clone-plan origins are wrapped - lazily on first access and cached. + Terrain-backed scenes forward the terrain importer's proxy so both accessors + always share one storage owner. Clone-plan origins are wrapped lazily on + first access and re-wrapped only if a new plan is published. """ if self._terrain is not None: - return self._terrain.env_origins_wp - if self._clone_origins_wp is None or self._clone_origins_wp.ptr != self.env_origins.data_ptr(): - self._clone_origins_wp = wp.from_torch(self.env_origins, dtype=wp.vec3f) - return self._clone_origins_wp + return self._terrain.env_origins_pa + if self._clone_origins_pa is None or self._clone_origins_pa.warp.ptr != self.env_origins.data_ptr(): + self._clone_origins_pa = ProxyArray(wp.from_torch(self.env_origins, dtype=wp.vec3f)) + return self._clone_origins_pa @property def terrain(self) -> TerrainImporter | None: diff --git a/source/isaaclab/isaaclab/terrains/terrain_importer.py b/source/isaaclab/isaaclab/terrains/terrain_importer.py index 7ebfa58d4e5a..b1e0b8cdb6de 100644 --- a/source/isaaclab/isaaclab/terrains/terrain_importer.py +++ b/source/isaaclab/isaaclab/terrains/terrain_importer.py @@ -16,6 +16,7 @@ import isaaclab.sim as sim_utils from isaaclab.markers import VisualizationMarkers from isaaclab.markers.config import FRAME_MARKER_CFG +from isaaclab.utils.warp import ProxyArray from .utils import create_prim_from_mesh @@ -74,16 +75,6 @@ class TerrainImporter: terrain_prim_paths: list[str] """A list containing the USD prim paths to the imported terrains.""" - terrain_origins: torch.Tensor | None - """Sub-terrain origins [m], shape ``(num_rows, num_cols, 3)``. - - If terrain origins is not None, the environment origins are computed based on the terrain origins. - Otherwise, the environment origins are computed based on the grid spacing. - """ - - env_origins: torch.Tensor - """Environment origins [m], shape ``(num_envs, 3)``.""" - def __init__(self, cfg: TerrainImporterCfg): """Initialize the terrain importer. @@ -102,14 +93,14 @@ def __init__(self, cfg: TerrainImporterCfg): self.cfg = cfg self.device = sim_utils.SimulationContext.instance().device # type: ignore - # create buffers for the terrains + # create buffers for the terrains. Each buffer is one ProxyArray owning the + # storage that both the public Torch accessors and the Warp views alias, so + # consumer-held views stay pointer-stable across reconfiguration. self.terrain_prim_paths = list() - self.terrain_origins = None - self.env_origins = None # assigned later when `configure_env_origins` is called - self._terrain_levels_wp = None - self._terrain_types_wp = None - self._terrain_origins_wp = None - self._env_origins_wp = None + self._terrain_origins_pa: ProxyArray | None = None + self._env_origins_pa: ProxyArray | None = None # assigned when `configure_env_origins` is called + self._terrain_levels_pa: ProxyArray | None = None + self._terrain_types_pa: ProxyArray | None = None # private variables self._terrain_flat_patches = dict() @@ -178,16 +169,36 @@ def terrain_names(self) -> list[str]: return [f"'{path.split('/')[-1]}'" for path in self.terrain_prim_paths] @property - def terrain_levels_wp(self) -> wp.array(dtype=wp.int64): - """Cached zero-copy Warp view of terrain levels, shape ``(num_envs,)``.""" - if self._terrain_levels_wp is None: + def terrain_origins(self) -> torch.Tensor | None: + """Sub-terrain origins [m], shape ``(num_rows, num_cols, 3)``; ``None`` for grid-like origins.""" + return self._terrain_origins_pa.torch if self._terrain_origins_pa is not None else None + + @property + def env_origins(self) -> torch.Tensor | None: + """Per-env origins [m], shape ``(num_envs, 3)``; ``None`` until origins are configured.""" + return self._env_origins_pa.torch if self._env_origins_pa is not None else None + + @property + def terrain_levels(self) -> torch.Tensor | None: + """Per-env terrain levels, shape ``(num_envs,)``, int64; ``None`` for grid-like origins.""" + return self._terrain_levels_pa.torch if self._terrain_levels_pa is not None else None + + @property + def terrain_types(self) -> torch.Tensor | None: + """Per-env terrain types, shape ``(num_envs,)``, int64; ``None`` for grid-like origins.""" + return self._terrain_types_pa.torch if self._terrain_types_pa is not None else None + + @property + def terrain_levels_pa(self) -> ProxyArray: + """Terrain levels proxy, shape ``(num_envs,)``, int64; use ``.warp`` / ``.torch`` views.""" + if self._terrain_levels_pa is None: raise RuntimeError("Terrain levels are unavailable for grid-like environment origins.") - return self._terrain_levels_wp + return self._terrain_levels_pa @property - def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): - """Cached zero-copy Warp view of environment origins [m], shape ``(num_envs,)``.""" - return self._env_origins_wp + def env_origins_pa(self) -> ProxyArray | None: + """Environment origins proxy [m], shape ``(num_envs,)`` vec3; use ``.warp`` / ``.torch`` views.""" + return self._env_origins_pa """ Operations - Visibility. @@ -348,23 +359,26 @@ def configure_env_origins(self, origins: np.ndarray | torch.Tensor | None = None if isinstance(origins, np.ndarray): origins = torch.from_numpy(origins) # store the origins - self.terrain_origins = self._assign_pointer_stable( - self.terrain_origins, origins.to(self.device, dtype=torch.float).contiguous() + self._terrain_origins_pa = self._assign_proxy_stable( + self._terrain_origins_pa, origins.to(self.device, dtype=torch.float).contiguous(), wp.vec3f ) # compute environment origins - self.env_origins = self._assign_pointer_stable( - self.env_origins, self._compute_env_origins_curriculum(self.cfg.num_envs, self.terrain_origins) + self._env_origins_pa = self._assign_proxy_stable( + self._env_origins_pa, + self._compute_env_origins_curriculum(self.cfg.num_envs, self.terrain_origins), + wp.vec3f, ) else: - self.terrain_origins = None + self._terrain_origins_pa = None + self._terrain_levels_pa = None + self._terrain_types_pa = None # check if env spacing is valid if self.cfg.env_spacing is None: raise ValueError("Environment spacing must be specified for configuring grid-like origins.") # compute environment origins - self.env_origins = self._assign_pointer_stable( - self.env_origins, self._compute_env_origins_grid(self.cfg.num_envs, self.cfg.env_spacing) + self._env_origins_pa = self._assign_proxy_stable( + self._env_origins_pa, self._compute_env_origins_grid(self.cfg.num_envs, self.cfg.env_spacing), wp.vec3f ) - self._configure_warp_origin_views() def update_env_origins(self, env_ids: torch.Tensor, move_up: torch.Tensor, move_down: torch.Tensor) -> None: """Update environment origins through the Torch ID-based interface. @@ -436,10 +450,10 @@ def update_env_origins_mask( move_up, move_down, rng_state, - self._terrain_levels_wp, - self._terrain_types_wp, - self._terrain_origins_wp, - self._env_origins_wp, + self._terrain_levels_pa.warp, + self._terrain_types_pa.warp, + self._terrain_origins_pa.warp, + self._env_origins_pa.warp, self.max_terrain_level, ], device=self.device, @@ -450,37 +464,26 @@ def update_env_origins_mask( """ @staticmethod - def _assign_pointer_stable(current: torch.Tensor | None, new: torch.Tensor) -> torch.Tensor: - """Copy into the existing buffer when layouts match so cached views stay valid.""" + def _assign_proxy_stable(current: ProxyArray | None, new: torch.Tensor, dtype) -> ProxyArray: + """Copy into the existing proxy storage when layouts match so consumer-held views stay valid. + + Args: + current: The proxy currently owning the buffer, or ``None``. + new: Freshly computed values to store. + dtype: Warp dtype of the proxy's array (e.g. ``wp.vec3f``). + + Returns: + The proxy holding the new values; :paramref:`current` when its layout matches. + """ if ( current is not None - and current.shape == new.shape - and current.dtype == new.dtype - and current.device == new.device + and current.torch.shape == new.shape + and current.torch.dtype == new.dtype + and current.torch.device == new.device ): - current.copy_(new) + current.torch.copy_(new) return current - return new - - def _configure_warp_origin_views(self) -> None: - """Create persistent zero-copy Warp views after configuring the Torch buffers. - - Existing views are kept when reconfiguration reuses the underlying Torch - storage, so consumers holding these views stay valid. - """ - if self._env_origins_wp is None or self._env_origins_wp.ptr != self.env_origins.data_ptr(): - self._env_origins_wp = wp.from_torch(self.env_origins, dtype=wp.vec3f) - if self.terrain_origins is None: - self._terrain_levels_wp = None - self._terrain_types_wp = None - self._terrain_origins_wp = None - return - if self._terrain_levels_wp is None or self._terrain_levels_wp.ptr != self.terrain_levels.data_ptr(): - self._terrain_levels_wp = wp.from_torch(self.terrain_levels, dtype=wp.int64) - if self._terrain_types_wp is None or self._terrain_types_wp.ptr != self.terrain_types.data_ptr(): - self._terrain_types_wp = wp.from_torch(self.terrain_types, dtype=wp.int64) - if self._terrain_origins_wp is None or self._terrain_origins_wp.ptr != self.terrain_origins.data_ptr(): - self._terrain_origins_wp = wp.from_torch(self.terrain_origins, dtype=wp.vec3f) + return ProxyArray(wp.from_torch(new.contiguous(), dtype=dtype)) def _compute_env_origins_curriculum(self, num_envs: int, origins: torch.Tensor) -> torch.Tensor: """Compute the origins of the environments defined by the sub-terrains origins.""" @@ -494,15 +497,17 @@ def _compute_env_origins_curriculum(self, num_envs: int, origins: torch.Tensor) # store maximum terrain level possible self.max_terrain_level = num_rows # define all terrain levels and types available - self.terrain_levels = self._assign_pointer_stable( - getattr(self, "terrain_levels", None), + self._terrain_levels_pa = self._assign_proxy_stable( + self._terrain_levels_pa, torch.randint(0, max_init_level + 1, (num_envs,), device=self.device), + wp.int64, ) - self.terrain_types = self._assign_pointer_stable( - getattr(self, "terrain_types", None), + self._terrain_types_pa = self._assign_proxy_stable( + self._terrain_types_pa, torch.div(torch.arange(num_envs, device=self.device), (num_envs / num_cols), rounding_mode="floor").to( torch.long ), + wp.int64, ) # create tensor based on number of environments env_origins = torch.zeros(num_envs, 3, device=self.device) diff --git a/source/isaaclab/test/scene/test_interactive_scene.py b/source/isaaclab/test/scene/test_interactive_scene.py index c9b4010cc0e2..749f58591e20 100644 --- a/source/isaaclab/test/scene/test_interactive_scene.py +++ b/source/isaaclab/test/scene/test_interactive_scene.py @@ -135,16 +135,16 @@ def test_reset_to_env_ids_input_types(device, setup_scene): @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) -def test_env_origins_wp_is_cached_zero_copy_view(device, setup_scene): +def test_env_origins_pa_is_cached_zero_copy_view(device, setup_scene): """The scene should expose one persistent Warp view of its Torch origins.""" make_scene, sim = setup_scene scene = InteractiveScene(make_scene(num_envs=4)) sim.reset() - assert scene.env_origins_wp is scene.env_origins_wp - assert scene.env_origins_wp.dtype == wp.vec3f - assert scene.env_origins_wp.shape == (scene.num_envs,) - assert scene.env_origins_wp.ptr == scene.env_origins.data_ptr() + assert scene.env_origins_pa.warp is scene.env_origins_pa.warp + assert scene.env_origins_pa.warp.dtype == wp.vec3f + assert scene.env_origins_pa.warp.shape == (scene.num_envs,) + assert scene.env_origins_pa.warp.ptr == scene.env_origins.data_ptr() def test_reset_dispatches_warp_mask_to_supported_entities() -> None: diff --git a/source/isaaclab/test/terrains/test_terrain_importer.py b/source/isaaclab/test/terrains/test_terrain_importer.py index 4ced50804d2e..9b53f80d9167 100644 --- a/source/isaaclab/test/terrains/test_terrain_importer.py +++ b/source/isaaclab/test/terrains/test_terrain_importer.py @@ -54,11 +54,11 @@ def test_terrain_importer_env_origins(device, env_spacing, num_envs): # check that the cached Warp view shares the canonical Torch storage assert isinstance(terrain_importer.env_origins, torch.Tensor) - assert terrain_importer.env_origins_wp.dtype == wp.vec3f - assert terrain_importer.env_origins_wp.shape == (num_envs,) + assert terrain_importer.env_origins_pa.warp.dtype == wp.vec3f + assert terrain_importer.env_origins_pa.warp.shape == (num_envs,) assert terrain_importer_origins.shape == (num_envs, 3) - assert terrain_importer.env_origins_wp is terrain_importer.env_origins_wp - assert terrain_importer_origins.data_ptr() == terrain_importer.env_origins_wp.ptr + assert terrain_importer.env_origins_pa.warp is terrain_importer.env_origins_pa.warp + assert terrain_importer_origins.data_ptr() == terrain_importer.env_origins_pa.warp.ptr # obtain env origins using Lab's grid_transforms lab_grid_origins, _ = lab_cloner.grid_transforms(num_envs, spacing=env_spacing, device=sim.device) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py index 38c82ddf0fe0..e2eb17835d4a 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py @@ -279,8 +279,12 @@ def device(self): @property def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): - """Warp-owned environment origins [m] provided by the scene.""" - return self.scene.env_origins_wp + """Warp view of the scene-owned environment origins [m], shape ``(num_envs,)``. + + The scene hands out a :class:`~isaaclab.utils.warp.ProxyArray`; this boundary + accessor resolves it once so Warp-native consumers see a plain Warp array. + """ + return self.scene.env_origins_pa.warp def resolve_env_mask( self, diff --git a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py index 86a352958a15..4e452de96a28 100644 --- a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py +++ b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py @@ -386,13 +386,13 @@ def __init__(self, assets: dict, env_origins, sensors=None): if env_origins is None: # Action-term tests don't need origins; keep the attributes present but empty. self.env_origins = None - self.env_origins_wp = None + self.env_origins_pa = None elif isinstance(env_origins, wp.array): - self.env_origins_wp = env_origins + self.env_origins_pa = ProxyArray(env_origins) self.env_origins = wp.to_torch(env_origins) else: self.env_origins = env_origins - self.env_origins_wp = wp.from_torch(env_origins, dtype=wp.vec3f) + self.env_origins_pa = ProxyArray(wp.from_torch(env_origins, dtype=wp.vec3f)) self.sensors = sensors or {} self.articulations = dict(assets) self.rigid_objects = {} diff --git a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py index 1b21c97d06d0..aba271420a69 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py @@ -81,7 +81,7 @@ class _Env: env.action_manager = MockActionManagerWarp(action_wp[0], action_wp[1]) env.num_envs = NUM_ENVS env.device = DEVICE - env.env_origins_wp = scene.env_origins_wp + env.env_origins_wp = scene.env_origins_pa.warp env.episode_length_buf = episode_length_buf env.step_dt = 0.02 env.max_episode_length_s = 10.0 @@ -126,7 +126,7 @@ class _Env: env.scene = MockScene({"robot": asset}, origins) env.num_envs = NUM_ENVS env.device = DEVICE - env.env_origins_wp = env.scene.env_origins_wp + env.env_origins_wp = env.scene.env_origins_pa.warp env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + seed, device=DEVICE) return env, data, asset diff --git a/source/isaaclab_experimental/test/managers/test_curriculum_manager.py b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py index bf9a30f299ac..9f95f32d412d 100644 --- a/source/isaaclab_experimental/test/managers/test_curriculum_manager.py +++ b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py @@ -19,6 +19,7 @@ from isaaclab.managers import ManagerTermBase as StableManagerTermBase from isaaclab.terrains import TerrainImporter +from isaaclab.utils.warp import ProxyArray class _GlobalCurriculumTerm(StableManagerTermBase): @@ -124,8 +125,8 @@ def env_origins(self) -> torch.Tensor: return self.terrain.env_origins @property - def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): - return self.terrain.env_origins_wp + def env_origins_pa(self) -> ProxyArray: + return self.terrain.env_origins_pa def __getitem__(self, name: str): return self._entities[name] @@ -148,18 +149,16 @@ def __init__(self, terrain: TerrainImporter, root_positions: np.ndarray, command @property def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): - """Return the scene's persistent Warp origin view.""" - return self.scene.env_origins_wp + """Boundary accessor mirroring :attr:`ManagerBasedEnvWarp.env_origins_wp`.""" + return self.scene.env_origins_pa.warp def _init_terrain_buffers(terrain: TerrainImporter) -> None: """Mirror the buffer initializers of :meth:`TerrainImporter.__init__`.""" - terrain.terrain_origins = None - terrain.env_origins = None - terrain._terrain_levels_wp = None - terrain._terrain_types_wp = None - terrain._terrain_origins_wp = None - terrain._env_origins_wp = None + terrain._terrain_origins_pa = None + terrain._env_origins_pa = None + terrain._terrain_levels_pa = None + terrain._terrain_types_pa = None def _make_terrain(levels: list[int]) -> TerrainImporter: @@ -197,15 +196,15 @@ def test_terrain_importer_exposes_persistent_warp_views(): assert isinstance(terrain.env_origins, torch.Tensor) assert isinstance(terrain.terrain_levels, torch.Tensor) assert isinstance(terrain.terrain_types, torch.Tensor) - assert terrain.terrain_origins.data_ptr() == terrain._terrain_origins_wp.ptr - assert terrain.env_origins.data_ptr() == terrain.env_origins_wp.ptr - assert terrain.terrain_levels.data_ptr() == terrain.terrain_levels_wp.ptr - assert terrain.terrain_types.data_ptr() == terrain._terrain_types_wp.ptr + assert terrain.terrain_origins.data_ptr() == terrain._terrain_origins_pa.warp.ptr + assert terrain.env_origins.data_ptr() == terrain.env_origins_pa.warp.ptr + assert terrain.terrain_levels.data_ptr() == terrain.terrain_levels_pa.warp.ptr + assert terrain.terrain_types.data_ptr() == terrain._terrain_types_pa.warp.ptr terrain.terrain_levels[0] = 2 - assert terrain.terrain_levels_wp.numpy()[0] == 2 + assert terrain.terrain_levels_pa.warp.numpy()[0] == 2 replacement_levels = wp.array([1, 0, 2, 1], dtype=wp.int64, device="cpu") - wp.copy(terrain.terrain_levels_wp, replacement_levels) + wp.copy(terrain.terrain_levels_pa.warp, replacement_levels) wp.synchronize() torch.testing.assert_close(terrain.terrain_levels, torch.tensor([1, 0, 2, 1])) @@ -221,25 +220,25 @@ def test_terrain_importer_grid_origins_cache_only_environment_view(): assert terrain.terrain_origins is None assert isinstance(terrain.env_origins, torch.Tensor) assert terrain.env_origins.shape == (4, 3) - assert terrain.env_origins_wp.shape == (4,) - assert terrain.env_origins.data_ptr() == terrain.env_origins_wp.ptr - assert terrain._terrain_levels_wp is None - assert terrain._terrain_types_wp is None - assert terrain._terrain_origins_wp is None + assert terrain.env_origins_pa.warp.shape == (4,) + assert terrain.env_origins.data_ptr() == terrain.env_origins_pa.warp.ptr + assert terrain.terrain_levels is None + assert terrain.terrain_types is None + assert terrain.terrain_origins is None def test_terrain_importer_reconfigure_preserves_cached_views(): """Reconfiguring origins must reuse buffer storage so cached views stay valid.""" terrain = _make_terrain([0, 1, 2, 1]) - env_origins_wp = terrain.env_origins_wp - levels_wp = terrain.terrain_levels_wp + env_origins_wp = terrain.env_origins_pa.warp + levels_wp = terrain.terrain_levels_pa.warp env_origins_ptr = terrain.env_origins.data_ptr() shifted_origins = terrain.terrain_origins.clone() + 1.0 terrain.configure_env_origins(shifted_origins) - assert terrain.env_origins_wp is env_origins_wp - assert terrain.terrain_levels_wp is levels_wp + assert terrain.env_origins_pa.warp is env_origins_wp + assert terrain.terrain_levels_pa.warp is levels_wp assert terrain.env_origins.data_ptr() == env_origins_ptr torch.testing.assert_close(terrain.terrain_origins, shifted_origins) @@ -253,8 +252,8 @@ def test_terrain_importer_mask_update_preserves_sparse_and_stable_semantics(): rng_state = wp.array(np.arange(6, dtype=np.uint32) + 71, device="cpu") rng_before = rng_state.numpy().copy() origins_before = terrain.env_origins.clone() - levels_wp = terrain.terrain_levels_wp - origins_wp = terrain.env_origins_wp + levels_wp = terrain.terrain_levels_pa.warp + origins_wp = terrain.env_origins_pa.warp levels_ptr = levels_wp.ptr origins_ptr = origins_wp.ptr assert terrain.terrain_levels.data_ptr() == levels_ptr @@ -276,10 +275,10 @@ def test_terrain_importer_mask_update_preserves_sparse_and_stable_semantics(): torch.testing.assert_close(terrain.env_origins[4:], origins_before[4:]) assert np.all(rng_state.numpy()[:4] != rng_before[:4]) np.testing.assert_array_equal(rng_state.numpy()[4:], rng_before[4:]) - assert terrain.terrain_levels_wp is levels_wp - assert terrain.env_origins_wp is origins_wp - assert terrain.terrain_levels_wp.ptr == levels_ptr - assert terrain.env_origins_wp.ptr == origins_ptr + assert terrain.terrain_levels_pa.warp is levels_wp + assert terrain.env_origins_pa.warp is origins_wp + assert terrain.terrain_levels_pa.warp.ptr == levels_ptr + assert terrain.env_origins_pa.warp.ptr == origins_ptr levels_before_stable_update = terrain.terrain_levels.clone() env_ids = torch.tensor([1, 3], dtype=torch.int64) @@ -291,10 +290,10 @@ def test_terrain_importer_mask_update_preserves_sparse_and_stable_semantics(): expected_levels = levels_before_stable_update.clone() expected_levels[env_ids] -= 1 torch.testing.assert_close(terrain.terrain_levels, expected_levels) - assert terrain.terrain_levels_wp is levels_wp - assert terrain.env_origins_wp is origins_wp - assert terrain.terrain_levels_wp.ptr == levels_ptr - assert terrain.env_origins_wp.ptr == origins_ptr + assert terrain.terrain_levels_pa.warp is levels_wp + assert terrain.env_origins_pa.warp is origins_wp + assert terrain.terrain_levels_pa.warp.ptr == levels_ptr + assert terrain.env_origins_pa.warp.ptr == origins_ptr def test_curriculum_manager_updates_masked_levels_and_logs_without_host_compaction(monkeypatch: pytest.MonkeyPatch): diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py index fedf35f23237..3500e1929d6b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py @@ -385,7 +385,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.rpy = wp.zeros((self.num_envs), dtype=wp.vec3f, device=self.sim.device) self.angle_to_target = wp.zeros((self.num_envs), dtype=wp.float32, device=self.sim.device) self.dof_pos_scaled = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device) - self.env_origins = self.scene.env_origins_wp + self.env_origins = self.scene.env_origins_pa.warp self.actions_mapped = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device) # Initial states and targets diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/inhand_manipulation_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/inhand_manipulation_warp_env.py index 2950143cc114..59a130a8b63e 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/inhand_manipulation_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/inhand_manipulation_warp_env.py @@ -582,7 +582,7 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa self.y_unit_vec = wp.vec3f(0.0, 1.0, 0.0) self.z_unit_vec = wp.vec3f(0.0, 0.0, 1.0) - self.env_origins = self.scene.env_origins_wp + self.env_origins = self.scene.env_origins_pa.warp # --------------------------------------------------------------------- # Warp buffers diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py index 1b96ad087c1c..f59c34696154 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py @@ -75,7 +75,7 @@ def __init__(self, cfg: CurriculumTermCfg, env: ManagerBasedRLEnv): self._root_pos_w_wp = self._asset.data.root_pos_w.warp self._env_origins_wp = env.env_origins_wp self._command_wp = env.command_manager.get_command_wp("base_velocity") - self._terrain_levels_wp = self._terrain.terrain_levels_wp + self._terrain_levels_wp = self._terrain.terrain_levels_pa.warp self._move_up_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) self._move_down_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) self._terrain_half_length = float(self._terrain.cfg.terrain_generator.size[0]) * 0.5 From 33eac2eaf917ec77f380edfa8475915d1bdd80f0 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 22 Jul 2026 18:02:41 -0700 Subject: [PATCH 57/58] Document rough-terrain velocity coverage on the warp frontend The terrain-levels curriculum runs as a Warp-native update, so the rough velocity tasks (AnymalD, Cassie, G1, H1, UnitreeGo2) adapt to the warp frontend. Note Digit as the one uncovered rough task (no warp twin for its desired_contacts reward). --- .../physical-backends/newton/warp-environments.rst | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst index 610d2af98fa1..5ada16e04bce 100644 --- a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst +++ b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst @@ -57,11 +57,15 @@ MDP term for its warp twin). Select the Newton solver explicitly with - ``Isaac-Reach-Franka``, ``Isaac-Reach-UR10`` - ``Isaac-Velocity-Flat-AnymalD``, ``-Cassie``, ``-G1``, ``-H1``, ``-UnitreeGo2`` (and their ``-Play`` variants) +- ``Isaac-Velocity-Rough-AnymalD``, ``-Cassie``, ``-G1``, ``-H1``, ``-UnitreeGo2`` + (and their ``-Play`` variants) — the terrain-levels curriculum runs as a + Warp-native update on the :class:`~isaaclab.terrains.TerrainImporter` Warp + origin views A missing twin is a hard error listing the affected terms, so a partially covered task fails at build time rather than silently changing behavior. -Rough-terrain velocity tasks remain unsupported until -:class:`~isaaclab.terrains.TerrainImporter` gains Warp APIs. +``Isaac-Velocity-Rough-Digit`` is the one rough task not yet covered: its +``desired_contacts`` reward term has no Warp twin. Quick Start From c9713641cd89afac4245477234c51af32639a2cc Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 22 Jul 2026 18:02:42 -0700 Subject: [PATCH 58/58] Pin rough velocity tasks in the frontend conversion sweep --- .../test/envs/test_frontend_cfg_conversion.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index 82ec8489de9c..e854739f78c1 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -55,6 +55,16 @@ "Isaac-Velocity-Flat-H1-Play", "Isaac-Velocity-Flat-UnitreeGo2", "Isaac-Velocity-Flat-UnitreeGo2-Play", + "Isaac-Velocity-Rough-AnymalD", + "Isaac-Velocity-Rough-AnymalD-Play", + "Isaac-Velocity-Rough-Cassie", + "Isaac-Velocity-Rough-Cassie-Play", + "Isaac-Velocity-Rough-G1", + "Isaac-Velocity-Rough-G1-Play", + "Isaac-Velocity-Rough-H1", + "Isaac-Velocity-Rough-H1-Play", + "Isaac-Velocity-Rough-UnitreeGo2", + "Isaac-Velocity-Rough-UnitreeGo2-Play", ] # Manager-based warp tasks are exactly those whose entry point is the shared warp