Axis

Completed in 2026 at the age of 19

PyTorch DINOv2 SE(3) Lie Groups Transformers Python
View Repository

Overview

Axis is a PyTorch transformer-based robot learning framework designed for end-effector control using SE(3) Lie group representations. By treating robot control as a sequence modeling problem, Axis learns from diverse datasets like DROID and Fractal via imitation learning to predict SE(3) twist actions.

Interactive Architecture Diagram

Hover over the components to explore the inner workings and tensor shapes of the Axis V2 model.

Vision

Image Window

Vision Input & Encoder

Shape: (B, W, 3, 224, 224)

Encoder: DINOv2 ViT-S/14

Encodes multi-camera visual input into dense 768D representations per frame.

Proprioception

13D SE(3) State

Proprioceptive State

Shape: (B, W, 13)

Contains flattened rotation matrix (9D), translation (3D), and gripper state (1D).

Mapped to 128D via an MLP.

Goal

Semantic Embedding

Goal Embedding

Shape: (B, 38)

Structured semantic goal vector including 13D start and end poses. Mapped to 128D via an MLP.

Token Construction

Per-Frame Concatenation

Tokenization

Shape: (B, W, 1024)

Concatenates Proprio (128D) + Vision (768D) + Goal (128D) into a single 1024D token per timestep.

Axis Transformer

8 Layers, 16 Heads, RoPE

Transformer Backbone

Shape: (B, W, 1024)

Uses Rotary Position Embeddings (RoPE) for improved sequence modeling across the temporal window `W`.

Last Token Extraction

Causal Bottleneck

Bottleneck

Shape: (1, 1024)

Extracts only the final output token output_tokens[:, -1, :] for future prediction.

Action Decoder

Future Trajectory

SE(3) Twist Actions

Shape: (B, ChunkSize, 7)

Predicts a sequence of Lie algebra twists [ω, v, gripper] using an Action Chunking approach.

Confidence Decoder

Reliability Score

Self-Supervised Confidence

Shape: (B, 1)

A simple MLP prediction trained to estimate the model's own accuracy (action loss) for the current step.

Key Features

  • SE(3) Twists: Actions are geometrically consistent, represented as 7D Lie algebra twists [ω, v, gripper] via PyPose.
  • Action Chunking: Predicts W future actions in parallel, bypassing the need for autoregressive rollout.
  • Endpoint Chordal Loss: Utilizes a trig-free Frobenius norm for task-oriented SE(3) loss optimization.
  • Temporal Ensembling: Overlapping chunk averaging yields smooth and robust action output.
  • Self-Supervised Confidence: The model actively learns to predict its own accuracy to provide a confidence signal.
Back to Projects