action_utils.py — 动作和图像处理的工具函数集合

文件路径: verl/experimental/vla/envs/action_utils.py 模块路径: verl.experimental.vla.envs.action_utils

文件概述

动作和图像处理的工具函数集合，为环境交互提供数据转换支持。包括动作预处理、图像操作和视频保存。

核心函数

1. prepare_actions - 动作预处理

def prepare_actions(simulator_type, raw_chunk_actions, num_action_chunks, action_dim):
    """将模型输出的原始动作转换为环境可接受的格式

    模型输出: (batch, num_chunks * action_dim) 的扁平向量
    环境需要: (batch, num_chunks, action_dim) 的结构化动作
    """
    actions = raw_chunk_actions.reshape(-1, num_action_chunks, action_dim)
    return actions

2. 图像处理函数

def resize_image(image: np.ndarray, size: tuple) -> np.ndarray:
    """将图像缩放到指定尺寸"""
    return cv2.resize(image, size)

def center_crop_image(image: Image) -> Image:
    """中心裁剪图像为正方形

    VLA 模型通常需要正方形输入图像。
    """
    width, height = image.size
    crop_size = min(width, height)
    left = (width - crop_size) // 2
    top = (height - crop_size) // 2
    return image.crop((left, top, left + crop_size, top + crop_size))

3. 可视化工具

def tile_images(images: list, grid_shape: tuple) -> np.ndarray:
    """将多张图像拼成网格，用于可视化"""
    ...

def put_text_on_image(image: np.ndarray, text: str) -> np.ndarray:
    """在图像上叠加文本标注"""
    ...

def save_rollout_video(frames: list, path: str, fps: int = 10):
    """将帧序列保存为视频文件

    用于记录机器人执行过程，便于调试和评估。
    """
    ...

4. 通用工具

def to_tensor(data, device="cpu"):
    """将 numpy 数组或列表转换为 PyTorch tensor"""
    if isinstance(data, np.ndarray):
        return torch.from_numpy(data).to(device)
    return torch.tensor(data).to(device)

核心函数列表

名称	说明
prepare_actions	动作格式转换
to_tensor	numpy -> tensor
tile_images	图像网格拼接
put_text_on_image	图像文本叠加
save_rollout_video	帧序列保存为视频
resize_image	图像缩放
center_crop_image	中心裁剪

与其他模块的关系

• 被 naive_rollout_rob.py 的 process_input 使用（图像预处理）
• 被 LiberoEnv 和 IsaacEnv 使用（动作格式转换）
• 被视频记录功能使用

小结

这是一个纯工具函数文件，提供了 VLA 系统中常用的数据转换操作。