Capturing sensor frames during training

Capturing sensor frames during training#

Isaac Lab supports saving image-like outputs from scene sensors during training using the CaptureEnvSensors Gymnasium wrapper. When --capture_env_sensors is set to a positive value, Isaac Lab iterates over the environment’s InteractiveScene sensors and writes each sensor’s data.output tensors (for example, rgb, depth, or normals from a Camera). The flag value is the number of parallel environment views to tile into each saved frame grid. This differs from Recording video clips during training, which records a single perspective viewport clip of the scene.

The capture flags are registered on the shared training entrypoints for RSL-RL, RL-Games, Stable Baselines3, and skrl. They are not available on play scripts.

This feature can be enabled using the following command line arguments with ./isaaclab.sh train:

--capture_env_sensors: number of parallel environments to include in each saved frame grid (default: 0, which disables capture)
--capture_env_sensors_length: length of each captured sensor frame window in per-episode steps (default: 200). Set to 0 to disable saving even when --capture_env_sensors is positive.
--capture_env_sensors_interval: interval between captured sensor frame windows in per-episode steps (default: 2000)
--capture_env_sensors_format: output format, either tensorboard (default) or file

Enabling sensor capture automatically enables camera rendering during training, the same as --video, which slows down both startup and runtime performance. Sensor capture can also be used at the same time as --video when both per-sensor frames and a perspective viewport clip are needed.

Example usage:

./isaaclab.sh train --rl_library rsl_rl --task=Isaac-Reorient-Cube-Shadow-Camera-Direct --capture_env_sensors 4 --capture_env_sensors_length 100 --capture_env_sensors_interval 2000 --capture_env_sensors_format file

The captured frames will be saved in the same directory as the training checkpoints, under logs/<rl_workflow>/<task>/<run>/sensor_frames/train. Here <run> is the training run directory (the same timestamped folder used for checkpoints and run.json), not the per-episode index used in output tags and file names.

Overview#

The sensor capture feature wraps the training environment through wrap_training_capture in scripts/reinforcement_learning/common.py and saves frames on reset and step when the current per-episode step falls inside a capture window. For each image-like scene sensor, the wrapper:

reads sensor.data.output and skips sensors whose output is not a dictionary of tensors
resolves each output to a torch.Tensor, including ProxyArray buffers through their .torch accessor
selects the first capture_num_envs parallel environments from each tensor
replaces non-finite values with zero before display normalization
normalizes each output for display using normalize_camera_output_for_display() and tiles the selected views into a single image grid with make_camera_output_grid()
writes the grid to TensorBoard or to disk as a PNG, depending on --capture_env_sensors_format

Sensors with None for a given data type are skipped. Only tasks that define scene sensors with image-like outputs (such as camera-based RL environments) produce captured frames.

Command-line options#

The training entrypoints register the capture flags in add_common_train_args:

    parser.add_argument(
        "--capture_env_sensors",
        type=int,
        default=0,
        help="Number of environment views to capture from each image-like scene sensor.",
    )
    parser.add_argument(
        "--capture_env_sensors_length",
        type=int,
        default=200,
        help="Length of each captured sensor frame window (in steps).",
    )
    parser.add_argument(
        "--capture_env_sensors_interval",
        type=int,
        default=2000,
        help="Interval between captured sensor frame windows (in steps).",
    )
    parser.add_argument(
        "--capture_env_sensors_format",
        choices=["tensorboard", "file"],
        default="tensorboard",
        help="Format used to save the captured sensor frames.",
    )

Capture schedule#

A frame is saved on every environment reset and after each step while the current per-episode step is inside an active capture window. The episode step counter resets to 0 on every environment reset. A step is inside the window when:

step % capture_env_sensors_interval < capture_env_sensors_length

For example, with --capture_env_sensors_length 100 and --capture_env_sensors_interval 2000, the wrapper saves episode steps 0-99 after each reset. If an episode lasts longer than 2000 steps, another capture window opens at episode steps 2000-2099, then 4000-4099, and so on.

Note

This schedule uses a per-episode step counter. Recording video clips during training instead keys off the Gymnasium RecordVideo wrapper’s global environment-step counter across episodes.

Episode indexing#

Each environment reset increments an episode index that starts at 1 on the first reset. Captured frames are grouped by this index so consecutive episodes can be compared:

TensorBoard tags use <sensor_name>/<data_type>/episode_<index> with a five-digit zero-padded episode index (for example, episode_00001)
File output uses episode_<index>_step_<step>.png under each sensor and data-type directory, with a five-digit episode index and an eight-digit step index (for example, episode_00001_step_00000042.png)

The <step> value is the per-episode step counter (reset to 0 on every environment reset). In TensorBoard, images are logged with global_step set to the total number of environment steps since training started.

Image processing#

Before tiling, each selected environment view is passed through normalize_camera_output_for_display(), which maps common camera data types to a [0, 1] float range suitable for PNG export:

RGB-like outputs are scaled by 255
Depth-like outputs (depth, distance_to_camera, distance_to_image_plane) are scaled by their per-frame maximum
albedo keeps the first three channels and scales by 255
normals are remapped from [-1, 1] to [0, 1]

make_camera_output_grid() arranges the capture_num_envs views into a roughly square grid (nrow = round(sqrt(num_envs))) before the image is written.

Output formats#

TensorBoard (default). Each sensor output is logged as a separate image series per episode with the tag <sensor_name>/<data_type>/episode_<index>. View the captures alongside other training metrics by pointing TensorBoard at the training log directory or the sensor_frames/train subdirectory:

./isaaclab.sh -p -m tensorboard.main --logdir logs/rsl_rl/Isaac-Reorient-Cube-Shadow-Camera-Direct/<run>/sensor_frames/train

File. PNG images are written per sensor, data type, episode, and step:

<sensor_name>/<data_type>/episode_<index>_step_<step>.png

Sensor and data-type names are sanitized for file output so they are safe as path components. For example, a frame from the first episode at step 42 for sensor front/camera and data type rgb is written as front_camera/rgb/episode_00001_step_00000042.png.

Summary#

`--capture_env_sensors_format`	Output location	How to view
`tensorboard` (default)	`sensor_frames/train` event files	TensorBoard image tab (tagged by sensor, data type, and episode index)
`file`	`sensor_frames/train/<sensor>/<data_type>/episode_<index>_step_<step>.png`	Any image viewer or filesystem browser