\documentclass[10pt]{article}

\usepackage[margin=0.75in]{geometry}
\usepackage{times}
\usepackage{microtype}
\usepackage{graphicx}
\usepackage{booktabs}
\usepackage{multirow}
\usepackage{amsmath,amssymb}
\usepackage{hyperref}
\usepackage{xcolor}
\usepackage{caption}
\usepackage{tabularx}
\usepackage{subcaption}
\usepackage{wrapfig}


\title{\vspace{-0.2in}\textbf{Pixel-Aligned Robot Actions (PARA) are More Data-Efficient and Robust to New Viewpoints and Environments}}
\author{Anonymous Authors}
\date{}

\begin{document}
\maketitle
\vspace{-0.2in}

\begin{abstract}
Robot policies trained from images often fail under distribution shift: modest changes in object placement or camera viewpoint
can drastically degrade performance, even when the underlying task is unchanged.
We argue that a major contributor is \emph{action parameterization}: common formulations regress global robot coordinates
(joints or end-effector pose) from an image embedding, which couples the policy to nuisance factors (viewpoint, background, scene layout).
In contrast, modern image encoders are highly shift-equivariant and excel at learning \emph{pixel-aligned} correspondences.
We propose \textbf{KeyGrip}, a simple formulation that predicts robot actions as \emph{image feature selection}:
the policy outputs a heatmap over pixels indicating where the gripper should project in the image.
We then lift the selected 2D keypoint into robot space for inverse kinematics.
Our second contribution is a low-variance lifting parameterization: we predict \emph{table-aligned height} rather than depth or full 3D coordinates,
leading to improved camera-invariance and generalization.
On a real-robot benchmark with grasping and dense reaching trajectories, KeyGrip improves viewpoint-OOD and object-position-OOD success by
\textbf{2--4$\times$} over joint and end-effector regression baselines at equal data budgets, and consistently outperforms point-track regression.
\end{abstract}

\vspace{-0.05in}
\section{Introduction}
Robotic manipulation policies trained with behavioral cloning or RL commonly overfit to the training distribution.
In practice, this appears as brittle failures when (i) the object is shifted a few centimeters, or (ii) the camera is moved slightly.
While representation learning has improved image features dramatically, we observe that \emph{the action head often remains a bottleneck}:
predicting global robot coordinates from an image embedding encourages shortcut solutions and ties the policy to viewpoint-specific cues.

A key intuition is that actions like ``grasp the mug handle'' are fundamentally \emph{local} in the image.
Image encoders such as DINO-style transformers are trained to produce spatially meaningful, shift-equivariant features.
Yet many robot policies discard this structure by pooling features globally and regressing joint angles or Cartesian poses.

\textbf{KeyGrip} restores this alignment.
Instead of regressing coordinates in robot space, we predict a dense heatmap over pixels identifying where the gripper should land in the image.
This pixel-aligned formulation inherits the inductive bias of the encoder: translating the object in the image translates the heatmap.
We then lift the keypoint to 3D and solve IK / a short motion primitive.

\paragraph{Contributions.}
\begin{itemize}
    \item \textbf{Pixel-aligned action prediction:} formulate robot actions as selecting a gripper projection in the image via heatmaps,
    leveraging shift-equivariance for improved generalization.
    \item \textbf{Height-based lifting:} predict table-aligned height to lift the 2D keypoint to 3D, reducing variance vs. depth or full 3D regression.
    \item \textbf{Empirical study:} show strong gains on viewpoint and object-position generalization for grasping and dense trajectories, across data budgets.
\end{itemize}

\begin{figure}[t!]
\centering
\includegraphics[width=\linewidth]{figs/PARA_Overview_comparison_trainingdata.pdf}
\vspace{-0.05in}
\caption{\textbf{KeyGrip overview.} A fine-tuned DINOv2 encoder produces a pixel heatmap for the gripper projection and a scalar height.
The chosen pixel is lifted to 3D and executed via IK / a motion primitive.}
\label{fig:overview}
\vspace{-0.12in}
\end{figure}

\section{Related Work}
\paragraph{Vision-based policy learning.}
Behavior cloning from RGB images commonly predicts joint angles, end-effector deltas, or SE(3) poses.
These heads are simple but can be sensitive to camera changes and background cues, especially with limited data.

\paragraph{Pixel-aligned representations in 3D vision.}
A large body of work in 3D reconstruction and correspondence estimation uses dense, pixel-aligned outputs:
depth, optical flow, and point maps are predicted in image coordinates, then lifted geometrically.
KeyGrip adapts this principle to robot action prediction.

\paragraph{3D representations for robot generalization.}
Many approaches attempt to build viewpoint-invariant 3D state from multi-view reconstruction, explicit point clouds, or learned 3D features.
In practice, training robust 3D pipelines can be fragile and compute-heavy, and can introduce new failure modes (calibration, occlusions, scale drift).
KeyGrip instead leverages strong 2D encoders and a minimal lifting step.

\paragraph{Point tracks and coordinate regression.}
Point tracking is a useful intermediate representation, but most point-track policies regress coordinates (2D/3D) directly.
We compare to a track-regression baseline and show that \emph{classification over pixels} (heatmaps) is more robust under shift.

\section{Method}
\subsection{Problem Setup}
We assume a single RGB camera with known intrinsics and a fixed support surface (e.g., table) defining a world ``up'' axis.
From a demonstration dataset $\mathcal{D}=\{(I_t, a_t)\}$, we train a policy that outputs an action for grasping
or a short horizon trajectory.

\paragraph{Action as a gripper projection.}
For grasping, the target can be represented by a 2D keypoint $\hat{p}=(u,v)$ indicating the gripper center projected into the image,
plus an approach yaw $\hat{\psi}$ (optional) and a scalar height $\hat{h}$.
For dense reaching, we predict a sequence of keypoints $\{\hat{p}_k\}_{k=1}^K$ and heights $\{\hat{h}_k\}_{k=1}^K$.

\subsection{Pixel-Aligned Heatmap Head}
Given an image $I$, a vision encoder $f_\theta$ produces a spatial feature map $F \in \mathbb{R}^{H\times W\times C}$.
A lightweight head $g_\phi$ outputs a heatmap logits map $Z \in \mathbb{R}^{H\times W}$.
We train with a target heatmap $H$ constructed from the labeled pixel location $p^\star$:
\begin{equation}
\mathcal{L}_{\text{hm}} = \text{CE}\big(\text{softmax}(Z),\, H\big),
\end{equation}
where $H$ is a narrow Gaussian around $p^\star$ (normalized).
At inference, we take $\hat{p}=\arg\max_{u,v} Z_{u,v}$ or a soft-argmax.

\subsection{Lifting: Height vs Depth}
To execute in robot space, we lift the predicted pixel to 3D.
A common approach predicts depth $d$ along the camera ray and unprojects $P=\Pi^{-1}(p,d)$.
We instead predict \emph{height} $h$ along the world up axis (table-normal), then intersect the camera ray with the plane $z=h$:
\begin{equation}
\hat{P} = \Pi^{-1}\big(\hat{p}, \hat{h}\big).
\end{equation}
Intuitively, $h$ is less sensitive to camera pose and has lower variance for tabletop manipulation.

\begin{figure}[t!]
\centering
\includegraphics[width=0.98\linewidth]{figs/PARA_Method_Overview_ray.pdf}
\vspace{-0.05in}
\caption{\textbf{Height-based lifting.} Depth is view-dependent along the ray; height is table-aligned and encourages camera-invariance.}
\label{fig:height}
\vspace{-0.12in}
\end{figure}

\subsection{Execution}
We use inverse kinematics to obtain joint targets for $\hat{P}$ (and optional yaw), then execute a short, fixed-duration primitive:
approach $\rightarrow$ close gripper $\rightarrow$ lift / retract.
For dense reaching, we execute a sequence of IK waypoints.

\section{Experiments}
\subsection{Tasks and Evaluation}
We evaluate two tasks:
\textbf{(i) Grasping:} pick a single target object.
\textbf{(ii) Dense reaching trajectory:} move the end-effector along a demonstrated path (e.g., push, slide, or reach-to-contact).

We test two distribution shifts:
\textbf{Object-position OOD:} object translated/rotated on the table outside training placements at a fixed camera view.
\textbf{Viewpoint OOD:} camera moved to unseen viewpoints while object placement remains in-distribution.

Metrics:
\textbf{Grasp success} (\%) and \textbf{trajectory error} (cm) averaged over 50 trials per condition.



\subsection{Baselines}
\begin{itemize}
    \item \textbf{Joint regression:} global pooled features $\rightarrow$ joint angles.
    \item \textbf{EEF regression:} pooled features $\rightarrow$ end-effector pose delta.
    \item \textbf{Point-track regression:} predict a 2D point coordinate (or tracked point) directly with an $L_2$ loss.
    \item \textbf{KeyGrip variants:} heatmap + height (ours), heatmap + depth, heatmap + full 3D (x,y,z).
\end{itemize}



\vspace{4pt}
\footnotesize
\textbf{Notes:} Train Loc. = in-distribution (training workspace + camera range).
New Env. = novel backgrounds/lighting/layout.
New View = novel camera viewpoints. Replace placeholder numbers with your results.



\subsection{Implementation Details}
Backbone: DINOv2 ViT (pretrained).
We fine-tune the backbone and train a lightweight conv head for the heatmap and an MLP for height.
Input resolution: $224\times224$.
Data: teleoperation demonstrations (robot) with optional auxiliary human-hand images for co-training.
Optimization: AdamW, 50 epochs, standard augmentation (crop, brightness, small rotations).

\section{Results}
\subsection{OOD Generalization vs Data Diversity}
Figures~\ref{fig:curves_view}--\ref{fig:curves_obj} show grasp success as we increase the number of distinct training viewpoints / object positions.
KeyGrip scales much better with limited diversity: even with 1--2 training views, it maintains usable OOD performance.

%\begin{wrapfigure}{r}{0.52\textwidth}
%    \centering
%    \includegraphics[width=0.48\textwidth]{figs/keygrip_intuition2.pdf}
%    \caption{\textbf{Viewpoint-OOD grasp success} vs training viewpoint diversity (fake/hopeful data).}
%    \label{fig:curves_view}
%\end{wrapfigure}




\begin{figure}[t]
\centering
\includegraphics[width=0.99\linewidth]{figs/PARA_Results.pdf}
\vspace{-0.05in}
\caption{\textbf{Generalization axes.} We vary object position and camera viewpoint separately to measure OOD robustness.}
\label{fig:ood_axes}
\vspace{-0.12in}

\end{figure}

\begin{figure}[t]
\centering
\includegraphics[width=0.99\linewidth]{figs/PARA_PointTracksInference.pdf}
\vspace{-0.05in}
\caption{\textbf{Example Inference} We vary object position and camera viewpoint separately to measure OOD robustness.}
\label{fig:ood_axes}
\vspace{-0.12in}
\end{figure}

\begin{table}[t!]
\centering
\small
\setlength{\tabcolsep}{6pt}
\renewcommand{\arraystretch}{1.15}
\caption{\textbf{Task completion score (\%) across tasks and generalization regimes.}
Each entry reports the mean over evaluation episodes (averaged over $S$ seeds).}
\label{tab:results_ppc_multi}
\begin{tabularx}{\linewidth}{l *{6}{>{\centering\arraybackslash}X}}
\toprule
\textbf{Method}
& \textbf{Pick\&Place Cup (PPC)}
& \textbf{PPC New Viewpoint}
& \textbf{PPC New Viewpoint F.T.}
& \textbf{PPC New Environment}
& \textbf{Wipe Table}
& \textbf{Fold Towel} \\
\midrule
\textbf{PARA (Ours)}           & -- & -- & -- & -- & -- & -- \\
PARA Depth                    & -- & -- & -- & -- & -- & -- \\
ACT                           & -- & -- & -- & -- & -- & -- \\
Motion Tracks Height          & -- & -- & -- & -- & -- & -- \\
Motion Tracks Depth           & -- & -- & -- & -- & -- & -- \\
CLIPort Height                & -- & -- & -- & -- & -- & -- \\
CLIPort Depth                 & -- & -- & -- & -- & -- & -- \\
\bottomrule
\end{tabularx}
\end{table}


\subsection{Main Table: Grasp + Trajectory}
Table~\ref{tab:main} summarizes performance at a fixed data budget (200 demos, 4 train views / 4 train positions).
KeyGrip improves grasp success substantially and reduces trajectory error.

% Requires: \usepackage{booktabs,tabularx,makecell}
% Optional (nice): \usepackage{siunitx}  % for aligned numbers


\subsection{Ablation: Height vs Depth vs Full 3D}
Table~\ref{tab:lift} isolates the lifting choice under identical heatmap prediction.
Height reduces error and improves robustness, consistent with lower-variance supervision.





\subsection{Optional: Co-training with Human-Hand Data}
We explore augmenting robot demonstrations with cheap human-hand data that labels pixel keypoints (no robot calibration needed).
Figure~\ref{fig:cotrain} shows consistent gains, suggesting KeyGrip benefits from additional pixel-aligned supervision.





\section{Discussion}
Why does KeyGrip help?
(1) \textbf{Locality:} predicting a pixel avoids compressing the scene into a single vector.
(2) \textbf{Equivariance:} shifting the object shifts the heatmap, reducing reliance on background shortcuts.
(3) \textbf{Easier supervision:} heatmaps turn regression into classification, often improving calibration and data-efficiency.
(4) \textbf{Low-variance lifting:} height is closer to task-relevant geometry for tabletop manipulation.

Limitations include dependence on a visible target region and a known support plane; future work can relax these via learned plane estimation or multi-view cues.



\section{Conclusion}
We presented KeyGrip, a simple pixel-aligned action formulation that significantly improves generalization under object-position and viewpoint shifts.
By predicting where the gripper should project in the image and lifting via a height parameterization, KeyGrip leverages strong image features
and avoids brittle global coordinate regression. We hope this encourages broader use of pixel-aligned supervision for robust robot learning.

\begin{wrapfigure}{r}{0.55\linewidth}
    \centering
    \vspace{-10pt}
    \includegraphics[width=\linewidth]{figs/height_illustration.pdf}
    \vspace{-0.05in}
    \caption{\textbf{Height vs Depth Prediction Along Ray.}}
    \label{fig:ood_axes}
    \vspace{-10pt}
\end{wrapfigure}


\vspace{0.1in}
{\small
\paragraph{Acknowledgements.} Placeholder.
}

\end{document}
