Introduction

This dataset contains the robotic manipulation data used in Magma pretraining. For fair comparison, we followed OpenVLA to use the data mix "siglip-224px+mx-oxe-magic-soup".

The dataset is organized by following source datasets, with each source containing one or more arrow files:

Features

In addition to the default features, we extracted the visual traces of future 16 frames for each frame. The dataset contains the following fields:

• dataset_name: Original source dataset name
• image: Image of the robot scene (binary)
• task_string: Description of the task
• frame_index: Index of the frame in the video
• traj_index: Index of the trajectory in the dataset
• action: Robot action vector (serialized numpy array)
• trace: Robot trajectory trace (serialized numpy array)
• trace_visibility: Visibility mask for the trace (serialized numpy array)

Dataset Loading

Full Dataset Load

from datasets import load_dataset
dataset = load_dataset("MagmaAI/Magma-OXE-ToM", streaming=True, split="train")

Individual Dataset Load

or specify a dataset by:

from datasets import load_dataset
dataset = load_dataset("MagmaAI/Magma-OXE-ToM", data_dir="austin_buds_dataset_converted_externally_to_rlds", streaming=True, split="train")

Sample Decoding

# Helper function to deserialize binary fields
def deserialize_array(bytes_data):
    return pickle.loads(bytes_data)

# Helper function to convert binary image data to PIL Image
def bytes_to_image(image_bytes):
    return Image.open(io.BytesIO(image_bytes))

for i, example in enumerate(dataset):   
    # decode the image: 256 x 256 x 3
    image = bytes_to_image(example['image'])
    # decode action: 1 x 7
    action = deserialize_array(example['action'])
    # decode trace: 1 x 17 x 256 x 2
    trace = deserialize_array(example['trace'])
    # decode trace visibility: 1 x 17 x 256 x 1
    trace_visibility = deserialize_array(example['trace_visibility'])