A Generalist Dynamics Model for Control

TL;DR

This work investigates transformer sequence models as dynamics models (TDMs) for control, demonstrating strong cross-environment generalization in few-shot and zero-shot settings while also delivering accurate single-environment predictions for MPC. By tokenizing trajectories and integrating TDMs into MPC with random shooting (and proposal) planners, the authors show TDMs can outperform baselines and even specialist models in data-efficient generalization scenarios. The key contributions include establishing few-shot and zero-shot cross-environment generalization, comparing generalist pre-training strategies, and illustrating the value of planning-based use of dynamics models over direct policy generalization. The findings suggest TDMs as a promising foundation model for robotics, capable of leveraging broad prior experience to accelerate learning and adaptation across diverse tasks and morphologies.

Abstract

We investigate the use of transformer sequence models as dynamics models (TDMs) for control. We find that TDMs exhibit strong generalization capabilities to unseen environments, both in a few-shot setting, where a generalist TDM is fine-tuned with small amounts of data from the target environment, and in a zero-shot setting, where a generalist TDM is applied to an unseen environment without any further training. Here, we demonstrate that generalizing system dynamics can work much better than generalizing optimal behavior directly as a policy. Additional results show that TDMs also perform well in a single-environment learning setting when compared to a number of baseline models. These properties make TDMs a promising ingredient for a foundation model of control.

Figures (13)

• Figure 1: Schematic overview of the data regimes for which we show experimental results. These regimes are characterized by how much data from the target environment is available to the agent, and how much (potentially generalizable) experience has been collected in other environments. The experiments both demonstrate that TDMs are capable single-environment models (marked purple) and generalize across environments (marked yellow). If sufficient data from the target environment is available, we can learn a single-environment specialist model (section \ref{['sec:expert_model_results']}). If there are only small amounts of data from the target environment, but more data from other environments, a generalist model can be pre-trained and then fine-tuned on the target environment (section \ref{['sec:finetuning_generalist_results']}). Finally, if we are able to train a generalist model on large amounts of data from different environments, we can zero-shot apply this model to our target environment without fine-tuning (section \ref{['sec:zero-shot-generalization-results']}). We also show an example for unsuccessful generalization (no color) in section \ref{['sec:walker_generalization_negative_example']}.
• Figure 2: Illustration of the tokenization for $n=3$ and $m=2$. Starting from $o_1$, performing action $a_1$ will result in the next observation $o_2$ and the reward $r_2$. The constant separator tokens $t_5$ and $t_{12}$ are inserted to indicate the start of a new environment step.
• Figure 3: The procedural walker universe.
• Figure 4: Performance of TDMs and baseline models when trained on data from the environment they are tested on. We observe that TDMs consistently outperform baselines. This finding is robust when switching the training distribution to a different task in the same environment (red lines for walker and humanoid). We also compare with the ground truth models (black line). We evaluate the models by doing MPC with a very basic random shooting planner. The planner uses $K=128$ samples for cartpole, $K=64$ samples for walker, and horizon $N=20$ for humanoid. For very short planner horizons $N$, the planner is too myopic, and for very long horizons, the number of samples $K$ is insufficient for the random shooting planner to consistently discover a near-optimal action sequence. Therefore, when keeping $K$ fixed, there is an intermediate sweet-spot planner horizon. We report mean values averaged over at least $4$ episodes, shaded areas indicate $68\%$ confidence intervals.
• Figure 5: Using the TDM for MPC with a proposal policy for humanoid stand. The subfigures correspond to different levels of additive noise $\sigma$. Best results are obtained for moderate additive noise (this ensures that the bias of the proposal policy is not washed out) and larger horizons $N$ (this ensures that the planner does not become too myopic). The resulting MPC agent both works better than the pure proposal policy (red line), and needs less imaginary samples $K$ than the random shooting planner (see Fig. \ref{['fig:same_domain_results_humanoid']}). We report mean values averaged over at least $4$ episodes, shaded areas indicate $68\%$ confidence intervals.