Scaling Manipulation Learning with Visual Kinematic Chain Prediction

There are numerous techniques in machine learning and computer vision that can successfully learn a single general-purpose model from multiple diverse datasets [1]. In robotics, however, despite the recent advances in multi-task learning that enable a single policy to perform various tasks through imitation learning with language instructions [2, 3], these methods are typically constrained to a single robot and workspace. Some recent work, such as RT-X [4, 5], leverages a large vision-language model (VLM) to directly train policies with Behavioral Cloning (BC) on the Open-X Embodiment [5], a collection of datasets crowd-sourced from various environments and robots. However, these techniques require a non-trivial action normalization procedure to bridge the gap between the different action spaces in diverse setups, such as end-effector poses, joint positions, and velocities in various world frames. This manual engineering procedure is currently custom-designed for each training dataset, which affects the generalization and interpretability of these models.

Therefore, a key question is: can we find an action representation that is precise and universal for various setups and robots? To this end, we propose the visual kinematic chain, which is the projection of the robot’s high-dimensional kinematic structure into the image plane as pixels. Instead of predicting the low-level robot actions, we propose to learn and visually forecast the movement of robots’ kinematic chains in the image plane. This visual approach provides a unified action representation for different robots, and requires no manual adjustment since the visual kinematic chains can be automatically obtained from the robot’s model and camera parameters.

To forecast the kinematic structures of various robots, we render the kinematic chains into point sets by sampling points along links and performing optimal matching [6] to minimize the earth moving distance [7, 8] between the predicted point sets and ground-truth kinematic structures. Further, we propose the Visual Kinematics Transformer (VKT), a convolution-free architecture that is solely based on the attention mechanism in the RGB space, which supports an arbitrary number of camera viewpoints. Our VKT is trained with a single objective of forecasting kinematic structures through optimal point set matching without seeing any low-level robot actions. VKT is deployed to a specific environment by simply training a tiny head while freezing the backbone. VKT demonstrates superior performance over BC transformers as a general agent on Calvin (success length 1.08 vs 0.33), RLBench (success rate 61.7% vs 24.3%), Open-X, and real robot manipulation tasks.

Robot Learning from Diverse Environments. Learning a single universal multi-task policy on various robots, cameras, and task configurations is a challenging problem that is receiving increasing attention [9]. One recent strategy to solving this problem consists in training a network that takes as input a robot’s kinematic structure [10, 11]. For instance, MetaMorph [12] tokenizes the robot’s kinematic tree, which is then used as a transformer prompt to produce per-joint action. However, these methods are only evaluated in simple walking environments of Mujoco [13] without any variations in the task, the camera viewpoint, and the workspace. Several recent works have also looked into leveraging a large vision-language model (VLM) to directly train robots, through behavior cloning, on a collection of crowd-sourced datasets such as Open-X Embodiment [14, 15, 5]. Despite their impressive results, these methods require a non-trivial action normalization procedure to reduce the gap between action spaces such as end-effector poses, joint positions, and velocities in various world frames. However, this manual engineering procedure is custom-designed for each training dataset, which severely affects the generalization and interpretability of these models. In contrast, our method directly learns and visually predicts the movement of robots’ kinematic chains.

Transformers for Manipulation Learning. Transformers are widely used in manipulation learning [4, 16]. However, existing architectures require a fixed number of camera viewpoints, which is not guaranteed when learning is performed across various environments. Both RT-1/2 and RT-X select one canonical view from each Open-X dataset [5, 4, 17]. RVT [18] supports multiple but predefined viewpoints. PerAct [3] trains transformers over point-clouds and is not limited by camera viewpoints. But methods based on point clouds still suffer from the lack of datasets with depth information, as well as scale discrepancies [19, 20]. In contrast, our visual kinematics transformer (VKT) is a convolution-free architecture that is solely based on attention mechanisms in the RGB space, without the need for depth data. VKT also supports an arbitrary number of camera viewpoints.

Intermediate Action Representations. Instead of only using low-level robot commands, incorporating intermediate actions into control policies has been shown to increase data efficiency [21]. RT-H is a hierarchical policy that predicts language motions such as “move arm forward” [22]. SWIM uses affordance maps to unify grasping actions of humans and robots [21]. General Flow employs a labeling process to translate end-effector motions into point-cloud movements [23]. IMOP [24] uses invariant regions to describe key manipulation poses. RT-Trajectory predicts the future 2D trajectory of the end-effector’s center in a given image [25]. In contrast, our method forecasts the movements of the entire kinematic chain in multiple camera viewpoints. This representation provides a visually grounded action abstraction, while capturing detailed and small motions at the same time.

We consider the problem of learning a single universal multi-task policy from diverse robots and workspaces through a unified action space. Our key insight lies in the use of the visual kinematic chain, which is the projection of the robot’s high-dimensional kinematic structure into a set of pixels in the image 2D plane. The visual kinematic chain can be analytically generated with camera parameters and robot models without any manual engineering. We show that this simple representation allows us to train a single visual language-conditioned policy over multiple environments. The VLM policy is trained with a single objective of forecasting the future visual kinematic chain movements without knowing any low-level robot actions. Our proposed transformer architecture is built entirely with attention layers and can predict consistent multi-view chains from an arbitrary number of camera viewpoints. To deploy the multi-task policy in a specific environment, one simply needs to freeze the transformer and train a tiny head to project the actions into actual robot commands. An overview of our framework is illustrated in Figure 1.

Fitting any Structure with Point-Set Matching. The kinematic structures of various robots differ significantly. To forecast different structures, we propose to render them to point sets by uniformly sampling points along links, as illustrated in Figure 2. Let $P=\{p_{i}\}_{i=1}^{M}$ denote the point-set rendered from the ground-truth configuration of a kinematic chain, and $Q=\{q_{i}\}_{i=1}^{N}$ denote the point-set predicted by VKT, where $p_{i}\in\mathbb{R}^{2},q_{i}\in\mathbb{R}^{2}$. We first compute the pairwise euclidean distance matrix $C\in\mathbb{R}^{M\times N}$, where $C_{ij}=\parallel p_{i}-q_{j}\parallel$. We then minimize the earth-moving distance $\operatorname{EMD}(P,Q)=\sum_{i,j}\gamma^{*}_{ij}C_{ij}$, where $\gamma^{*}$ is a stochastic assignment matrix that matches points in $P$ and $Q$. We find $\gamma^{*}$ using the Sinkhorn-Knopp algorithm to solve the following optimization problem online for each mini-batch [26, 8].

Visual Kinematics Transformer (VKT). Our VKT network accepts a language instruction and an RGB image as input and predicts a sequence of $T$ point-sets. The dimensions of the kinematics tokens are $T\times D$, where $D$ denotes the embedding size, and each token predicts one point-set. The first point-set in the predicted sequence describes the current state of the kinematic chain in the given RGB image. The remaining point-sets are predictions of the future states of the kinematic chain in the next $T-1$ time-steps, i.e., future robot movements. The network is trained to match each point-set to the ground-truth kinematic structure at the corresponding time-step by minimizing the earth-moving distance. The same process is applied to multiple RGB images taken from different viewpoints. The network is also trained to make the future point-sets predicted from different viewpoints consistent with each other. The kinematics tokens are concatenated with the text tokens as query tokens. The concatenated query and vision tokens are forwarded to our proposed multi-view dual attention block, to interweave the information between query and vision tokens and across multiple camera viewpoints.

The multi-view dual attention block performs dual attentions between query and vision tokens within each single view, indicated by the green blocks in Figure 3, and multi-view attentions for query and vision tokens independently but across multiple camera viewpoints, shown as the blue blocks in Figure 3. This dual attention enables the query token to learn the visual kinematic chain movements. The multi-view attention enables a spatially-consistent prediction across multiple viewpoints. The output kinematics tokens are projected to $T$ point-sets using a small MLP network. Each point-set contains $N$ 2D points on the input image. The VKT architecture is illustrated in Figure 3.

Compared to existing robot learning transformers [18, 17], our proposed VKT is convolution-free and built solely with attention layers. Therefore, VKT supports an arbitrary number of camera viewpoints, which is an important advantage as camera setups are different among various environments.

Deploying to a Specific Environment. We freeze the trained VKT, drop the point-set prediction branch, and only use the kinematics tokens. We apply a 1D convolution to project the $T$ kinematics tokens into low-level robot actions for the next $T$ time-steps. Unlike the VKT backbone which is trained from multiple environments, the 1D convolution head is trained separately for each environment. The corresponding proposed architecture is illustrated in Figure 4.

We evaluate VKT on Calvin [28] and RLBench [29], two standard benchmarks in language-conditioned multi-task manipulation learning. Further, we evaluate VKT in a subset of Open-X Embodiment [4] and real robot experiments. We aim to answer the following questions: (1) Can visual kinematics forecasting improve specialized agents in each environment? (2) Can VKT serve as a strong general agent in multiple environments? (3) Can VKT be efficiently trained with real-world demonstrations and solve manipulation tasks with real robots?

Setup. We adopt the setting of RT-X [5] for imitation learning from RGB images without depth. RLBench contains various task categories, such as pick-and-place, tool use, high-precision operations, screwing, tasks with visual occlusion, and long-term manipulation. RLBench provides five RGBD cameras. We use the RGB signal from the front, left, and right cameras. We choose 17 tasks based on camera visibility and the task categorization of Hiveformer [30]. We use 100 recorded trajectories per task for training, evaluate the trained agent on 25 independent trials, and report the average success rate. Calvin requires policies to accomplish multiple tasks sequentially in each episode. We follow the convention of [28] of training on the pre-collected demonstrations and reporting the average number of successful steps per episode (success length). Examples of Calvin and RLBench tasks are shown in Figure 5. Additional visualizations are provided in the supplementary material.

Implementation Details. We use CLIP ViT-B/16 [27] as text and vision encoders. We resize all RGB inputs to $224\times 224$. We use 8 multi-view dual attention blocks ($L=8$) and predict visual kinematics chains for a horizon of 20 time-steps ($T=20$) and set size $D$ to $512$ in all experiments.

We have shown that visual kinematics forecasting of quasi-static robot movements can be a building block technique for developing general agents across diverse robot learning environments. However, there are still some open questions that are worth discussing. For example, we did not utilize wrist cameras in VKT because it does not take images of the robot. This can be resolved with RGBD wrist camera and virtual view projection [18]. Also, some public real-robot datasets lack robot URDF files and reliable camera parameters, which creates several challenges for applying our method. Finally, it is worth noting that one can re-purpose VKT to forecast not only robot movements, but also human poses and movements in order to leverage abundant human demonstration videos for robot learning.