# 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 contextlib
import inspect
import logging
import math
import os
from abc import abstractmethod
from dataclasses import MISSING
from typing import Any, ClassVar
import gymnasium as gym
import numpy as np
import torch
# import omni.kit.app
# import omni.log
# import omni.physx
import warp as wp
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.utils.spaces import sample_space, spec_to_gym_space
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_experimental.envs.interactive_scene_warp import InteractiveSceneWarp
from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache
# from isaacsim.core.simulation_manager import SimulationManager
# from isaacsim.core.version import get_version
# import logger
logger = logging.getLogger(__name__)
DEBUG_TIMER_STEP = os.environ.get("DEBUG_TIMER_STEP", "0") == "1"
"""Enable outer step() timer only. Set DEBUG_TIMER_STEP=1 env var to enable."""
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
[docs]
class DirectRLEnvWarp(DirectRLEnv):
"""The superclass for the direct workflow to design environments.
This class implements the core functionality for reinforcement learning (RL)
environments. It is designed to be used with any RL library. The class is designed
to be used with vectorized environments, i.e., the environment is expected to be run
in parallel with multiple sub-environments.
While the environment itself is implemented as a vectorized environment, we do not
inherit from :class:`gym.vector.VectorEnv`. This is mainly because the class adds
various methods (for wait and asynchronous updates) which are not required.
Additionally, each RL library typically has its own definition for a vectorized
environment. Thus, to reduce complexity, we directly use the :class:`gym.Env` over
here and leave it up to library-defined wrappers to take care of wrapping this
environment for their agents.
Note:
For vectorized environments, it is recommended to **only** call the :meth:`reset`
method once before the first call to :meth:`step`, i.e. after the environment is created.
After that, the :meth:`step` function handles the reset of terminated sub-environments.
in a vectorized environment.
"""
is_vector_env: ClassVar[bool] = True
"""Whether the environment is a vectorized environment."""
metadata: ClassVar[dict[str, Any]] = {
"render_modes": [None, "human", "rgb_array"],
# "isaac_sim_version": get_version(),
}
"""Metadata for the environment."""
[docs]
def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs):
"""Initialize the environment.
Args:
cfg: The configuration object for the environment.
render_mode: The render mode for the environment. Defaults to None, which
is similar to ``"human"``.
Raises:
RuntimeError: If a simulation context already exists. The environment must always create one
since it configures the simulation context and controls the simulation.
"""
# check that the config is valid
cfg.validate()
# store inputs to class
self.cfg = cfg
# store the render mode
self.render_mode = render_mode
# initialize internal variables
self._is_closed = False
# set the seed for the environment
if self.cfg.seed is not None:
self.cfg.seed = self.seed(self.cfg.seed)
else:
logger.warning("Seed not set for the environment. The environment creation may not be deterministic.")
# create a simulation context to control the simulator
if SimulationContext.instance() is None:
self.sim: SimulationContext = SimulationContext(self.cfg.sim)
else:
raise RuntimeError("Simulation context already exists. Cannot create a new one.")
# make sure torch is running on the correct device
if "cuda" in self.device:
torch.cuda.set_device(self.device)
# print useful information
print("[INFO]: Base environment:")
print(f"\tEnvironment device : {self.device}")
print(f"\tEnvironment seed : {self.cfg.seed}")
print(f"\tPhysics step-size : {self.physics_dt}")
print(f"\tRendering step-size : {self.physics_dt * self.cfg.sim.render_interval}")
print(f"\tEnvironment step-size : {self.step_dt}")
if self.cfg.sim.render_interval < self.cfg.decimation:
msg = (
f"The render interval ({self.cfg.sim.render_interval}) is smaller than the decimation "
f"({self.cfg.decimation}). Multiple render calls will happen for each environment step."
"If this is not intended, set the render interval to be equal to the decimation."
)
logger.warning(msg)
# generate scene
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._setup_scene()
# attach_stage_to_usd_context()
print("[INFO]: Scene manager: ", self.scene)
# 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
if self.cfg.events:
self.event_manager = EventManager(self.cfg.events, self)
# apply USD-related randomization events
if "prestartup" in self.event_manager.available_modes:
self.event_manager.apply(mode="prestartup")
# play the simulator to activate physics handles
# note: this activates the physics simulation view that exposes TensorAPIs
# note: when started in extension mode, first call sim.reset_async() and then initialize the managers
# if builtins.ISAAC_LAUNCHED_FROM_TERMINAL is False:
# print("[INFO]: Starting the simulation. This may take a few seconds. Please wait...")
with Timer("[INFO]: Time taken for simulation start", "simulation_start"):
# since the reset can trigger callbacks which use the stage,
# we need to set the stage context here
with use_stage(self.sim.stage):
self.sim.reset()
# update scene to pre populate data buffers for assets and sensors.
# this is needed for the observation manager to get valid tensors for initialization.
# this shouldn't cause an issue since later on, users do a reset over all the
# environments so the lazy buffers would be reset.
self.scene.update(dt=self.physics_dt)
# check if debug visualization is has been implemented by the environment
source_code = inspect.getsource(self._set_debug_vis_impl)
self.has_debug_vis_implementation = "NotImplementedError" not in source_code
self._debug_vis_handle = None
# extend UI elements
# we need to do this here after all the managers are initialized
# this is because they dictate the sensors and commands right now
if self.sim.has_gui and self.cfg.ui_window_class_type is not None:
self._window = self.cfg.ui_window_class_type(self, window_name="IsaacLab")
else:
# if no window, then we don't need to store the window
self._window = None
# allocate dictionary to store metrics
self.extras = {}
# initialize data and constants
# -- counter for simulation steps
self._sim_step_counter = 0
# -- counter for curriculum
self.common_step_counter = 0
# -- init buffers
self._episode_length_buf_wp = wp.zeros(self.num_envs, dtype=wp.int32, device=self.device)
self._episode_length_buf_torch = wp.to_torch(self._episode_length_buf_wp)
self.reset_terminated = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device)
self.reset_time_outs = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device)
self.reset_buf = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device)
self._ALL_ENV_MASK = wp.ones(self.num_envs, dtype=wp.bool, device=self.device)
# Expected bindings:
self.torch_obs_buf: torch.Tensor = None
self.torch_reward_buf: torch.Tensor = None
self.torch_reset_terminated: torch.Tensor = None
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()
# setup the action and observation spaces for Gym
self._configure_gym_env_spaces()
# setup noise cfg for adding action and observation noise
if self.cfg.action_noise_model:
self._action_noise_model: NoiseModel = self.cfg.action_noise_model.class_type(
self.cfg.action_noise_model, num_envs=self.num_envs, device=self.device
)
if self.cfg.observation_noise_model:
self._observation_noise_model: NoiseModel = self.cfg.observation_noise_model.class_type(
self.cfg.observation_noise_model, num_envs=self.num_envs, device=self.device
)
# perform events at the start of the simulation
if self.cfg.events:
# we print it here to make the logging consistent
print("[INFO] Event Manager: ", self.event_manager)
if "startup" in self.event_manager.available_modes:
self.event_manager.apply(mode="startup")
# set the framerate of the gym video recorder wrapper so that the playback speed of the produced
# video matches the simulation
self.metadata["render_fps"] = 1 / self.step_dt
# print the environment information
print("[INFO]: Completed setting up the environment...")
def __del__(self):
"""Cleanup for the environment."""
# Suppress errors during Python shutdown to avoid noisy tracebacks
# Note: contextlib may be None during interpreter shutdown
if contextlib is not None:
with contextlib.suppress(ImportError, AttributeError, TypeError):
self.close()
"""
Properties.
"""
@property
def num_envs(self) -> int:
"""The number of instances of the environment that are running."""
return self.scene.num_envs
@property
def physics_dt(self) -> float:
"""The physics time-step (in s).
This is the lowest time-decimation at which the simulation is happening.
"""
return self.cfg.sim.dt
@property
def step_dt(self) -> float:
"""The environment stepping time-step (in s).
This is the time-step at which the environment steps forward.
"""
return self.cfg.sim.dt * self.cfg.decimation
@property
def device(self):
"""The device on which the environment is running."""
return self.sim.device
@property
def max_episode_length_s(self) -> float:
"""Maximum episode length in seconds."""
return self.cfg.episode_length_s
@property
def max_episode_length(self):
"""The maximum episode length in steps adjusted from s."""
return math.ceil(self.max_episode_length_s / (self.cfg.sim.dt * self.cfg.decimation))
@property
def episode_length_buf(self) -> torch.Tensor:
"""The episode length buffer as a torch tensor.
This is a view of the underlying warp array ``_episode_length_buf_wp``.
Setting this property copies values in-place to preserve the shared
memory link with warp kernels.
"""
return self._episode_length_buf_torch
@episode_length_buf.setter
def episode_length_buf(self, value: torch.Tensor):
self._episode_length_buf_torch.copy_(value)
"""
Operations.
"""
def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[VecEnvObs, dict]:
"""Resets all the environments and returns observations.
This function calls the :meth:`_reset_idx` function to reset all the environments.
However, certain operations, such as procedural terrain generation, that happened during initialization
are not repeated.
Args:
seed: The seed to use for randomization. Defaults to None, in which case the seed is not set.
options: Additional information to specify how the environment is reset. Defaults to None.
Note:
This argument is used for compatibility with Gymnasium environment definition.
Returns:
A tuple containing the observations and extras.
"""
# set the seed
if seed is not None:
self.seed(seed)
# reset state of scene
self._reset_idx(self._ALL_ENV_MASK)
# update articulation kinematics
self.scene.write_data_to_sim()
# if sensors are added to the scene, make sure we render to reflect changes in reset
if hasattr(self.sim, "has_rtx_sensors") and self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset:
self.sim.render()
# if self.cfg.wait_for_textures and self.sim.has_rtx_sensors():
# while SimulationManager.assets_loading():
# self.sim.render()
# return observations
self._get_observations()
# store the returned buffer so RslRlVecEnvWrapper.get_observations() can read env.obs_buf
self.obs_buf = {"policy": self.torch_obs_buf.clone()}
return self.obs_buf, self.extras
@Timer(name="env_step", msg="Step took:", enable=DEBUG_TIMER_STEP or DEBUG_TIMERS)
def step(self, action: torch.Tensor) -> VecEnvStepReturn:
"""Execute one time-step of the environment's dynamics.
The environment steps forward at a fixed time-step, while the physics simulation is decimated at a
lower time-step. This is to ensure that the simulation is stable. These two time-steps can be configured
independently using the :attr:`DirectRLEnvCfg.decimation` (number of simulation steps per environment step)
and the :attr:`DirectRLEnvCfg.sim.physics_dt` (physics time-step). Based on these parameters, the environment
time-step is computed as the product of the two.
This function performs the following steps:
1. Pre-process the actions before stepping through the physics.
2. Apply the actions to the simulator and step through the physics in a decimated manner.
3. Compute the reward and done signals.
4. Reset environments that have terminated or reached the maximum episode length.
5. Apply interval events if they are enabled.
6. Compute observations.
Args:
action: The actions to apply on the environment. Shape is (num_envs, action_dim).
Returns:
A tuple containing the observations, rewards, resets (terminated and truncated) and extras.
"""
action = action.to(self.device)
# add action noise
if self.cfg.action_noise_model:
action = self._action_noise_model(action)
# process actions, #TODO pass the torch tensor directly.
with Timer(name="pre_physics", msg="Pre-physics step took:", enable=DEBUG_TIMERS):
self._pre_physics_step(
wp.from_torch(action)
) # Creates a tensor and discards it. Not graphable unless training loop reuses the same pointer.
# check if we need to do rendering within the physics loop
# note: hoisted out of the decimation loop; is_rendering does live settings lookups
is_rendering = self.sim.is_rendering
# perform physics stepping
with Timer(name="physics_loop", msg="Physics loop took:", enable=DEBUG_TIMERS):
for _ in range(self.cfg.decimation):
self._sim_step_counter += 1
# 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)
# 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.
with Timer(name="write_data_to_sim_loop", msg="Write data to sim (loop) took:", enable=DEBUG_TIMERS):
self.scene.write_data_to_sim()
with Timer(name="simulate", msg="Newton simulation step took:", enable=DEBUG_TIMERS):
self.sim.step(render=False)
# 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.
if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering:
self.sim.render()
# update buffers at sim dt
with Timer(name="scene_update", msg="Scene update took:", enable=DEBUG_TIMERS):
self.scene.update(dt=self.physics_dt)
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.
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)
# Visualization hook — runs after CUDA graph scope. Override in subclass
# to update markers or other non-graphable visual elements.
with Timer(name="visualize", msg="Visualize took:", enable=DEBUG_TIMERS):
self._post_step_visualize()
# return observations, rewards, resets and extras
# store the returned buffer so RslRlVecEnvWrapper.get_observations() can read env.obs_buf
self.obs_buf = {"policy": self.torch_obs_buf.clone()}
return (
self.obs_buf,
self.torch_reward_buf,
self.torch_reset_terminated,
self.torch_reset_time_outs,
self.extras,
)
def _post_step_visualize(self) -> None:
"""Hook for updating visualization markers after CUDA graph scope.
Override in subclass to update markers or other non-graphable visual
elements (e.g., those requiring wp.to_torch + .cpu().numpy()).
This runs every step, outside any CUDA graph capture.
"""
pass
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.
def _step_warp_end_pre(self) -> None:
"""Capturable portion before write_data_to_sim (pure warp kernels)."""
wp.launch(
add_to_env,
dim=self.num_envs,
inputs=[
self._episode_length_buf_wp,
1,
],
)
self._get_dones()
self._get_rewards()
# -- reset envs that terminated/timed-out and log the episode information
self._reset_idx(mask=self.reset_buf)
def _step_warp_end_post(self) -> None:
"""Capturable portion after write_data_to_sim (pure warp kernels)."""
# if sensors are added to the scene, make sure we render to reflect changes in reset
# if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset:
# self.sim.render()
# TODO We could split it out.
# post-step: step interval event
# if self.cfg.events:
# if "interval" in self.event_manager.available_modes:
# self.event_manager.apply(mode="interval", dt=self.step_dt)
# update observations
self._get_observations()
# add observation noise
# note: we apply no noise to the state space (since it is used for critic networks)
# if self.cfg.observation_noise_model:
# self.obs_buf["policy"] = self._observation_noise_model(self.obs_buf["policy"])
@staticmethod
def seed(seed: int = -1) -> int:
"""Set the seed for the environment.
Args:
seed: The seed for random generator. Defaults to -1.
Returns:
The seed used for random generator.
"""
# set seed for replicator
try:
import omni.replicator.core as rep
rep.set_global_seed(seed)
except ModuleNotFoundError:
pass
# set seed for torch and other libraries
return configure_seed(seed)
def render(self, recompute: bool = False) -> np.ndarray | None:
"""Run rendering without stepping through the physics.
By convention, if mode is:
- **human**: Render to the current display and return nothing. Usually for human consumption.
- **rgb_array**: Return an numpy.ndarray with shape (x, y, 3), representing RGB values for an
x-by-y pixel image, suitable for turning into a video.
Args:
recompute: Whether to force a render even if the simulator has already rendered the scene.
Defaults to False.
Returns:
The rendered image as a numpy array if mode is "rgb_array". Otherwise, returns None.
Raises:
RuntimeError: If mode is set to "rgb_data" and simulation render mode does not support it.
In this case, the simulation render mode must be set to ``RenderMode.PARTIAL_RENDERING``
or ``RenderMode.FULL_RENDERING``.
NotImplementedError: If an unsupported rendering mode is specified.
"""
# run a rendering step of the simulator
# if we have rtx sensors, we do not need to render again sim
if not (hasattr(self.sim, "has_rtx_sensors") and self.sim.has_rtx_sensors()) and not recompute:
self.sim.render()
# decide the rendering mode
if self.render_mode == "human" or self.render_mode is None:
return None
elif self.render_mode == "rgb_array":
# 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(
f"Cannot render '{self.render_mode}' when the simulation render mode is"
f" '{render_mode_name}'. Please set the simulation render mode"
" to:'PARTIAL_RENDERING' or"
" 'FULL_RENDERING'."
)
# create the annotator if it does not exist
if not hasattr(self, "_rgb_annotator"):
import omni.replicator.core as rep
# create render product from the main Kit viewport camera
_cam_prim_path = "/OmniverseKit_Persp"
_resolution = (1280, 720)
self._render_product = rep.create.render_product(_cam_prim_path, _resolution)
self._render_resolution = _resolution
# create rgb annotator -- used to read data from the render product
self._rgb_annotator = rep.AnnotatorRegistry.get_annotator("rgb", device="cpu")
self._rgb_annotator.attach([self._render_product])
# obtain the rgb data
rgb_data = self._rgb_annotator.get_data()
# convert to numpy array
rgb_data = np.frombuffer(rgb_data, dtype=np.uint8).reshape(*rgb_data.shape)
# return the rgb data
# note: initially the renerer is warming up and returns empty data
if rgb_data.size == 0:
return np.zeros((self._render_resolution[1], self._render_resolution[0], 3), dtype=np.uint8)
else:
return rgb_data[:, :, :3]
else:
raise NotImplementedError(
f"Render mode '{self.render_mode}' is not supported. Please use: {self.metadata['render_modes']}."
)
def close(self):
"""Cleanup for the environment."""
if not self._is_closed:
# close entities related to the environment
# note: this is order-sensitive to avoid any dangling references
if self.cfg.events:
del self.event_manager
del self.scene
# # clear callbacks and instance
# if float(".".join(get_version()[2])) >= 5:
# if self.cfg.sim.create_stage_in_memory:
# # detach physx stage
# omni.physx.get_physx_simulation_interface().detach_stage()
# self.sim.stop()
# self.sim.clear()
# self.sim.clear_all_callbacks()
self.sim.clear_instance()
# destroy the window
if self._window is not None:
self._window = None
# update closing status
self._is_closed = True
"""
Operations - Debug Visualization.
"""
def set_debug_vis(self, debug_vis: bool) -> bool:
"""Toggles the environment debug visualization.
Args:
debug_vis: Whether to visualize the environment debug visualization.
Returns:
Whether the debug visualization was successfully set. False if the environment
does not support debug visualization.
"""
# check if debug visualization is supported
if not self.has_debug_vis_implementation:
return False
# toggle debug visualization objects
self._set_debug_vis_impl(debug_vis)
# toggle debug visualization handles
if debug_vis:
# create a subscriber for the post update event if it doesn't exist
if self._debug_vis_handle is None:
self._debug_vis_handle = self.sim.vis_marker_registry.add_debug_vis_callback(self)
else:
# remove the subscriber if it exists
self.sim.vis_marker_registry.clear_debug_vis_callback(self)
# return success
return True
"""
Helper functions.
"""
def _configure_gym_env_spaces(self):
"""Configure the action and observation spaces for the Gym environment."""
# show deprecation message and overwrite configuration
if self.cfg.num_actions is not None:
logger.warning("DirectRLEnvCfg.num_actions is deprecated. Use DirectRLEnvCfg.action_space instead.")
if isinstance(self.cfg.action_space, type(MISSING)):
self.cfg.action_space = self.cfg.num_actions
if self.cfg.num_observations is not None:
logger.warning(
"DirectRLEnvCfg.num_observations is deprecated. Use DirectRLEnvCfg.observation_space instead."
)
if isinstance(self.cfg.observation_space, type(MISSING)):
self.cfg.observation_space = self.cfg.num_observations
if self.cfg.num_states is not None:
logger.warning("DirectRLEnvCfg.num_states is deprecated. Use DirectRLEnvCfg.state_space instead.")
if isinstance(self.cfg.state_space, type(MISSING)):
self.cfg.state_space = self.cfg.num_states
# set up spaces
self.single_observation_space = gym.spaces.Dict()
self.single_observation_space["policy"] = spec_to_gym_space(self.cfg.observation_space)
self.single_action_space = spec_to_gym_space(self.cfg.action_space)
# batch the spaces for vectorized environments
self.observation_space = gym.vector.utils.batch_space(self.single_observation_space["policy"], self.num_envs)
self.action_space = gym.vector.utils.batch_space(self.single_action_space, self.num_envs)
# optional state space for asymmetric actor-critic architectures
self.state_space = None
if self.cfg.state_space:
self.single_observation_space["critic"] = spec_to_gym_space(self.cfg.state_space)
self.state_space = gym.vector.utils.batch_space(self.single_observation_space["critic"], self.num_envs)
# instantiate actions (needed for tasks for which the observations computation is dependent on the actions)
self.actions = sample_space(self.single_action_space, self.sim.device, batch_size=self.num_envs, fill_value=0)
def _reset_idx(self, mask: wp.array | None = None):
"""Reset environments based on a boolean mask.
Args:
mask: Boolean mask indicating which environments to reset.
Shape is (num_envs,). If None, all environments are reset.
"""
if mask is None:
mask = self._ALL_ENV_MASK
self.scene.reset(env_ids=None, env_mask=mask)
# apply events such as randomization for environments that need a reset
# if self.cfg.events:
# if "reset" in self.event_manager.available_modes:
# env_step_count = self._sim_step_counter // self.cfg.decimation
# self.event_manager.apply(mode="reset", env_ids=env_ids, global_env_step_count=env_step_count)
# reset noise models
# if self.cfg.action_noise_model:
# self._action_noise_model.reset(env_ids)
# if self.cfg.observation_noise_model:
# self._observation_noise_model.reset(env_ids)
# reset the episode length buffer
wp.launch(
zero_mask_int32,
dim=self.num_envs,
inputs=[
mask,
self._episode_length_buf_wp,
],
)
"""
Implementation-specific functions.
"""
def _setup_scene(self):
"""Setup the scene for the environment.
This function is responsible for creating the scene objects and setting up the scene for the environment.
The scene creation can happen through :class:`isaaclab.scene.InteractiveSceneCfg` or through
directly creating the scene objects and registering them with the scene manager.
We leave the implementation of this function to the derived classes. If the environment does not require
any explicit scene setup, the function can be left empty.
"""
pass
@abstractmethod
def _pre_physics_step(self, actions: wp.array) -> None:
"""Pre-process actions before stepping through the physics.
This function is responsible for pre-processing the actions before stepping through the physics.
It is called before the physics stepping (which is decimated).
Args:
actions: The actions to apply on the environment. Shape is (num_envs, action_dim).
"""
raise NotImplementedError(f"Please implement the '_pre_physics_step' method for {self.__class__.__name__}.")
@abstractmethod
def _apply_action(self) -> None:
"""Apply actions to the simulator.
This function is responsible for applying the actions to the simulator. It is called at each
physics time-step. Must be pure warp (no torch ops) to be CUDA graph capturable.
"""
raise NotImplementedError(f"Please implement the '_apply_action' method for {self.__class__.__name__}.")
@abstractmethod
def _get_observations(self) -> dict:
"""Compute and return the observations for the environment.
Returns:
The observations dictionary, e.g. ``{"policy": tensor}``.
"""
raise NotImplementedError(f"Please implement the '_get_observations' method for {self.__class__.__name__}.")
def _get_states(self) -> VecEnvObs | None:
"""Compute and return the states for the environment.
The state-space is used for asymmetric actor-critic architectures. It is configured
using the :attr:`DirectRLEnvCfg.state_space` parameter.
Returns:
The states for the environment. If the environment does not have a state-space, the function
returns a None.
"""
return None # noqa: R501
@abstractmethod
def _get_rewards(self) -> None:
"""Compute the rewards for the environment.
Writes results into the reward buffer (e.g., ``self.reward_buf``).
"""
raise NotImplementedError(f"Please implement the '_get_rewards' method for {self.__class__.__name__}.")
@abstractmethod
def _get_dones(self) -> None:
"""Compute the done flags for the environment.
Writes results into the done buffers (e.g., ``self.reset_terminated``, ``self.reset_time_outs``).
"""
raise NotImplementedError(f"Please implement the '_get_dones' method for {self.__class__.__name__}.")
def _set_debug_vis_impl(self, debug_vis: bool):
"""Set debug visualization into visualization objects.
This function is responsible for creating the visualization objects if they don't exist
and input ``debug_vis`` is True. If the visualization objects exist, the function should
set their visibility into the stage.
"""
raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.")
