Action Map Policy: Pixel Classification for High-Precision Robot Manipulation

Problem: Curse of Dimensionality in Robot Action Spaces

Robot actions are high-dimensional (6–7 DoF per timestep). Discretizing each dimension into just 10 bins yields a joint action space of 10⁶ possibilities, and finer discretization or longer temporal horizons explode the space further. This mirrors the bitter lesson of the curse of dimensionality, making direct cross-entropy prediction over raw action tokens infeasible.

AMP's Core Idea: From 3D Regression to 2D Pixel Classification

AMP introduces a new action representation . It first expresses 3D motion (orientation + translation) as reversible 3D keypoint trajectories, then projects these keypoints onto multiple camera image planes. The 3D trajectory becomes a set of 2D pixel trajectories in image space. Instead of regression or diffusion, AMP predicts a per-pixel probability distribution (heatmap) for each keypoint at each future timestep, turning action prediction into a pixel-level classification problem trained with cross-entropy loss.

Key Advantages

Unified input/output space: Visual observations and action distributions both live in image space, analogous to token space in LLMs.

Native multi-modality: Cross-entropy naturally captures multiple plausible action modes (e.g., grasping different objects) without mode collapse.

Single forward pass: The full temporal action distribution is produced in one inference (~13.80 ms), unlike iterative diffusion policies (~93.53 ms for 16-step DDIM).

Architecture: X-Net Encoder–Decoder with Multi-View Transformer

The model uses an X-Net shaped encoder–decoder. Left branches encode one in-hand and two side-view images into latent feature maps, flattened to tokens. A Multi-View Transformer with in-image and cross-image attention layers fuses features across views. Right branches decode heatmaps at the original image resolution for each keypoint across the temporal horizon. Training is end-to-end in image space without explicit 3D trajectory supervision. At inference, argmax on heatmaps yields 2D keypoint trajectories, which are triangulated using camera intrinsics/extrinsics to recover 3D trajectories.

Experiment 1: Fine-Grained Perception & Multi-Modal Expression

Four visually identical blocks on a table; a laser pointer highlights one. Only 40 demonstrations (10 per block) with minimal spatial variation. Diffusion Policy suffers mode collapse, failing to condition on the laser signal. AMP achieves 100% success , demonstrating fine-grained visual perception and faithful multi-modal action distribution modeling.

Experiment 2: Real-World 3D Closed-Loop Breakfast Tasks

Three long-horizon tasks: Make Coffee (open lid, insert capsule, place cup, press button), Toast Bread (insert two slices, push lever), Steam Egg (place three eggs, close lid). Each requires high precision and stable closed-loop control over multiple stages.

Results (Table 1)

AMP outperforms Diffusion Policy (DiffPo) and ACT by 50–70% absolute success rate .

AMP maintains strong spatial generalization and precision under large spatial variation.

Inference latency: 13.80 ms (AMP) vs 93.53 ms (DiffPo, 16-step DDIM).

Additional Analyses

Precision vs. resolution: 3D reconstruction error scales linearly with image resolution; at 224×224, AMP achieves ~ 1 mm position error and 1.3° rotation error .

Extensive simulation experiments and ablation studies dissect action representation, architecture, and training design.

Compared to World Action Model (WAM) , AMP uses compact pixel-level keypoint heatmaps instead of high-dimensional image/video generation plus inverse dynamics, yielding lighter computation and faster inference.

Project page: https://haojhuang.github.io/amp_page/ | Paper:

https://arxiv.org/abs/2607.10706