Raising Body Ownership in End-to-End Visuomotor Policy Learning via Robot-Centric Pooling

In plain words

A robot policy that looks at a camera image to decide its next move can get confused when other objects, or a second robot, sit in the frame. It may track the wrong arm, or react to something that only looks like its own gripper. This paper asks whether the policy can learn to tell "that is me" apart from "that is something else" just by looking at the training data it already has.

The method changes one step in the image pipeline. Instead of averaging image features over the whole picture, the model asks which part of the picture matches the robot's current joint and gripper state, and it keeps only the features from that part. To teach this matching without labels, the training data is augmented: the robot's arm is pasted onto different backgrounds, and a copy of another robot or object is sometimes added. Two views of the same robot pose should give similar answers, and views of different poses should give different ones. The matching step is trained together with the normal action-prediction loss, so no extra data has to be collected.

On reaching tasks in simulation and on a real robot, the authors report that policies built this way keep working when distractors appear, including a second robot of the same type. They also report that the policies are less affected when the camera image is shifted sideways.

TL;DR

• Robot-centric Pooling (RcP) replaces the final spatial pooling of a ResNet18 policy with a proprioception-guided cross-attention and a hard argmax mask §3.1.
• It is trained with a MoCo-style InfoNCE loss on augmented images, jointly with an L1 policy loss, using lambda = 0.05 §3.2 §3.3 §4.1.
• Author-reported simulated reaching success under unseen distractors: RcP 76.4%, Spatial-Softmax 31.9%, R3M 8.3% tab. 1 §4.2.1.
• Author-reported real-robot reaching success with a self-distractor: RcP 92%, Spatial-Softmax 14%, R3M 2% §4.3.
• No official code was found. Re-implementation requires resolving several unstated details listed under Pitfalls.

Problem

Standard end-to-end visuomotor policies fuse image and proprioception by concatenating features after a global or spatial pooling of the image. They do not learn to separate the robot from its surroundings or from other robots, so distractors shift their actions §1 §2. The paper argues this is a missing "body ownership" ability at the visuomotor level and measures it through reaching success under distractors §1 §4.2.1. The term is used as an analogy for the robot-level behaviour measured in the experiments, not as a claim about human cognition §1.

The method assumes a fixed second-person camera, a calibrated-enough view of the manipulator for augmentation, and a dataset with proprioception and actions for each frame §3.2 §4.2.

Method

The model has four parts, all trained on the same image-proprioception-action dataset §3:

1. RcP image encoder. A ResNet18 with the last average-pooling and classification layers removed produces a spatial feature map. After LayerNorm, two learned linear projections give keys k and values v, one per spatial location §3.1 eq. 1–3.
2. Proprioception query. The proprioceptive vector p is expanded with sine and cosine features at L frequency levels, then passed through a two-layer perceptron Q to give a 1 x d query q §3.1 eq. 4 eq. 5. A decoder with the mirrored architecture reconstructs the query, which encourages a meaningful latent code §3.1.
3. Alignment and mask. The alignment score of each location is the softmax of q against the keys §3.1 eq. 6. The context vector h is the score-weighted sum of the values §3.1 eq. 7. For the policy path, an argmax over the scores gives a one-hot mask that selects one location §3.1 eq. 8.
4. Policy path. A second ResNet18, with two extra 3x3 convolutions to enlarge the receptive field, gives a feature map. It is multiplied element-wise by the one-hot mask to give the pooled image feature g §3.3 eq. 10. g is concatenated with p and passed to a two-layer perceptron that predicts the action with an L1 loss §3.3 eq. 11.

The contrastive part trains the RcP encoder so that context vectors for the same robot pose agree across augmentations and differ across poses §3.2:

• Augmentation. The manipulator region is cropped with a roughly calibrated camera and pasted onto two random backgrounds at random positions. A self-distractor is cropped from another training image and pasted onto one of the two copies, leaving most of the robot visible. Random scaling and colour jitter are applied §3.2 fig. 3.
• Two encoders. The RcP image encoder is copied. The online copy gets the gradient. The momentum copy has weights updated as a moving average, and the proprioception encoder and decoder are not copied §3.2.
• Contrastive objective. Both context vectors pass through one shared linear map W, giving z and z+. The loss is InfoNCE against a queue of past projected negatives, with temperature tau §3.2 eq. 9.
• Joint loss. The total loss is the policy loss plus lambda times the sum of the contrastive loss and the reconstruction loss §3.3 eq. 12. The method can also be pre-trained alone §3.3.

Pseudocode: RcP forward pass

# Robot-centric Pooling forward pass  [§3.1] [eq. 1-8]
# Input:  image (H, W, 3); proprioception vector (P,)
# Output: one-hot mask (1, h*w); context vector (1, d)
def robot_centric_pooling(image, proprio):
    # Backbone features, layer-normalised, then keys and values
    # [§3.1] [eq. 1-3]
    features = layer_norm(resnet18(image))         # (h, w, c)
    keys = project_k(flatten(features))            # (h*w, d)
    values = project_v(flatten(features))          # (h*w, d)
    # Frequency encoding, then a two-layer perceptron Q  [§3.1] [eq. 4-5]
    encoded = frequency_encode(proprio)            # L levels: not stated
    query = mlp_q(encoded)                         # (1, d)
    # Alignment score per location: softmax of query against keys
    # [§3.1] [eq. 6]
    scores = softmax(query @ transpose(keys))      # (1, h*w)
    # Context vector: score-weighted sum of values  [§3.1] [eq. 7]
    context = scores @ values                      # (1, d)
    # One-hot mask at the highest-scoring location  [§3.1] [eq. 8]
    mask = one_hot(argmax(scores), h * w)          # (1, h*w)
    return mask, context

Pseudocode: pooled policy feature and loss

# Policy feature and total loss  [§3.3] [eq. 10-12]
# Input:  mask (1, h*w); conv_features (h, w, d); action; proprio; L_moco
# Output: scalar loss
def policy_loss(mask, conv_features, action, proprio, l_moco, recon):
    # Element-wise product of the mask with the feature map  [§3.3] [eq. 10]
    masked = reshape(mask, (h, w, 1)) * conv_features   # (h, w, d)
    # Reduction over space to a d-vector is not written in the paper;
    # summing over locations is the only reduction that gives g in R^d.
    pooled = sum_over_space(masked)                     # (d,)  inferred
    # L1 regression on the concatenated feature and state  [§3.3] [eq. 11]
    prediction = mlp_policy(concat(pooled, proprio))
    l_policy = l1(prediction, action)
    # Joint objective with preferential weight lambda  [§3.3] [eq. 12]
    return l_policy + lam * (l_moco + recon)            # lam = 0.05 [§4.1]

Pseudocode: one contrastive training step

# One training step: contrastive pair, momentum update, queue  [§3.2] [eq. 9]
# Input:  images, proprio, queue of K projected negatives
# Output: contrastive loss; updated momentum weights and queue
def contrastive_step(images, proprio, queue):
    # Two augmented copies of each image: arm pasted on backgrounds  [§3.2]
    anchor, positive = augment_pair(images)
    # Online encoder on the anchor; momentum encoder on the positive  [§3.2]
    _, h_anchor = robot_centric_pooling(anchor, proprio)
    _, h_positive = robot_centric_pooling_momentum(positive, proprio)
    # Shared linear projection W  [§3.2]
    z = linear_w(h_anchor)
    z_positive = linear_w(h_positive)
    # InfoNCE against the queue, temperature tau = 0.1  [§4.1] [eq. 9]
    loss = info_nce(z, z_positive, queue, tau=0.1)
    # Momentum update: theta_bar <- m * theta_bar + (1 - m) * theta  [§3.2]
    momentum_weights = m * momentum_weights + (1 - m) * weights  # m = 0.95
    # Enqueue the positive projection; queue length 4096  [§3.2] [§4.1]
    queue = enqueue(queue, z_positive)
    return loss, queue

Hyperparameters

Results

Author-reported. Simulation is a UR5 multi-instance reaching task with a second-person camera. Real experiments use a Franka Panda with a side second-person view. Baselines are ResNet18 with Spatial-Softmax (ImageNet pre-trained) and ResNet18 with average pooling and R3M pre-training §4.1.

Simulation, distractor success rate (%), mean ± std over 3 seeds x 50 trajectories tab. 1 §4.2.1

The best SSM case is 62.7% against the Franka at back-left, and SSM is at 0% for the Franka and Plant in front tab. 1 §4.2.1. RcP is lowest for the self-distractor at back-left (62.7%), which the authors attribute to the 7 x 7 final feature map resolution §4.2.1.

Real robot reaching, success rate (%) §4.3

Real results are single numbers. No variance is reported in the text.

Translation error after the self-distractor fig. 5a §4.3. The authors report RcP deviations mostly within about 1 cm, and wider and larger errors for SSM and R3M. The exact values are only in the figure.

Rolled-shift robustness fig. 5b §4.4. RcP degrades gradually over shifts from -10% to +35%, and outperforms the baselines in both simulation and real experiments. R3M is similar in simulation but is sensitive to the sim-to-real gap. SSM peaks sharply at the 7% shift used in training §4.4. The text gives no numbers for this curve; read them from the figure.

Paper vs code

No official code.

The paper's LaTeX source contains one link, in a comment in the abstract file: https://github.com/Zheyu-Zhuang/RcP 00_abstract.tex. It is not in the compiled paper. The repository does not exist: cloning it returned "Repository not found". The first author's GitHub account (lower-case zheyu-zhuang) has seven public repositories, none of them for RcP. Searches run: the repository link above; GitHub repository search for "robot-centric pooling visuomotor" and "RcP" within that user; no project page is linked. The arXiv metadata contains no comment field with a link.

Because no code was found, the points below cannot be checked against an implementation. They are gaps, not discrepancies.

Pitfalls

• Cosine similarity claim. The text says the alignment scores are cosine similarities §3.1, but the equation is a softmax of the raw dot product q k^T with no normalisation or temperature §3.1 eq. 6. Follow the equation unless the authors say otherwise; the two differ in scale and in how query norms affect the scores.
• No gradient through the argmax. The policy path uses a hard one-hot mask eq. 8. The equations do not show a gradient path from the policy loss to the query, keys or values. The contrastive loss uses the soft context h eq. 7, not the mask, so the query and keys seem to learn from the contrastive and reconstruction terms only. This is inferred from the equations and is not stated in the paper.
• Reduction to g in R^d. The equation for g multiplies the one-hot mask with a spatial feature map and gives g in R^d, but does not write the spatial reduction §3.3 eq. 10. Summing over space is the only reading that gives a d-vector. This is inferred.
• Reconstruction target. L_recon is a mean squared error that reconstructs the encoded proprioception query §3.3, but the decoder input and the target are not stated.
• Queue contents. The text says the momentum context h+ updates the queue §3.2, while eq. 9 uses projected vectors z-. Whether the queue stores h+ or W h+ is not stated.
• Queue initialisation. How the queue is filled before the first step is not stated.
• Warm-up shift. The first 5 epochs use "slightly shifted" images §4.1. Their size and distribution are not given.
• Arm crop. Augmentation crops the manipulator "using a roughly calibrated camera" §3.2. How the crop mask is computed (projection of joint states, segmentation, or manual) is not stated. Without it, the augmentation cannot be reproduced.
• Self-distractor placement. The paste location must keep most of the robot visible §3.2, but the visibility threshold is not stated.
• Random scaling and jitter. Scaling range and jitter strength are not stated; jitter is cited to robomimic §3.2.
• Mask resolution. The authors say the 7 x 7 final feature map may limit separating overlapping robot and distractor §4.2.1. The receptive-field argument for the extra 3x3 convolutions depends on this map §3.3.
• Dimension d. d is used for keys, values, query, the projection W and the policy feature, but never given a value §3.1 §3.2.
• Frequency levels. L is used in the encoding but never set §3.1.
• Two copies of ResNet18. The policy encoder and the RcP encoder are both ImageNet-initialised ResNet18s §4.1. Whether they share weights is not stated; the equations use separate parameters psi and phi §3.1 §3.3.
• Robot-centric claim. The paper labels the mechanism "body ownership". The experiments measure reaching success under distractors and pixel shifts. Read the terminology as a name for those measurements §1.
• Ablations. The paper has no ablation over lambda, the argmax mask, the queue size, the warm-up, or the momentum. Only the three-way comparison with baselines is reported §4.
• Baseline pre-training. The two baselines differ in pooling and pre-training together (SSM: ImageNet, R3M: Ego4D), so the comparison does not isolate the pooling change §4.1.
• Real-world sample size. The real-world experiment has one attempt per network, with and without the distractor, for each of 50 target locations §4.3. No confidence intervals are reported for those numbers.
• Simulation sample size. Three seeds, each with 50 trajectories, per distractor and location §4.2.1. Seed selection for the real-world experiment uses simulated reaching performance §4.3.
• Real-world fine-tuning. The paper says the training protocol matches the simulation protocol §4.3. Fine-tuning schedule and epochs are not separately stated.
• Pixel shift direction. Test-time shifts are horizontal rolls that wrap around the image §4.4. Vertical shifts are not tested.

Lineage

• Builds on: MoCo, the momentum-encoder contrastive framework (cited as he2020_moco) §2 §3.2; ResNet18 §3.1; InfoNCE §3.2; and the cross-attention formulation of the transformer (cited as vaswani2017attention) §3.1. None of these has a card in this wiki.
• Compared against: Spatial-Softmax, the pooling used in robomimic and diffusion policy §4.1; R3M, pre-trained on Ego4D §4.1.
• Related: Spatial-Softmax was introduced by finn2016_spatial_encoder_vs §4.1. The FullGrad saliency tool used for visualisation is cited as srinivas2019fullgrad fig. 1.
• Superseded by: not stated in this wiki.