Implementation:ARISE Initiative Robomimic VisUtils

Knowledge Sources	robomimic
Domains	Robotics, Visualization, Debugging
Last Updated	2026-02-15 08:00 GMT

Overview

Concrete tool for converting, saving, and visualizing image observations and depth maps during training and debugging provided by the robomimic utilities module.

Description

The vis_utils module provides five utility functions for image observation debugging. image_tensor_to_numpy converts processed image tensors (CHW, normalized) back to NumPy arrays (HWC, uint8) via the observation unprocessing pipeline. image_to_disk and image_tensor_to_disk save images to files. visualize_image_randomizer creates a matplotlib grid comparing original images with their N randomized variants for visual inspection of data augmentation. depth_to_rgb applies matplotlib's "hot" colormap to depth maps for false-color visualization.

Usage

Use these functions during model development and debugging to visually inspect how observations are processed, how data augmentation randomizers transform images, and how depth sensors render. They are not intended for production training loops.

Code Reference

Source Location

• Repository: robomimic
• File: robomimic/utils/vis_utils.py
• Lines: 1-113

Signature

def image_tensor_to_numpy(image):
    """
    Converts processed image tensors to numpy for saving to disk or video.

    Args:
        image (torch.Tensor): images of shape [..., C, H, W]

    Returns:
        image (np.array): converted images of shape [..., H, W, C] and type uint8
    """


def image_to_disk(image, fname):
    """
    Writes an image to disk.

    Args:
        image (np.array): image of shape [H, W, 3]
        fname (str): path to save image to
    """


def image_tensor_to_disk(image, fname):
    """
    Writes an image tensor to disk, indexing out leading batch dimensions.

    Args:
        image (torch.Tensor): image of shape [..., C, H, W]
        fname (str): path to save image to
    """


def visualize_image_randomizer(original_image, randomized_image, randomizer_name=None):
    """
    Visualizes before/after of an image-based input randomizer.

    Args:
        original_image: batch of original images shaped [B, H, W, 3]
        randomized_image: randomized images shaped [B, N, H, W, 3]
        randomizer_name (str, optional): name of the randomizer for the title
    """


def depth_to_rgb(depth_map, depth_min=None, depth_max=None):
    """
    Convert depth map to RGB using matplotlib hot colormap.

    Args:
        depth_map (np.array): depth values of shape [H, W] or [H, W, 1]
        depth_min (float, optional): minimum depth for normalization
        depth_max (float, optional): maximum depth for normalization

    Returns:
        rgb (np.array): false-color RGB image of shape [H, W, 3], uint8
    """

Import

import robomimic.utils.vis_utils as VisUtils

# Or import individual functions:
from robomimic.utils.vis_utils import image_tensor_to_numpy, visualize_image_randomizer, depth_to_rgb

I/O Contract

Inputs

Name	Type	Required	Description
image	torch.Tensor	Yes	Processed image tensor of shape [..., C, H, W] (for tensor functions)
image	np.array	Yes	NumPy image of shape [H, W, 3] (for image_to_disk)
fname	str	Yes	File path for saving images to disk
original_image	np.array	Yes	Batch of original images [B, H, W, 3] (for visualize_image_randomizer)
randomized_image	np.array	Yes	Randomized images [B, N, H, W, 3] (for visualize_image_randomizer)
depth_map	np.array	Yes	Depth map of shape [H, W] or [H, W, 1] (for depth_to_rgb)
depth_min	float	No	Minimum depth for normalization (default: auto from data)
depth_max	float	No	Maximum depth for normalization (default: auto from data)

Outputs

Name	Type	Description
image_tensor_to_numpy returns	np.array (uint8)	Unprocessed image of shape [..., H, W, C] in [0, 255]
image_to_disk returns	None	Saves image file to fname
image_tensor_to_disk returns	None	Saves first element of batch to fname
visualize_image_randomizer returns	None	Displays matplotlib grid comparing original and randomized images
depth_to_rgb returns	np.array (uint8)	False-color RGB image [H, W, 3] using hot colormap

Usage Examples

Inspecting Processed Images

import robomimic.utils.vis_utils as VisUtils

# During training, inspect what the model sees
# image_tensor is a processed tensor from the dataloader [B, C, H, W]
numpy_images = VisUtils.image_tensor_to_numpy(image_tensor)

# Save first image to disk for inspection
VisUtils.image_tensor_to_disk(image_tensor, "debug_observation.png")

Visualizing Data Augmentation

import numpy as np
from robomimic.utils.vis_utils import visualize_image_randomizer

# original_images: [B, H, W, 3] batch of images
# randomized_images: [B, N, H, W, 3] batch with N augmented copies each
visualize_image_randomizer(
    original_image=original_images,
    randomized_image=randomized_images,
    randomizer_name="CropRandomizer",
)

Depth Map Visualization

import numpy as np
from robomimic.utils.vis_utils import depth_to_rgb

# depth_map: [H, W] array of depth values
rgb_depth = depth_to_rgb(depth_map, depth_min=0.0, depth_max=5.0)
# rgb_depth is now a [H, W, 3] uint8 image with hot colormap

Related Pages

Implements Principle

• Principle:ARISE_Initiative_Robomimic_Observation_Visualization_Debugging