Learning what you can do before doing anything

TL;DR

This work tackles learning an agent's action space from purely visual observations by training a stochastic video-prediction model with a minimality constraint and introducing a composability loss to yield a disentangled, composable latent representing actions. The CLASP framework further grounds this latent space to real actions via a lightweight bijection, enabling action-conditioned video prediction and planning in the learned space with significantly fewer labeled examples. Empirical results on synthetic reacher data and the BAIR robot pushing dataset show that CLASP learns action-structured representations, improves prediction quality, and enables planning with data efficiency and robustness to appearance changes. Overall, the method demonstrates that action spaces can be learned from passive observations, offering a data-efficient route to imitation and control in environments with diverse visuals.

Abstract

Intelligent agents can learn to represent the action spaces of other agents simply by observing them act. Such representations help agents quickly learn to predict the effects of their own actions on the environment and to plan complex action sequences. In this work, we address the problem of learning an agent's action space purely from visual observation. We use stochastic video prediction to learn a latent variable that captures the scene's dynamics while being minimally sensitive to the scene's static content. We introduce a loss term that encourages the network to capture the composability of visual sequences and show that it leads to representations that disentangle the structure of actions. We call the full model with composable action representations Composable Learned Action Space Predictor (CLASP). We show the applicability of our method to synthetic settings and its potential to capture action spaces in complex, realistic visual settings. When used in a semi-supervised setting, our learned representations perform comparably to existing fully supervised methods on tasks such as action-conditioned video prediction and planning in the learned action space, while requiring orders of magnitude fewer action labels. Project website: https://daniilidis-group.github.io/learned_action_spaces

Figures (15)

• Figure 1: Using latent composition to recover actions from passive data. a) Two sequences starting from different initial states but changing according to the same actions. Without requiring labels, our model learns to represent the action in sequences like these identically. We train a representation $z$ to capture the dynamics of the scene and its compositional structure: applying ($z_1$ and $z_2$) should have the same effect as applying the composed representation $g(z_1, z_2)$. These properties capture the fact that effector systems, such as a robot arm, use the same composable action space in many different states. b) The learned action space $z$ recovered by our method (PCA visualization). Points are colored by the true action $u$: true actions can be easily decoded from $z$, validating that the structure of the action space has been captured.
• Figure 2: Components of the proposed architecture. Left: The stochastic video prediction model, shown for one timestep. During training, we estimate the latent variable ${z}_t$ using the approximate inference network ($\text{MLP}_{\text{infer}}, \text{CNN}_e$) from the current and previous image. At test time, we produce $z_t$ using the prior distribution $p({z}) \sim \mathcal{N}(0, I)$. Future frames are estimated by passing $z_{t}$ together with images $x_{t-1}$ through the generative network (LSTM, $\text{CNN}_d$). Please refer to Appendices \ref{['app:svp']} and \ref{['app:expp']} for architectural details. Right: Composability training. Latent samples $z$ are concatenated pairwise and passed through the composition network $\text{MLP}_{\text{comp}}$ that defines a distribution over $\nu$ in the trajectory space. A sampled value of $\nu$ is decoded into an image through the same generative network (LSTM and CNN$_d$) and matched to the final image in the composed sequence.
• Figure 3: Visualization of the learned action space, $z$, on the reacher dataset. Each of the 1000 points depicts a value of $z$ for a different frame pair from the dataset. We plot the projection of $z$ onto the first two principal components of the data. Each point is colored by the value of the ground truth rotation, in radians, depicted in the two images used to infer $z$ for that point. a) The latent space learned by the baseline model has no discernible correspondence to the ground truth actions. b) Our method learns a latent space with a clear correspondence to the ground truth actions. In the Appendix, Fig. \ref{['fig:pca_app']} further investigates why the baseline fails to produce a disentangled representation.
• Figure 4: Illustration of how the learned representation can be used for a) action-conditioned prediction by inferring the latent variable, $z_t$, from the action, and b) visual servoing by solving the control problem in latent space through iterated rollouts and then mapping the latent variable to robot control actions, $u_t$.
• Figure 5: Transplantation of action representations $z$ from one sequence to another. We infer action representations from the donor sequence and use them to create the recipient sequences from a different initial state. a) the reacher dataset. The previous frame is superimposed onto each frame to illustrate the movement. b) the BAIR dataset. The previous and the current position of the end effector are annotated in each frame (red and blue dots, respectively) to illustrate the movement. c) reacher with varying backgrounds. d) reacher with varying agent shape. The synchronization of movement in the sequences suggests that the learned action representation is disentangled from static content. Best viewed on a screen. Additional generated videos are available at: https://daniilidis-group.github.io/learned_action_spaces/.