robotic-transformer-pytorch

This project provides a PyTorch implementation of the RT-1 (Robotic Transformer) architecture. The model typically uses a MaxViT vision backbone to process video sequences and outputs action logits based on text instructions.

Training Mode

When calling the model during training, pass the video tensor and a list of instructions. The output shape is (batch, frames, actions, bins).

Inference with Classifier-Free Guidance

For evaluation/inference, you can use classifier-free guidance by setting the cond_scale parameter. This allows you to control the strength of the conditional instruction during generation.

import torch
from robotic_transformer_pytorch import MaxViT, RT1

# Initialize the vision backbone
vit = MaxViT(
    num_classes = 1000,
    dim_conv_stem = 64,
    dim = 96,
    dim_head = 32,
    depth = (2, 2, 5, 2),
    window_size = 7,
    mbconv_expansion_rate = 4,
    mbconv_shrinkage_rate = 0.25,
    dropout = 0.1
)

# Initialize the RT1 model
model = RT1(
    vit = vit,
    num_actions = 11,
    depth = 6,
    heads = 8,
    dim_head = 64,
    cond_drop_prob = 0.2
)

# Prepare input data
video = torch.randn(2, 3, 6, 224, 224)  # (batch, channels, frames, height, width)
instructions = [
    'bring me that apple sitting on the table',
    'please pass the butter'
]

# Training forward pass
train_logits = model(video, instructions) # Output shape: (2, 6, 11, 256)

# Inference with classifier-free guidance
model.eval()
eval_logits = model(video, instructions, cond_scale = 3.)