KINet: Unsupervised Forward Models
for Robotic Pushing Manipulation

The keypoint detector ($f_{\mathrm{kp}}$, Fig. 2) is a mapping from RGB visual observations of the scene to a lower-dimensional set of $K$ keypoint coordinates $\{x_{t}^{k}\}_{k=1\dotsc K}=f_{\mathrm{kp}}(I_{t})$. The keypoint coordinates are learned by capturing the spatial appearance of the objects in an unsupervised manner. Specifically, the keypoint detector receives a pair of initial and current images $(I_{0},I_{t})$ and uses a convolutional encoder to compute a $K$-dimensional feature map for each image $\Phi(I_{0}),\Phi(I_{t})\in\mathbb{R}^{H^{\prime}\times W^{\prime}\times K}$. The expected number of keypoints is set by the dimension $K$. Next, the feature maps are marginalized into a 2D keypoint coordinate $\{x_{0}^{k},x_{t}^{k}\}_{k=1\dotsc K}\in\mathbb{R}^{2}$. We use a convolutional image reconstruction model ($f_{\mathrm{rec}}$, Fig. 2) with skip connections to inpaint the current image frame using the initial image and the predicted keypoint coordinates $\hat{I}_{t}=f_{\mathrm{rec}}(I_{0},\{x_{0}^{k},x_{t}^{k}\}_{k=1\dotsc K})$. With this formulation, $f_{\mathrm{kp}}$ and $f_{\mathrm{rec}}$ create a bottleneck to encode the visual observation in a temporally consistent lower-dimensional set of keypoints [14].

After factorizing the scene into $K$ keypoints, we build a graph $\mathcal{G}_{t}=(\mathcal{V}_{t},\mathcal{E}_{t},\mathbf{Z})$ (undirected, no self-loop) where keypoints and their pairwise relations are the graph nodes and edges. Keypoint positional and visual information are encoded into embedding of nodes $\{\mathbf{n}_{t}^{k}\}_{k=1\dotsc K}\in\mathcal{V}_{t}$ and edges $\{\mathbf{e}_{t}^{ij}\}\in\mathcal{E}_{t}$. We define the adjacency matrix to specify the graph connectivity as $\mathbf{Z}\in\mathbb{R}^{K\times K}$ where $z_{ij}\in[0,1]$ is the probability of the edge $\mathbf{e}^{ij}\in\mathcal{E}$. At timestep $t$, node embeddings encode keypoint positional and visual features $\mathbf{n}_{t}^{k}=\big{[}\Phi_{t}^{k},x_{t}^{k}\big{]}$. Edge embeddings contain relative positional information of each node pair $\mathbf{e}_{t}^{ij}=\big{[}x_{t}^{i}-x_{t}^{j},\|x_{t}^{i}-x_{t}^{j}\|_{2}^{2}\big{]}$.

Existing graph-based forward modeling methods construct the scene graph such that each node represents an object. The node embeddings are positional and visual features extracted using ground-truth object state information [11, 24, 32]. Our framework, however, does not rely on any ground-truth information supervision and only uses RGB images.

We extend the existing graph-based approaches [2, 25] and propose probabilistic graph-based forward models. The core of our forward model is probabilistic message-passing using edge-level latent variables $z_{ij}\in\mathbb{R}^{d}$ that represent the edge probabilities. A posterior network $p_{\theta}$ infers the elements of the adjacency matrix given the scene graph representation. In particular, we model the posterior as $p_{\theta}(z_{ij}|\mathcal{G}_{t})=\sigma(f_{\mathrm{enc}}([\mathbf{n}_{t}^{i},\mathbf{n}_{t}^{j}]))$ where $\sigma(.)$ is the sigmoid function.

The forward model $\hat{\mathcal{G}}_{t+1}=f_{\mathrm{fwd}}(\mathcal{G}_{t},u_{t},\mathbf{Z})$ predicts the graph representation at the next timestep by taking as input the current graph representation and the action vector ($f_{\mathrm{fwd}}$, Fig. 2) and performing $L$ rounds of probabilistic message-passing. Fig. 2-$f_{\mathrm{fwd}}$ illustrates the probabilistic message-passing operation. A single round of message-passing operation in the forward model can be described as,

where the edge-specific function $f_{\mathrm{edge}}$ first updates edge embeddings, then the node-specific function $f_{\mathrm{node}}$ updates each node embedding ${\mathbf{n^{\prime}}}^{k}$ by aggregating its neighboring nodes $N(k)$ information. Specifically, the updated edges are aggregated using the edge probabilities from the inferred adjacency matrix ${\mathbf{Z}}$ as weights. The action $u_{t}$ is also an input to the aggregation step. After performing the $L$ rounds of message-passing, the resulting updated node embeddings estimate the graph’s nodes at the next timestep $\{\hat{\mathbf{n}}_{t+1}\}$.

Recent models for forward prediction rely on fully connected graphs for message passing [15, 24, 32]. Our model, however, learns to probabilistically sample the neighbor information. Intuitively, this adaptive sampling allows the network to efficiently aggregate the most relevant neighboring information. This is specifically essential in our model as keypoints could provide redundant information if they are in close proximity.

The state decoder ($f_{\mathrm{state}}$, Fig. 2) transforms the predicted node embeddings of the updated graph $\hat{\mathcal{G}}_{t+1}$ to a first-order difference $\{\Delta\hat{x}_{t+1}^{k}\}=f_{\mathrm{state}}(\{\mathbf{\hat{n}}_{t+1}^{k}\})$, which is integrated once to predict the position of the keypoints in the next timestep $\{\hat{x}_{t+1}^{k}\}=\{x_{t}^{k}\}+\{\Delta\hat{x}_{t+1}^{k}\}$. To reconstruct the image at the next timestep, we borrow the reconstruction model $f_{\mathrm{rec}}$ from the keypoint detector $\hat{I}_{t+1}=f_{\mathrm{rec}}(I_{0},\{x_{0}^{k},\hat{x}_{t+1}^{k}\})$.

A learned KINet model can be used to perform model-based planning. We extend the Model Predictive Control (MPC) method [7] and propose the GraphMPC algorithm based on graph embeddings. The graph representation of the next timestep is estimated for multiple sampled actions using $f_{\mathrm{fwd}}$. The optimal action is selected such that it produces the most similar graph representation to the goal graph representation $\mathcal{G}^{\mathrm{goal}}$. We describe our GraphMPC algorithm with a horizon of $T$ as: ${u_{t}^{*}={\arg\max}\{\mathcal{S}\big{(}\mathcal{G}^{\mathrm{goal}},f_{\mathrm{fwd}}(\mathcal{G}^{t},\{u_{t:T}\})\big{)}\};\ t\in[0,T]}$. Unlike performing conventional MPC only with respect to positional states, GraphMPC allows for accurately bringing the objects to a goal configuration both explicitly (i.e., position) and implicitly (i.e., pose, orientation, and visual appearance).

We apply our approach to learn a forward model for multi-object manipulation tasks. The task involves rearranging the objects to achieve a goal configuration using pushing actions. We use MuJoCo 2.0 [27] to generate two simulated environments. We also test on a real robot environment. Fig. 3 demonstrates these three experiment setups.

In both of the simulated environments, we generate 10K episodes of random pushing on multiple objects (1-5 objects) where a simplified robot end-effector applies randomized pushing for 60 timesteps per episode. We collect the 4-dimensional action vector (pushing start and end location) and RGB images before and after each action is applied. Images are obtained using an overhead (TopView) and an angled camera (SideView).

We compare our approach with existing methods for learning object-centric forward models:

All models are trained on a subset of the simulated data (BlockObjs and YCBObjs) with 3 objects (8K episodes: 80% training, 10% validation, and 10% testing sets). To evaluate generalization to a different number of objects, we use other subsets of data with 1, 2, 4, and 5 objects ($\sim 400$ episodes for each case). We train our model separately on images obtained from the overhead camera (TopView) and the angled camera (SideView). We set the expected number of keypoints $K=6$ for BlockObjs, $K=9$ for YCBObjs and RealYCB. For the message-passing operation, the number of rounds is set to $L=3$. Note that $f_{\mathrm{enc}}$, $f_{\mathrm{edge}}$, and $f_{\mathrm{node}}$ are MLPs with three 128-dim hidden layers.

First, we evaluate the forward prediction accuracy. Our model factorizes the observation into a set of keypoint representations and accurately estimates the future appearance of the scene conditioned on external action.

Using BlockObjs data, we quantify the effectiveness of our model in comparison with Forw, ForwInv, IN, and VisIN baselines (Table I). We separately train and examine each model on TopView and SideView images. The prediction error is computed as the average distance between the predicted and ground-truth positional states. VisIN performs the best among baseline models as it builds object representations with explicit ground-truth object positions and their visual features. Our model, on the other hand, achieves a comparable performance to VisIN while it does not rely on any supervision beyond the scene images. Forw and ForwInv baselines have similar supervision to ours but are significantly less accurate. This emphasizes the capability of our approach in learning a rich graph representation of the scene and an accurate forward model, while relaxing prevailing assumptions of the prior work on the structure of the environment and the availability of ground-truth state information. KINet - determ & no ctr will be discussed in Section V-D.

We design a robotic manipulation task of rearranging a set of objects to the desired goal state using MPC with pushing actions. For all models, we run 1K episodes with randomized object geometries and initial poses and a random goal configuration of objects. The planning horizon is set to $T=40$ timesteps in each episode. For our model, we perform GraphMPC based on graph embedding similarity as described in Section III-F. For all baseline models, we perform MPC directly on the distance to the goal. Fig. 4a shows MPC results of BlockObjs based on Top View observations. Our approach is consistently more accurate and faster in reaching the goal configuration compared to the baseline models.

One of our main motivations to learn a forward model in the keypoint space is to eliminate the dependency of model formulation on the number of objects in the scene which allows for generalization to an unseen number of objects, object geometries, background textures, etc.

Further, we test the performance of the model on unseen background textures (i.e., the table texture). Since the keypoint extraction relies on visual features of salient objects, our model is able to perform the control tasks by ignoring the background and assigning keypoints to the moving objects. Fig. 4b compares the MPC results for the KINet trained on a fixed white background ($\mathrm{fixed}_{train}$), zero-shot generalization of KINet trained on the fixed background to randomized backgrounds ($\mathrm{rand}_{gen}$), and KINet trained directly on randomized backgrounds ($\mathrm{rand}_{train}$). As expected, although the MPC converges, the final distance to the goal configuration is larger for $\mathrm{rand}_{gen}$. This final error is statistically at the same level of accuracy for $\mathrm{fixed}_{train}$ and $\mathrm{rand}_{train}$. Qualitative examples are included in Fig. 4.

We compare the performance of our framework with V-CDN [18] baseline which is also a keypoint-based model to learn the structure of physical systems and perform future predictions and potentially generalize to an unseen number of objects. Although V-CDN is formulated to extract the causal structure of a fixed system through visual observations, we pretrain its perception module on our randomized YCBObjs dataset for a direct comparison to our model. To perform a control task, we condition the V-CDN on the action by adding an encoding of the action vector to each keypoint embedding. Fig. 5 compares the MPC results (normalized to the number of objects) of our method and V-CDN both trained on a subset of 3 YCB objects. The MPC performance of both models is comparable for rearranging 3 objects (green lines). However, V-CDN is most accurate for the number of objects it has been trained on and significantly less accurate when generalizing to an unseen number of objects. We attribute this to two reasons: (1) Although using keypoints, V-CDN is formulated to infer explicit causal structure and physical attributes for the graph representation of fixed environments which does not necessarily carry over to unseen circumstances. (2) Unlike our method, V-CDN does not take into account the visual features in the model and only uses the keypoint positions.

We justify the major choices we made to formulate the model with ablation studies: the probabilistic graph representation (KINet - determ), and the contrastive loss (KINet - no ctr). As shown in Table I, the best forward prediction for both TopView and SideView images is achieved when the model is probabilistic and trained with a contrastive loss. The contrastive loss is an essential element in our approach to ensure the learned forward model is accurately action conditional. Also, with a probabilistic graph representation, our model achieves better generalization compared to the deterministic variant. This performance gap is more evident when generalizing to unseen geometries (Table II).

We ran additional experiments to examine the effect of the number of keypoints. Fig. 7-b shows the distance to the goal configuration after executing 30 planning steps with KINet, trained with different numbers of keypoints. Notably, for object rearrangement involving 3 YCB objects, KINet with fewer than 6 keypoints exhibited considerably larger errors. This indicates that in environments with complex visual features, optimal performance is achieved by a larger number of keypoints to ensure proper keypoint assignment (i.e., with $K=3$, Fig. 7-a shows one object remains unassigned).

Since the number of objects is not hard-coded into our formulation, KINet allows for training on a variable number of objects. We experimented with training on varying numbers of YCBObjs ($N$=1 to 5), as opposed to a fixed number ($N$=3). The results, shown in Figure 7-c, demonstrate that training KINet on a variable number of objects leads to a more balanced performance. As expected, for $N>3$, the final distance to the goal improved since the model was trained on these scenarios instead of generalizing. However, we did not observe a statistically significant change for $N\leq 3$.

After training the model in YCBObjs simulation, we test for the real robot performance (see Fig. 3). Fig. 8 (top row) shows examples of the control task. The framework successfully transferred to the real setting and was able to perform the RealYCB object rearrangement tasks. We observed that some of the detected keypoints were slightly less consistent compared to the simulation. We attribute this to the sim2real domain gap (e.g., the tabletop texture and camera noise).

We also tested for generalization to unseen objects (Fig. 8, bottom row). Our model generalized to unseen objects by appropriately distributing the keypoints across the objects. We noticed that in some instances a keypoint is attached to the tabletop due to its visual feature. Regardless, the control task is achieved as the majority of keypoints are attached to the objects. The GraphMPC algorithm selects the action based on maximizing graph similarity to the goal scene across all nodes which reduces the effect of a single inconsistent keypoint on the rearrangement task.

We also compare with V-CDN baseline in the real setting. Since the ground-truth positions are unknown in the real world, we use image similarity to quantify the accuracy of the rearrangement task. Specifically, we used the pretrained VGG [26] to measure the cosine similarity between the scene image and the goal image in the feature space to obtain MPC results for the real-world setting. As shown in Fig. 9, our model is consistently more accurate in rearranging the real scene and archives a significantly higher similarity to the goal image.