env_test_utils.py

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

"""Shared test utilities for Isaac Lab environments."""

import gc
import importlib
import os
import sys

import gymnasium as gym
import pytest
import torch

import isaaclab.sim as sim_utils
from isaaclab.app.settings_manager import get_settings_manager
from isaaclab.envs.mdp.actions.actions_cfg import OperationalSpaceControllerActionCfg
from isaaclab.envs.utils.spaces import sample_space
from isaaclab.sim import SimulationContext
from isaaclab.utils.version import get_isaac_sim_version

from isaaclab_tasks.utils.hydra import apply_overrides, collect_presets
from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry, parse_env_cfg

# Map of task IDs to the reason for marking the corresponding parametrized
# test cases as expected failures.  Tests that consume :func:`setup_environment`
# automatically pick up these marks via :class:`pytest.param`.
XFAIL_TASKS: dict[str, str] = {}


def _is_teleop_env(task_spec) -> bool:
    """Check if a task's environment config has teleop dependencies.

    Inspects the class hierarchy of the env config to check if any base
    class module defines ``_TELEOP_AVAILABLE``, indicating the environment
    uses isaacteleop / isaaclab_teleop.
    """
    env_cfg_entry_point = task_spec.kwargs.get("env_cfg_entry_point")
    if not isinstance(env_cfg_entry_point, str) or ":" not in env_cfg_entry_point:
        return False
    try:
        mod_name, attr_name = env_cfg_entry_point.split(":")
        mod = importlib.import_module(mod_name)
        cfg_cls = getattr(mod, attr_name, None)
        if cfg_cls is None:
            return False
        for cls in cfg_cls.__mro__:
            cls_module = sys.modules.get(cls.__module__)
            if cls_module is not None and hasattr(cls_module, "_TELEOP_AVAILABLE"):
                return True
    except (ImportError, AttributeError):
        pass
    return False


def _is_pickplace_stack_env(task_id: str) -> bool:
    """Check if a task is a PickPlace or Stack environment based on its ID."""
    return any(keyword in task_id for keyword in ("Place", "Stack", "NutPour", "ExhaustPipe"))


def _task_tier(task_spec) -> str | None:
    """Return ``"core"`` or ``"contrib"`` based on the task's env-config entry-point module.

    Core tasks register their config under :mod:`isaaclab_tasks.core` and contributed
    tasks under :mod:`isaaclab_tasks.contrib`. Returns ``None`` if the tier cannot be
    determined from the registered entry point.
    """
    entry = task_spec.kwargs.get("env_cfg_entry_point")
    if isinstance(entry, str):
        module = entry.split(":")[0]
        if module.startswith("isaaclab_tasks.core"):
            return "core"
        if module.startswith("isaaclab_tasks.contrib"):
            return "contrib"
    return None


def setup_environment(
    factory_envs: bool | None = None,
    multi_agent: bool | None = None,
    teleop_envs: bool | None = None,
    cartpole_showcase_envs: bool | None = None,
    pickplace_stack_envs: bool | None = None,
    tier: str | None = None,
) -> list[str]:
    """
    Acquire all registered Isaac environment task IDs with optional filters.

    Args:
        factory_envs:
            - True: include only Factory environments
            - False: exclude Factory environments
            - None: include both Factory and non-Factory environments
        multi_agent:
            - True: include only multi-agent environments
            - False: include only single-agent environments
            - None: include all environments regardless of agent type
        teleop_envs:
            - True: include only teleop environments (those requiring isaacteleop)
            - False: exclude teleop environments
            - None: include all environments regardless of teleop dependency
        cartpole_showcase_envs:
            - True: include only Cartpole Showcase environments
            - False: exclude Cartpole Showcase environments
            - None: include all environments regardless of showcase type
        pickplace_stack_envs:
            - True: include only PickPlace/Stack environments
            - False: exclude PickPlace/Stack environments
            - None: include all environments regardless of pick-place/stack type
        tier:
            - "core": include only core environments (registered under ``isaaclab_tasks.core``).
            - "contrib": include only contributed environments (registered under ``isaaclab_tasks.contrib``).
            - None: include all environments regardless of tier.

    Returns:
        A sorted list of task IDs matching the selected filters.
    """
    # disable interactive mode for wandb for automate environments
    os.environ["WANDB_DISABLED"] = "true"

    # acquire all Isaac environment names
    registered_tasks = []
    for task_spec in gym.registry.values():
        # only consider Isaac environments
        if "Isaac" not in task_spec.id:
            continue

        # apply core/contrib tier filter
        if tier is not None and _task_tier(task_spec) != tier:
            continue

        # TODO: factory environments cause tests to fail if run together with other envs,
        # so we collect these environments separately to run in a separate unit test.
        # apply factory filter
        if (factory_envs is True and ("Factory" not in task_spec.id and "Forge" not in task_spec.id)) or (
            factory_envs is False and ("Factory" in task_spec.id or "Forge" in task_spec.id)
        ):
            continue
        # if None: no filter

        # apply cartpole showcase filter
        if (cartpole_showcase_envs is True and "Showcase" not in task_spec.id) or (
            cartpole_showcase_envs is False and "Showcase" in task_spec.id
        ):
            continue
        # if None: no filter

        # apply pickplace/stack filter
        if pickplace_stack_envs is not None:
            is_pickplace_stack = _is_pickplace_stack_env(task_spec.id)
            if (pickplace_stack_envs is True and not is_pickplace_stack) or (
                pickplace_stack_envs is False and is_pickplace_stack
            ):
                continue
        # if None: no filter

        # apply teleop filter
        if teleop_envs is not None:
            is_teleop = _is_teleop_env(task_spec)
            if (teleop_envs is True and not is_teleop) or (teleop_envs is False and is_teleop):
                continue
        # if None: no filter

        # apply multi agent filter
        if multi_agent is not None:
            # parse config
            env_cfg = parse_env_cfg(task_spec.id)
            if (multi_agent is True and not hasattr(env_cfg, "possible_agents")) or (
                multi_agent is False and hasattr(env_cfg, "possible_agents")
            ):
                continue
        # if None: no filter

        registered_tasks.append(task_spec.id)

    # sort environments alphabetically
    registered_tasks.sort()

    # this flag is necessary to prevent a bug where the simulation gets stuck randomly when running many environments
    get_settings_manager().set_bool("/physics/cooking/ujitsoCollisionCooking", False)

    print(">>> All registered environments:", registered_tasks)

    # Wrap tasks listed in XFAIL_TASKS in pytest.param so the corresponding
    # parametrized test cases are reported as xfailed instead of failed.
    return [
        pytest.param(task_id, marks=pytest.mark.xfail(reason=XFAIL_TASKS[task_id], strict=False))
        if task_id in XFAIL_TASKS
        else task_id
        for task_id in registered_tasks
    ]


def _fire_all_interval_events_once(env) -> None:
    """Force every interval-mode event term to fire once.

    Invokes :meth:`~isaaclab.managers.EventManager.apply` with ``mode="interval"``
    and a ``dt`` larger than any plausible ``interval_range_s`` upper bound, so the
    trigger condition trips for every term in a single call. The manager re-samples
    ``time_left`` from each term's original ``interval_range_s`` after firing, so
    subsequent ``env.step()`` calls observe original interval timing.

    No-op for envs without an :class:`~isaaclab.managers.EventManager` or
    without any ``interval``-mode terms.

    Args:
        env: A constructed env instance.
    """
    event_manager = getattr(env.unwrapped, "event_manager", None)
    if event_manager is None:
        return
    if "interval" not in event_manager.available_modes:
        return
    # Pass a very large dt for (time_left -= dt) to be less than 1e-6
    event_manager.apply("interval", dt=1e9)


def _configure_osc_smoke_actions(env, actions: torch.Tensor) -> None:
    """Replace absolute OSC targets with matching task commands.

    Random samples from an unbounded action space are not valid absolute poses: their
    quaternions are not normalized and their positions may be unreachable. When an
    operational-space action term has a unique pose command for the same asset and
    body, this helper tracks that command with the controller's nominal gains.

    Args:
        env: A constructed manager-based environment.
        actions: The sampled actions to update in place.
    """
    action_manager = getattr(env.unwrapped, "action_manager", None)
    command_manager = getattr(env.unwrapped, "command_manager", None)
    if action_manager is None or command_manager is None:
        return

    action_offset = 0
    for term_name, term_dim in zip(action_manager.active_terms, action_manager.action_term_dim):
        action_term = action_manager.get_term(term_name)
        action_cfg = action_term.cfg
        if not isinstance(action_cfg, OperationalSpaceControllerActionCfg):
            action_offset += term_dim
            continue

        controller_cfg = action_cfg.controller_cfg
        if "pose_abs" not in controller_cfg.target_types or action_cfg.task_frame_rel_path is not None:
            action_offset += term_dim
            continue

        pose_commands = []
        for command_name in command_manager.active_terms:
            command_term = command_manager.get_term(command_name)
            if (
                getattr(command_term.cfg, "asset_name", None) == action_cfg.asset_name
                and getattr(command_term.cfg, "body_name", None) == action_cfg.body_name
                and command_term.command.shape[1] == 7
            ):
                pose_commands.append(command_term.command)
        if len(pose_commands) != 1:
            action_offset += term_dim
            continue

        command_offset = 0
        for target_type in controller_cfg.target_types:
            if target_type == "pose_abs":
                pose_command = pose_commands[0]
                actions[:, action_offset + command_offset : action_offset + command_offset + 3] = (
                    pose_command[:, :3] / action_cfg.position_scale
                )
                actions[:, action_offset + command_offset + 3 : action_offset + command_offset + 7] = (
                    pose_command[:, 3:] / action_cfg.orientation_scale
                )
                command_offset += 7
            elif target_type in ("pose_rel", "wrench_abs"):
                command_offset += 6

        if controller_cfg.impedance_mode in ("variable_kp", "variable"):
            stiffness = torch.as_tensor(controller_cfg.motion_stiffness_task, device=actions.device)
            actions[:, action_offset + command_offset : action_offset + command_offset + 6] = (
                stiffness / action_cfg.stiffness_scale
            )
            command_offset += 6
        if controller_cfg.impedance_mode == "variable":
            damping_ratio = torch.as_tensor(controller_cfg.motion_damping_ratio_task, device=actions.device)
            actions[:, action_offset + command_offset : action_offset + command_offset + 6] = (
                damping_ratio / action_cfg.damping_ratio_scale
            )

        action_offset += term_dim


def _run_environments(
    task_name,
    device,
    num_envs,
    num_steps=20,
    multi_agent=False,
    create_stage_in_memory=False,
    disable_clone_in_fabric=False,
    physics_preset_name: str | None = None,
):
    """Run all environments and check environments return valid signals.

    Args:
        task_name: Name of the environment.
        device: Device to use (e.g., 'cuda').
        num_envs: Number of environments.
        num_steps: Number of simulation steps.
        multi_agent: Whether the environment is multi-agent.
        create_stage_in_memory: Whether to create stage in memory.
        disable_clone_in_fabric: Whether to disable fabric cloning.
        physics_preset_name: Name of the physics preset to apply (e.g., 'newton_mjwarp').
            If None, uses the environment's default physics.
    """

    # skip test if stage in memory is not supported
    if get_isaac_sim_version().major < 5 and create_stage_in_memory:
        pytest.skip("Stage in memory is not supported in this version of Isaac Sim")

    # skip suction gripper environments as they require CPU simulation and cannot be run with GPU simulation
    if "Suction" in task_name and device != "cpu":
        return

    # skip these environments as they cannot be run with 32 environments within reasonable VRAM
    if num_envs == 32 and task_name in [
        "IsaacContrib-Stack-Cube-Franka-IK-Rel-Blueprint",
        "IsaacContrib-Stack-Cube-Instance-Randomize-Franka-IK-Rel",
        "IsaacContrib-Stack-Cube-Instance-Randomize-Franka",
        "IsaacContrib-PickPlace-G1-InspireFTP-Abs",
    ]:
        return

    # these environments are using SingleArticulation class, which need to be updated
    if "RmpFlow" in task_name or "Isaac-Stack-Cube-Galbot-Left-Arm-Gripper-Visuomotor" in task_name:
        return

    # skip these environments as they cannot be run with 32 environments within reasonable VRAM
    if "Visuomotor" in task_name and num_envs == 32:
        return

    # skip automate environments as they require cuda installation
    if task_name in ["IsaacContrib-AutoMate-Assembly-Direct", "IsaacContrib-AutoMate-Disassembly-Direct"]:
        return

    # skip skillgen environments as they require cuRobo installation;
    # tested separately via test_environments_skillgen.py
    if "Skillgen" in task_name:
        return

    print(f""">>> Running test for environment: {task_name}""")
    _check_random_actions(
        task_name,
        device,
        num_envs,
        num_steps=num_steps,
        multi_agent=multi_agent,
        create_stage_in_memory=create_stage_in_memory,
        disable_clone_in_fabric=disable_clone_in_fabric,
        physics_preset_name=physics_preset_name,
    )
    print(f""">>> Closing environment: {task_name}""")
    print("-" * 80)


def _check_random_actions(
    task_name: str,
    device: str,
    num_envs: int,
    num_steps: int = 20,
    multi_agent: bool = False,
    create_stage_in_memory: bool = False,
    disable_clone_in_fabric: bool = False,
    physics_preset_name: str | None = None,
):
    """Run random actions and check environments return valid signals.

    Args:
        task_name: Name of the environment.
        device: Device to use (e.g., 'cuda').
        num_envs: Number of environments.
        num_steps: Number of simulation steps.
        multi_agent: Whether the environment is multi-agent.
        create_stage_in_memory: Whether to create stage in memory.
        disable_clone_in_fabric: Whether to disable fabric cloning.
        physics_preset_name: Name of the physics preset to apply (e.g., 'newton_mjwarp').
            If None, uses the environment's default physics.
    """
    # create a new context stage, if stage in memory is not enabled
    if not create_stage_in_memory:
        sim_utils.create_new_stage()

    # reset the rtx sensors setting to False
    get_settings_manager().set_bool("/isaaclab/render/rtx_sensors", False)
    env = None
    try:
        # parse config
        env_cfg = parse_env_cfg(task_name, device=device, num_envs=num_envs)
        # apply physics preset override before creating the environment
        if physics_preset_name is not None:
            # parse_env_cfg already resolved PresetCfg wrappers to their default,
            # so we load the raw config to retrieve preset alternatives.
            raw_cfg = load_cfg_from_registry(task_name, "env_cfg_entry_point")
            presets = {"env": collect_presets(raw_cfg), "agent": {}}
            hydra_cfg = {"env": env_cfg.to_dict(), "agent": None}
            apply_overrides(env_cfg, None, hydra_cfg, [physics_preset_name], [], [], presets)
            # Re-apply num_envs since apply_overrides may have replaced
            # the scene config with the preset's default num_envs.
            if num_envs is not None:
                env_cfg.scene.num_envs = num_envs
        # set config args
        env_cfg.sim.create_stage_in_memory = create_stage_in_memory
        if disable_clone_in_fabric:
            env_cfg.scene.clone_in_fabric = False

        # filter based off multi agents mode and create env
        if multi_agent:
            if not hasattr(env_cfg, "possible_agents"):
                print(f"[INFO]: Skipping {task_name} as it is not a multi-agent task")
                return
        else:
            if hasattr(env_cfg, "possible_agents"):
                print(f"[INFO]: Skipping {task_name} as it is a multi-agent task")
                return

        # TODO: Selecting the MJWarp preset routes through the Newton backend, which does not yet
        # support multi-asset spawning; some combinations fail config validation here with a
        # ValueError. Consider filtering invalid combinations in setup_environment() rather than
        # forgiving them at runtime. See PR #5097 commit fb2c74a3862 for a workaround that caught
        # the error and called pytest.skip().
        env = gym.make(task_name, cfg=env_cfg)

        # disable control on stop
        env.unwrapped.sim._app_control_on_stop_handle = None  # type: ignore

        # reset environment
        obs, _ = env.reset()

        # check signal
        assert _check_valid_tensor(obs)

        _fire_all_interval_events_once(env)

        # simulate environment for num_steps
        with torch.inference_mode():
            for _ in range(num_steps):
                # sample actions according to the defined space
                if multi_agent:
                    actions = {
                        agent: sample_space(
                            env.unwrapped.action_spaces[agent], device=env.unwrapped.device, batch_size=num_envs
                        )
                        for agent in env.unwrapped.possible_agents
                    }
                else:
                    actions = sample_space(
                        env.unwrapped.single_action_space, device=env.unwrapped.device, batch_size=num_envs
                    )
                    _configure_osc_smoke_actions(env, actions)
                # apply actions
                transition = env.step(actions)
                # check signals
                for data in transition[:-1]:  # exclude info
                    if multi_agent:
                        for agent, agent_data in data.items():
                            assert _check_valid_tensor(agent_data), f"Invalid data ('{agent}'): {agent_data}"
                    else:
                        assert _check_valid_tensor(data), f"Invalid data: {data}"

    finally:
        # Always ensure cleanup happens, regardless of success or failure
        if env is not None:
            env.close()

        # Drop unreachable environment objects while the device is still alive. Warp arrays
        # free device memory from a finalizer, so collect them before the simulation teardown
        # destroys their streams.
        gc.collect()

        # Clear the simulation context singleton (also closes the USD context stage)
        SimulationContext.clear_instance()


def _check_valid_tensor(data: torch.Tensor | dict) -> bool:
    """Checks if given data does not have corrupted values.

    Args:
        data: Data buffer.

    Returns:
        True if the data is valid.
    """
    if isinstance(data, torch.Tensor):
        return not torch.any(torch.isnan(data))
    elif isinstance(data, (tuple, list)):
        return all(_check_valid_tensor(value) for value in data)
    elif isinstance(data, dict):
        return all(_check_valid_tensor(value) for value in data.values())
    else:
        raise ValueError(f"Input data of invalid type: {type(data)}.")