Pixel-UMI: Learning Environment-Aware Whole-Body Manipulation from Sensorized Human Demonstrations

Pixel-UMI: Environment-Aware Whole-Body Manipulation from Human Demonstrations§

Problem & Motivation§

Sensorized human demonstrations (head-mounted cameras + gripper-tracked hands) are a scalable data source for mobile manipulation, but transferring them to a robot suffers from embodiment mismatch: different kinematics produce different reachable poses, occlusions, and viewpoints. Pixel-UMI is a three-stage pipeline that addresses each mismatch source in turn: (1) pose mismatch via scene-aware retargeting, (2) viewpoint mismatch via error-aware masking, and (3) contact/dynamics mismatch via on-policy human corrections.

1. Environment-Aware Retargeting (EAR)§

Given a reconstructed scene geometry (mesh/point cloud) $\mathcal{S}$ and demonstrated head pose $T^{h}_t$ and gripper poses $T^{g}_t$, EAR solves for robot base and joint configuration $q_t$ that best match the demonstrated end-effector targets while respecting collision and kinematic limits:

$$q_t = \arg\min_{q \in \mathcal{Q}_{free}} \; \sum_{k} \| f_k(q) - T^{g,k}_t \|^2 + \lambda \, \mathcal{C}(q, \mathcal{S})$$

where $f_k(\cdot)$ is the robot forward kinematics for gripper $k$, and $\mathcal{C}(q,\mathcal{S})$ is a collision penalty against the reconstructed scene (so the robot avoids whole-body collisions when the human leaned into constrained spaces). Collision-aware base placement is derived from the head trajectory's reachability cone; a delay-calibrated action assignment aligns asynchronous human/robot timesteps by estimating the human's motor latency and shifting demonstration timestamps before pairing observations with actions.

2. Retargeting-Error-Aware Masking (REAM)§

Residual pose errors $e_t = \| f(q_t) - T^{g}_t\|$ remain after IK-style retargeting, so robot observations differ from the human's viewpoints at the same state. REAM computes a per-pixel / per-token reliability mask and masks mismatched views during policy training so the network does not overfit to views it cannot reproduce at deployment:

$$\mathcal{L}_{\text{pol}} = -\mathbb{E}\left[ \sum_i m_i \, \log \pi(a_t \mid o_t \odot m, q_t) \right], \quad m_i = \mathbb{1}[\, \text{view-alignment}(i) > \tau \,]$$

This is conceptually similar to masked imitation learning but the mask is derived from geometric retargeting error rather than occlusion heuristics.

3. DAgger On-Policy Corrections§

Estimated retargeted actions cannot capture contact-rich dynamics (deformable objects, tight bimanual coordination). Pixel-UMI runs the learned policy and collects human corrections on the robot under a DAgger-style aggregation to correct covariate shift:

$$\pi_{n+1} = \arg\min_\pi \; \mathbb{E}_{s \sim d_{\pi_n}}\!\left[ \ell(\pi(s), a^{\text{human}}_{\text{corrected}}(s)) \right]$$

This closes the loop between human-scale demonstrations and robot-scale dynamics.

Implementation Sketch§

import torch, torch.nn as nn

class Policy(nn.Module):
    def __init__(self, obs_dim, act_dim):
        super().__init__()
        self.vis = nn.Sequential(nn.Conv2d(3,32,5,2), nn.ReLU(), nn.Conv2d(32,64,5,2), nn.ReLU(), nn.Flatten())
        self.fuse = nn.Sequential(nn.Linear(64*7*7 + obs_dim, 512), nn.ReLU(), nn.Linear(512, act_dim))
    def forward(self, img, proprio, mask):
        feat = self.vis(img * mask)          # REAM: mask mismatched views
        return self.fuse(torch.cat([feat, proprio], -1))

def ear_ik_reachability(q_fn, scene, Tg_targets, lam=1.0):
    # gradient-based IK with collision penalty (EAR)
    q = torch.zeros(q_fn.dof, requires_grad=True)
    opt = torch.optim.Adam([q], lr=1e-2)
    for _ in range(200):
        loss = sum(((q_fn(q,k) - Tg_targets[k])**2).sum() for k in Tg_targets) \
               + lam * collision_penalty(q, scene)
        opt.zero_grad(); loss.backward(); opt.step()
    return q.detach()

# training step (masked IL + DAgger replay)
def train_step(policy, batch):
    img, proprio, mask, actions = batch
    pred = policy(img, proprio, mask)
    return nn.functional.mse_loss(pred, actions)

Results§

Evaluated on five tasks — articulated-object manipulation, deformable handling, constrained-space reaching, whole-body collision avoidance, and contact-rich bimanual coordination — averaging 87.0% task progress, demonstrating that scene-aware retargeting plus error-aware masking and targeted human corrections can span the human→robot embodiment gap for mobile whole-body manipulation.