# LIMT: Language-Informed Multi-Task
Visual World Models

Elie Aljalbout1,2,∗, Nikolaos Sotirakis1,3,∗, Patrick van der Smagt1,4,
Maximilian Karl1, Nutan Chen1

###### Abstract

Most recent successes in robot reinforcement learning involve learning a
specialized single-task agent. However, robots capable of performing multiple
tasks can be much more valuable in real-world applications. Multi-task
reinforcement learning can be very challenging due to the increased sample
complexity and the potentially conflicting task objectives. Previous work on
this topic is dominated by model-free approaches. The latter can be very
sample inefficient even when learning specialized single-task agents. In this
work, we focus on model-based multi-task reinforcement learning. We propose a
method for learning multi-task visual world models, leveraging pre-trained
language models to extract semantically meaningful task representations. These
representations are used by the world model and policy to reason about task
similarity in dynamics and behavior. Our results highlight the benefits of
using language-driven task representations for world models and a clear
advantage of model-based multi-task learning over the more common model-free
paradigm.

††footnotetext: 1During this work, all authors were affiliated with the
Machine Learning Research Lab at Volkswagen Group, Germany. 2E.A. is currently
with the Robotics and Perception Group, at the Department of Informatics of
the University of Zurich (UZH) and the Department of Neuroinformatics at UZH
and ETH Zurich, Switzerland. 3Technical University of Munich, Germany.
4Faculty of Informatics, Eötvös Loránd University, Budapest, Hungary. ∗Shared
first authorship.

> Keywords: Multi-Task Learning, Language-Conditioned World Models, Model-
> Based Reinforcement Learning

##  1 Introduction

Reinforcement learning (RL) methods have shown great potential in various
robotic control tasks such as manipulation and locomotion [1, 2]. The majority
of successes in this domain are in single-task settings, where the agent is
concerned with finding control policies for a single task. Ideally, a single
agent should be capable of performing various tasks and smoothly switching
between different task performances. This is especially important when
considering the high cost of robotic systems such as manipulators. The goal of
multi-task reinforcement learning (MTRL) is to learn such a single policy
capable of performing multiple tasks by jointly optimizing the individual task
objectives [3]. This joint training process can be beneficial, for instance,
in terms of bootstrapping the learning of complex tasks [4]. However, it can
be highly challenging, not only due to the increased complexity of the
problem, but also because different tasks can have conflicting objectives
leading to unstable training. The majority of research on MTRL considers
model-free methods. However, model-free RL methods are considerably sample-
inefficient even in single-task learning. This property is undesirable in
robot learning systems, where environment interactions are very expensive and
hard to obtain in the real world. Alternatively, it is possible to train
robotic control policies in simulation and transfer them to the real world.
However, this process is not trivial and presents multiple challenges [5, 1,
6, 7]. This efficiency problem becomes even more pronounced when dealing with
high-dimensional and complex observations, such as the ones encountered in
visual RL. The extension to the more complex multi-task setting can further
exacerbate these problems.

Model-based reinforcement learning (MBRL) methods tend to have superior sample
efficiency compared to the model-free approach [8, 9]. This boost in
efficiency is achieved by incorporating a model of the environment, which
allows the agent to simulate environment interactions for policy search. This
capability is crucial for multi-task learning as it enables the agent to
effectively share knowledge across tasks, reducing the overall amount of data
needed to achieve proficient performance on each task. By learning a shared
model that captures the dynamics relevant to multiple tasks, the agent can
develop a more holistic understanding of its embodiment and its environment
dynamics.

In this work, we propose a model-based vision-based method for multi-task
learning. Our method incorporates language-conditioning in the world model as
well as the actor and critic. These models are conditioned on language
embeddings from a pre-trained language model. By doing so, we leverage these
semantically meaningful task representations to boost the parameter sharing in
the world model and policy. We evaluate the proposed approach using multiple
robotic manipulation tasks from the CALVIN dataset [10]. We compare our work
to baselines from single-task MBRL, and model-free MTRL methods. Our
experiments demonstrate successful learning of a multi-task world model and
its usage for learning a policy for multiple manipulation tasks. We validate
the benefit of using language-embeddings as task representations for the world
model and the policy, and demonstrate a substantial performance improvement in
comparison to model-free MTRL.

Figure 1: We train a model-based actor-critic agent using a multi-task world
model. The actor, critic and the world model are conditioned on tokens from a
vision and proprioception tokenizer as well as a language embedding from a
pretrained language model. The latter encodes instructions for the different
tasks. The resulting agent can perform multiple tasks depending on the
instruction input.

##  2 Method

Previous work has shown the benefits of multi-task policy training for
bootstrapping the learning of harder tasks and increasing the individual
tasks’ sample efficiency [4]. We postulate that a similar benefit can be
observed for learning (visual) world models for robotics. Intuitively,
multiple tasks share similar dynamics and perception components. For instance,
a world model trained on opening a drawer has much in common with another
trained on opening cupboards. Training a common world model can leverage these
task similarities to boost sample efficiency by sharing data across tasks. To
allow the model to reason about task similarity, a semantically meaningful
task representation can be essential. Hence, we propose a model-based MTRL
approach, that trains a language-conditioned world model and actor-critic
agent. Our approach consists of four main components. A language model encodes
text instructions into structured embeddings. A tokenizer maps observations to
discrete token representations. A world model is a sequence model that
predicts future observations, rewards, and successes based on past
trajectories and task instruction embeddings. The actor-critic networks
respectively output control commands and value estimates based on the latent
state and embedding of the task description. We refer to our method as LIMT,
an acronym for language-informed multi-task visual world models. Our overall
approach is illustrated in Figure 1.

###  2.1 Tokenizer

Following the work in [11], we use a discrete autoencoder $(E,D)$ which is a
variant of Vector Quantized Variational Autoencoder (VQVAE) [12], equipped
with attention blocks as proposed in [13] and additionally trained with a
perceptual loss [13] [14] [15]. The reason we chose to use discrete
representations is that transformer networks, such as the one underlying our
world model, are particularly successful at modeling sequences of discrete
tokens [16] [17]. We further extend it to also handle proprioception data. The
encoder $E$ accepts an observation $(x,\theta)$ consisting of an image
observation $x$ and a $d$-dimensional proprioception vector $\theta$ and
converts these to $K=K_{x}+K_{\theta}$ tokens of dimension $d_{\mathrm{enc}}$.
Specifically, we use $K_{x}$ tokens to represent the image and $K_{\theta}$
tokens to represent proprioception. $E$ maintains two separate codebooks of
the same vocabulary size $N$ for the image and proprioception tokens
respectively, $C_{x}=\\{c_{x}^{i}\\}_{i=1}^{N}$,
$C_{\theta}=\\{c_{\theta}^{i}\\}_{i=1}^{N}$, where
$c_{x,\theta}^{i}\in\mathbb{R}^{d_{\mathrm{enc}}}$.

Concretely, the input of $E$ consists of observations $x\in\mathbb{R}^{H\times
W\times C}$, $\theta\in\mathbb{R}^{d}$, where $H,W,C$ denote the image height,
width, and number of channels respectively, and $d$ is the dimension of the
proprioception vectors. $E$ passes $x$ through a series of convolutional and
self-attention layers to obtain features $z_{x}\in\mathbb{R}^{h\times w\times
d_{\mathrm{enc}}}$, where $h\times w=K_{x}$. It then computes a quantized
embedding representation according to the nearest neighbor in $C_{x}$ using
the Euclidean distance

$q_{x}\left((z_{x})_{ij}\right)=\arg\min_{c\,\in\,C_{x}}{\left|(z_{x})_{ij}-c\,\right|}^{2}\quad,\
i\in[h]\,,\ j\in[w].$ (1)

$E$ spatially decomposes the image into $h\times w$ feature vectors and
assigns an element of the codebook to each one of them. Similarly, the
proprioception vector $\theta$ is linearly projected to a latent vector
$z_{\theta}\in\mathbb{R}^{d_{\mathrm{enc}}}$ via an affine layer, before being
quantized according to the codebook $C_{\theta}$

$q_{\theta}(z_{\theta})=\arg\min_{c\,\in\,C_{\theta}}{\left|z_{\theta}-c\,\right|}^{2}.$
(2)

The output $E(x,\theta)=(q_{x}(z_{x}),q_{\theta}(z_{\theta}))$ of the encoder
comprises the latent representation of the observations, which is subsequently
used as an input to the other components, namely the world model and the
actor-critic. Since the set of possible encoder outputs is discrete, we can
define the token representation of an observation $(x,\theta)$ as
$w=(w_{x},w_{\theta})$, where $w_{x},w_{\theta}\in\\{1,..,N\\}$. For training
purposes, a decoder $D$ maps these embeddings back into the observation space.
For the RGB images, $D$ uses a network consisting of multiple convolutional,
self-attention and upsampling layers, while the proprioception data is decoded
via a single linear layer. We train our discrete autoencoder using the loss

$\displaystyle L_{A}(E,D,x,\theta)=$ $\displaystyle\ \lVert
x-D(E(x))\rVert_{1}+\lVert D(E(\theta))-\theta\rVert_{2}^{2}$ (3)
$\displaystyle+$ $\displaystyle\ \lVert
sg(E(x,\theta))-z_{x,\theta}\rVert_{2}^{2}+\lVert
E(x,\theta)-sg(z_{x,\theta})\rVert_{2}^{2}$ (4) $\displaystyle+$
$\displaystyle\ L_{\mathrm{perceptual}}(x,D(E(x))),$ (5)

where $sg()$ is the stop-gradient operator. The right side of (3) denotes the
reconstruction loss for images and proprioception respectively, while the
terms in (4) constitute the commitment loss, which ensures that the
unquantized latent vectors are close to their corresponding discrete
representations. (5) is a perceptual loss [14] [15] shown in equation(12) in
the appendix, and computed with a pre-trained VGG16 CNN and given the ground-
truth and the reconstructed images as inputs.

###  2.2 Language Model

We use a pre-trained version of Sentence-BERT [18], MiniLM-L6-v2 a large
language model that has been tuned on semantic similarity, to encode natural
language directives. In [19], SBERT embeddings were found to be more suitable
for language-conditioned policy learning compared to alternatives such as BERT
[16] and CLIP [20] embeddings. We select this type of language model because
it has a semantically meaningful structure of the embedding space, where
encoded sentences can be compared using cosine similarity. Our hypothesis is
that the embeddings of language instructions describing the same task or tasks
with similar dynamics cluster nicely together and away of dissimilar tasks.
This is shown in Figure 3. This natural clustering can help our agent reason
about similarities among different tasks and thus learn skills and dynamics
faster.

###  2.3 Dynamics Learning

The world model, denoted as $G$, is an autoregressive transformer similar to
the one in [11]. In addition to actions and observation tokens, the
transformer is conditioned on language embeddings. Given a trajectory of $T$
timesteps $(x_{t},\theta_{t},a_{t},l^{\prime})_{t=\tau}^{T}$, where
$l^{\prime}$ denotes the language instruction, we first compute the sequence
$(w_{t},a_{t},l)_{t=\tau}^{T}$, where $w^{k}_{t}\in\\{1,..,N\\}^{K}$ is the
joint tokenized representation of the observations $x_{t},\theta_{t}$, and $l$
is the instruction embedding from our language model. $G$ predicts the next
observation tokens $\hat{w}_{t+1}$, the reward $\hat{r_{t}}$ and the episode
end $\hat{d_{t}}$

$\displaystyle\hat{w}_{t+1}\sim p_{G}(\hat{w}_{t+1}\mid w_{\leq t},a_{\leq
t},l)$ $\displaystyle\hat{r_{t}}\sim p_{G}(\hat{r_{t}}\mid w_{\leq t},a_{\leq
t},l)$ (6) $\displaystyle\hat{d_{t}}\sim p_{G}(\hat{d_{t}}\mid w_{\leq
t},a_{\leq t},l).$

The transformer predicts the tokens $\hat{w}_{t+1}$ at $t+1$ autoregressively

$\hat{w}_{t+1}^{k+1}\sim p_{G}(\hat{w}_{t+1}^{k+1}\mid w_{t+1}^{\leq
k},w_{\leq t},a_{\leq t},l).$ (7)

$G$ operates on a context window $H\in\mathbb{N}$, such that only the last $H$
timesteps influence the predictions. Besides the task-dependent reward and
termination condition, the prediction of the next latent state is also
conditioned on the language directive. The loss function for $G$ is

$\begin{split}L_{G}=\sum_{t=\tau}^{\tau+H}\left[\sum_{k=1}^{K}w_{t+1}^{k+1}\log
p_{G}(w_{t+1}^{k+1}\mid w_{t+1}^{\leq k},w_{\leq t},a_{\leq t},l)\right]\\\
+\rho\lVert r_{t}-\hat{r}_{t}\rVert_{2}^{2}+d_{t}\log p_{G}(\hat{d}_{t}\mid
w_{\leq t},a_{\leq t},l),\end{split}$ (8)

where the first term on the right side of the equation is the cross entropy
loss between predicted and ground-truth tokens. The second term is the reward
prediction. The third is a cross-entropy loss for the termination label.
$\rho\in\mathbb{R}$ is a hyperparameter for weighing reward loss against the
other terms.

###  2.4 Policy Learning

The actor and critic networks receive the observation token embeddings as
inputs, concatenated with the instruction embeddings, and the predicted
proprioception information. They are jointly trained in latent imagination as
in Dreamer [9]. Given an encoded instruction $l$ and observation token
embedding $q(z_{t_{0}})$ at timestep $t_{0}$, we start a rollout of length
$H$: For each $0<t<H$, the actor outputs an action $a_{t}$, based on $l,w_{t}$
and the world model predicts the next observation embedding $q(z_{t+1})$ and
reward $r_{t+1}$ which result from taking the action $a_{t}$. The critic is
trained to regress the $V^{\lambda}$ estimates, defined in equation (9) in the
appendix. To stabilize training, we maintain a target critic $\hat{v}_{\psi}$,
the weights of which we periodically update with those of our value function
$v_{\psi}$. $\hat{v}_{\psi}$ is used to compute the $V^{\lambda}$ returns, and
the objective of our value function $v_{\psi}$ is to estimate them.

###  2.5 Training Algorithm

We begin our training using the offline episodes,
$\mathcal{D}_{\mathrm{offline}}$ provided by CALVIN. In this work, we
concentrate on a subset of tasks $\mathcal{T}=\\{T_{i}\\}_{i}^{N}$. We gather
episodes where one of the tasks in $\mathcal{T}$ is completed in a separate
dataset $\mathcal{D}_{\mathrm{filtered}}$, used in the later stages of the
training process. First, we train our tokenizer on the image-proprioception
pairs of $\mathcal{D}_{\mathrm{offline}}$ for $N_{t}$ epochs until we observe
convergence. Subsequently, we train our world model on
$\mathcal{D}_{\mathrm{offline}}$ jointly with the tokenizer for another
$N_{w}$ epochs. Given that $\mathcal{D}_{\mathrm{offline}}$ contains numerous
episodes from tasks other than those in $\mathcal{T}$, we modify each of these
episodes during the training of our world model. Specifically, we relabel
every such episode with a task chosen from $\mathcal{T}$ with equal
probability, adjusting the ground-truth rewards and language instructions
accordingly. This increases the number of learning samples for our tasks of
interest and provides negative examples, which are otherwise not available in
the expert play data. In the next stage, we train the tokenizer and world
model for an additional $N_{f}$ epochs on $\mathcal{D}_{\mathrm{filtered}}$.
After that, we start performing rollouts using the actor’s policy, which we
collect in an online buffer $\mathcal{D}_{\mathrm{online}}$. During this
phase, we jointly train the tokenizer, the world model and the actor-critic on
both $\mathcal{D}_{\mathrm{online}}$ and $\mathcal{D}_{\mathrm{filtered}}$. We
balance sampling from the two datasets using an adaptive ratio
$p_{\mathrm{online}}$ denoting the proportion of samples from
$\mathcal{D}_{\mathrm{online}}$. $p_{\mathrm{online}}$ starts at 0 and grows
along with the size of $\mathcal{D}_{\mathrm{online}}$ until it reaches a
maximum ratio $p_{\mathrm{max}}$. Within each sampled episode, we select a
starting state from which our agent imagines a trajectory uniformly at random
across the time dimension. We randomize the task goal by taking the episode’s
true goal with provability $p$ and randomly sampling another goal from
$\mathcal{T}$ with probability $1-p$, using weights $w$. In this context,
specifying a goal task involves selecting a language instruction corresponding
to the task as an input to our policy.

##  3 Experimental Results

We design experiments to answer the following questions:

  1. (Q1)

Can we learn multi-task policies with language-informed world models and MBRL?

  2. (Q2)

Is MBRL more sample efficient than model-free RL in multi-task settings?

  3. (Q3)

Is multi-task training beneficial for MBRL for single and multi-task
performance?

  4. (Q4)

Are language embeddings good task representations for learning multi-task
world models?

Figure 2: We plot the success rate over a fixed budget of training epochs. We
consider model-free baselines based on multi-task soft-actor-critic (MT-SAC).
MT-SAC is trained on raw images (MT-SAC:raw) and another variant is trained on
image embeddings from our tokenizer (MT-SAC:token). We include a single-task
MBRL baseline (MBRL-ST) and variants of our method that use integer task
representations instead of language embeddings. The integer representations
are either only used in the actor critic (LIMT:nlac) or in both the world
model and actor-critic (LIMT:nl). LIMT and its variant have better sample
efficiency and success rate.

###  3.1 Setup

We train and deploy our agent on environment D of the CALVIN benchmark [10].
The agent controls a 7-DoF Franka Emika Panda robot, equipped with a parallel
gripper. It outputs actions in the form of relative Cartesian positions. The
environment consists of a table where the following static objects are
present: i) a drawer that can be opened or closed, ii) a slider that can be
moved left or right, iii) a button that toggles a green LED light iv) a switch
that can be flipped up to control a lightbulb. In addition to the static
objects, 3 colored rectangular blocks appear in the scene in different
positions and orientations. The agent’s state consists of RGB images from a
static camera, as well as 7-dimensional proprioception information indicating
the position of the end effector and gripper width. We focus and evaluate our
model on the following 6 tasks: open_drawer, close_drawer, move_slider_left,
move_slider_right, turn_on_lightbulb, turn_on_led. For each of the 6 tasks,
CALVIN [10] provides 7-14 distinct but semantically equivalent training
language instructions. We leverage the embeddings of those instructions to
train our agent. CALVIN also provides a single validation language directive
for each task, which is distinct but equivalent to the training instructions.
We use the validation directives when evaluating the performance of our agent.
During evaluation, we estimate our evaluation metrics, using 20 rollouts per
task. We then append the rollouts to our online training dataset
$\mathcal{D}_{\mathrm{online}}$ and continue training as described in Section
2.5.

###  3.2 Baselines

We compare our approach to the following baselines: a) MT-SAC:raw a model-free
RL algorithm extending the SAC algorithm [21] to multi-task settings. The
policies are learned from raw images. To ensure a fair comparison, we use the
same network architecture used in our method for all policy networks,
including the actors, the critics and the targets. b) MT-SAC:token extends the
raw version to use the latent states obtained from our tokenizer, as well as
the language embeddings as inputs to the actor and critic. c) MBRL-ST trains
separate single-task (ST) world models and policies for each task. d)
LIMT:nlac is based on our method, but replaces the instruction embeddings with
predefined integer task identifiers in the actor-critic networks. nlac refers
to having no language in the actor critic. e) LIMT:nl is based on our method,
but replaces the instruction embeddings with predefined integer task
identifiers in the world model as well as in the policy networks. nl refers to
having no language in the whole model. For a fair comparison, in the last two
baselines, we repeat the task identifier multiple times to ensure it has the
same dimension as the language embeddings.

###  3.3 Results

baseline | multi-task
---|---
| success rate (%)
MT-SAC:raw | 22.50
MT-SAC:token | 22.99
LIMT:nl | 24.01
LIMT:nlac | 27.15
LIMT (ours) | 52.12

Figure 3: (Left) We visualize the language embedding of our different task
instructions in 2-dimension using t-SNE [22]. The circles and triangles
represent instructions in the training and test datasets, respectively.
Language embeddings of task instructions with semantic similarity tend to
cluster well together. (Right) We compare the multi-task success rate of our
method to the studied baselines. LIMT outperforms all baselines in the multi-
task setting by a large margin. Figure 4: Our agent is able to switch between
different tasks during inference. We initialize our model with the instruction
to flip up the switch on the right side of the table. The frames of the first
row depict the agent trying to reach this goal, before the task description
changes (second row). The task is then switched to move the slider to the
left. After the agent has performed this task, we change the instruction
again. The last row depicts the agent trying to move the slider to the right.
By the end of the episode (last frame) the slider has returned to its initial
position.

Single-task performance. We compare the sample efficiency and success rate of
LIMT to the studied baselines. Figure 2 shows the success rate of the
different methods on individual tasks over training epochs. We compute the
average success rate of the policies on individual tasks at different
evaluation time steps. LIMT and its variant consistently show better sample
efficiency than the baselines. However, none of the baselines achieve a
satisfactory success rate on the move_slider_right and close_drawer tasks.
This illustrates the conflicting objectives problem common to MTRL, since
these two tasks might be conflicting with move_slider_left and open_drawer
respectively. However, one would expect that conditioning the policy on some
kind of task representation input would help alleviate this problem. One
explanation is that the embeddings for some tasks such as the move_slider
tasks can be very conflicting and non-discriminatory as shown in Figure 3.

To understand whether a model-based approach is beneficial for multi-task
learning, we first compare its performance to model-free baselines based on
MT-SAC as described in Section 3.2. The latter do not reach a similar success
rate as LIMT under the same training budget. To ensure that this comparison is
fair, in one variant of this baseline we use the tokenizer from our method as
a way to remove the complexity of also learning perception from reward only.
Both MT-SAC:token and MT-SAC:raw lag behind LIMT and its variants.
Furthermore, we aim to validate whether multi-task training is beneficial for
individual task performance in MBRL, as was previously shown for model-free
methods [4]. Hence, we compare LIMT to the MBRL-ST baseline. The latter does
not benefit from any kind of data or model sharing across tasks. For most
tasks, LIMT’s sample efficiency and success rate are substantially higher. One
exception is the open_drawer task where MBRL-ST shows better sample efficiency
than LIMT. However, such a behavior is to be expected since the single-task
policy can more easily excel at learning the one task it specializes in. In
fact, it is a positive sign for our method that this tendency is only seen in
one task. Additionally, we ablate the effect of using the instruction
embeddings from the pre-trained language model on LIMT. Replacing the
instruction embeddings with integer task identifiers in both the world model
and the actor critic (LIMT:nl) significantly deteriorates the performance in
all tasks. This validates our hypothesis on the importance of using
semantically useful task representations for bootstrapping the learning of the
dynamics model of similar tasks. When only removing the language embeddings
from the actor-critic (LIMT:nlac), we observe that the resulting agent is
competitive with the other baselines, but still lags behind the fully
language-informed version of LIMT.

Multi-task performance. In Table 3, we compare the multi-task success rate of
the different MTRL methods. We compute this metric by averaging the success
rate of each agent in all the tasks at evaluation time. LIMT achieves a
success rate higher than the model-free baselines by approximately 30% under
the same sample and update budget. Even when not using language instructions
at all in LIMT, the multi-task success rate is still higher than the model-
free baselines. Additionally, we can clearly observe the benefit of using
instruction embeddings as task representation to bootstrap the learning of the
multi-task world model and actor-critic. Not using these embeddings in the
actor-critic component decreases the success rate by 25% and not using it at
all by 28%. These results clearly illustrate the strength of our method for
multi-task policy learning.

Task Switching. In Figure 4, we illustrate the emerging capability of LIMT
agents to switch to performing new tasks during task execution. This behavior
can simply be achieved by feeding the policy a different instruction while the
agent is performing a given task. The figure illustrates a successful
switching between 3 different tasks without having to reset the agent or the
environment. We attribute this feature to the relabeling of task data for
sharing it across tasks.

##  4 Related Work

Multi-task reinforcement learning is concerned with learning one policy for
multiple tasks. Previous research on this topic highlighted its benefits, such
as bootstrapping the learning of more complex tasks [4], but also identified
some of its challenges, such as conflicting objectives [23]. To address these
challenges multiple methods have been proposed. Rusu et al. [24], Parisotto et
al. [25] proposed distilling a single multi-task policy from multiple
individually trained DQN policies. Hessel et al. [26] demonstrated the first
single agent surpassing human performance in the multi-task domain of Atari
games [27]. The proposed algorithm adapts the contribution of different tasks
on the agent’s updates in a way that reduces the bias toward specific tasks.
Kalashnikov et al. [4] presented a method for learning multiple manipulation
skills using off-policy RL and relabeling shared data across tasks. D’Eramo et
al. [28] examined the effect of sharing representation across multiple tasks
and demonstrated the benefits of that paradigm on improving the multi-task
learning and even its positive effect on individual task performance. Yang et
al. [29] proposed sharing the policy network while separately learning a
rerouting mechanism to choose which parts of the network are used for the
different tasks. Rosete-Beas et al. [30] proposed learning a latent action
plan using conditional autoencoders and learning low-level skills conditioned
on such low-level plans from an offline dataset, and a high-level policy
outputting such high-level actions.

Language in robot learning. Recent work integrated language in robotics for
different purposes. For instance, multiple methods have been proposed for
using pre-trained language models or vision-language models (VLM) as high-
level planners for robotics tasks [31, 32, 33, 34]. VLMs have also been used
as success detectors [35]. Other work explored the usage of language models as
a way to establish communication in multi-agent robotics tasks [36].
Karamcheti et al. [37] proposed an approach to learn language-informed latent
actions as a way to allow humans to influence policy actions. Driess et al.
[38] learn an embodied language model integrating sensor measurements and
grounding the resulting model with embodiment data. Multiple efforts have been
made to design language-conditioned policies for robotics [39, 40, 41, 42,
43]. Similarly, other work has leveraged language embeddings for goal-
conditioning of robot policies [44].

Inspired by previous efforts in model-free MTRL, our work leverages pretrained
language models to represent different tasks in multi-task policy learning
using a novel model-based approach.

##  5 Limitations

LIMT relies on precomputed language embeddings as a task representation on
which the policy is conditioned. As discussed in Section 3.3, these embeddings
can be conflicting for tasks with similar textual descriptions but different
dynamics. A possible way to alleviate this would be to use our sequence model
for finetuning the task embeddings via contrastive representation learning as
done in [19]. Furthermore, model-based trajectory generation only happens
during training, while at inference time LIMT samples actions from the policy
network without directly accessing the world model. While LIMT’s sample
efficiency still benefits from its world model, this limits the generalization
of our policy when encountering out-of-distribution states. In other works
[45], trajectory optimization and dynamics learning are more closely coupled.

##  6 Conclusion

We propose a method for learning language-informed multi-task visual world
models. We use a pre-trained language model to embed task instructions into a
semantically meaningful latent space and use these embeddings as task
representations. We then use these representations as inputs to the actor,
critic and dynamics models as a meaningful task discriminator. We train all
these components using environment interactions in multiple tasks. Our
experiments demonstrate the benefit of model-based training for obtaining
multi-task policies, as well as the importance of language conditioning as a
way to learn world models and policies in multi-task settings. Namely, our
method substantially outperforms the model-free multi-task baselines
demonstrating the benefit of model-based learning for learning multi-task
policies. In addition, we show that multi-task training can lead to a higher
success rate than single-task training in the model-based setting and under
the same data budget. Given the high sample complexity of multi-task policy
learning, our results provide a promising path towards more efficient learning
of such policies based on language-conditioned world models.

## References

  * Akkaya et al. [2019] I. Akkaya, M. Andrychowicz, M. Chociej, M. Litwin, B. McGrew, A. Petron, A. Paino, M. Plappert, G. Powell, R. Ribas, et al.  Solving rubik’s cube with a robot hand.  _arXiv preprint arXiv:1910.07113_ , 2019.
  * Hwangbo et al. [2019] J. Hwangbo, J. Lee, A. Dosovitskiy, D. Bellicoso, V. Tsounis, V. Koltun, and M. Hutter.  Learning agile and dynamic motor skills for legged robots.  _Science Robotics_ , 4(26):eaau5872, 2019.
  * Brunskill and Li [2013] E. Brunskill and L. Li.  Sample complexity of multi-task reinforcement learning.  In _Proceedings of the Twenty-Ninth Conference on Uncertainty in Artificial Intelligence_ , pages 122–131, 2013.
  * Kalashnikov et al. [2022] D. Kalashnikov, J. Varley, Y. Chebotar, B. Swanson, R. Jonschkowski, C. Finn, S. Levine, and K. Hausman.  Scaling up multi-task robotic reinforcement learning.  In _Conference on Robot Learning_ , pages 557–575. PMLR, 2022.
  * James et al. [2017] S. James, A. J. Davison, and E. Johns.  Transferring end-to-end visuomotor control from simulation to real world for a multi-stage task.  In S. Levine, V. Vanhoucke, and K. Goldberg, editors, _Proceedings of the 1st Annual Conference on Robot Learning_ , volume 78 of _Proceedings of Machine Learning Research_ , pages 334–343. PMLR, 13–15 Nov 2017.  URL https://proceedings.mlr.press/v78/james17a.html.
  * Peng et al. [2018] X. B. Peng, M. Andrychowicz, W. Zaremba, and P. Abbeel.  Sim-to-real transfer of robotic control with dynamics randomization.  In _2018 IEEE international conference on robotics and automation (ICRA)_ , pages 3803–3810. IEEE, 2018.
  * Aljalbout et al. [2024] E. Aljalbout, F. Frank, M. Karl, and P. van der Smagt.  On the role of the action space in robot manipulation learning and sim-to-real transfer.  _IEEE Robotics and Automation Letters_ , 2024.
  * Dulac-Arnold et al. [2019] G. Dulac-Arnold, D. Mankowitz, and T. Hester.  Challenges of real-world reinforcement learning, 2019.  URL https://openreview.net/forum?id=S1xtR52NjN.
  * Hafner et al. [2020] D. Hafner, T. Lillicrap, J. Ba, and M. Norouzi.  Dream to control: Learning behaviors by latent imagination.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=S1lOTC4tDS.
  * Mees et al. [2022] O. Mees, L. Hermann, E. Rosete-Beas, and W. Burgard.  Calvin: A benchmark for language-conditioned policy learning for long-horizon robot manipulation tasks.  _IEEE Robotics and Automation Letters_ , 7(3):7327–7334, 2022.  doi:10.1109/LRA.2022.3180108.
  * Micheli et al. [2023] V. Micheli, E. Alonso, and F. Fleuret.  Transformers are sample-efficient world models.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=vhFu1Acb0xb.
  * Van Den Oord et al. [2017] A. Van Den Oord, O. Vinyals, et al.  Neural discrete representation learning.  _Advances in neural information processing systems_ , 30, 2017.
  * Esser et al. [2021] P. Esser, R. Rombach, and B. Ommer.  Taming transformers for high-resolution image synthesis.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 12873–12883, June 2021.
  * Johnson et al. [2016] J. Johnson, A. Alahi, and L. Fei-Fei.  Perceptual losses for real-time style transfer and super-resolution.  In B. Leibe, J. Matas, N. Sebe, and M. Welling, editors, _Computer Vision – ECCV 2016_ , pages 694–711, Cham, 2016. Springer International Publishing.  ISBN 978-3-319-46475-6.
  * Larsen et al. [2016] A. B. L. Larsen, S. K. Sønderby, H. Larochelle, and O. Winther.  Autoencoding beyond pixels using a learned similarity metric.  In _Proceedings of the 33rd International Conference on International Conference on Machine Learning - Volume 48_ , ICML’16, page 1558–1566. JMLR.org, 2016.
  * Devlin et al. [2018] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Brown et al. [2020] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei.  Language models are few-shot learners, 2020.
  * Reimers and Gurevych [2019] N. Reimers and I. Gurevych.  Sentence-BERT: Sentence embeddings using Siamese BERT-networks.  In K. Inui, J. Jiang, V. Ng, and X. Wan, editors, _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_ , pages 3982–3992, Hong Kong, China, Nov. 2019. Association for Computational Linguistics.  doi:10.18653/v1/D19-1410.  URL https://aclanthology.org/D19-1410.
  * Mees et al. [2022] O. Mees, L. Hermann, and W. Burgard.  What matters in language conditioned robotic imitation learning over unstructured data.  _IEEE Robotics and Automation Letters_ , 7(4):11205–11212, 2022.  doi:10.1109/LRA.2022.3196123.
  * Radford et al. [2021] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, G. Krueger, and I. Sutskever.  Learning transferable visual models from natural language supervision.  In M. Meila and T. Zhang, editors, _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pages 8748–8763. PMLR, 18–24 Jul 2021.  URL https://proceedings.mlr.press/v139/radford21a.html.
  * Haarnoja et al. [2018] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor, 2018.
  * Van der Maaten and Hinton [2008] L. Van der Maaten and G. Hinton.  Visualizing data using t-sne.  _Journal of machine learning research_ , 9(11), 2008.
  * Liu et al. [2021] B. Liu, X. Liu, X. Jin, P. Stone, and Q. Liu.  Conflict-averse gradient descent for multi-task learning.  _Advances in Neural Information Processing Systems_ , 34:18878–18890, 2021.
  * Rusu et al. [2016] A. A. Rusu, S. G. Colmenarejo, C. Gulcehre, G. Desjardins, J. Kirkpatrick, R. Pascanu, V. Mnih, K. Kavukcuoglu, and R. Hadsell.  Policy distillation.  In _International Conference on Learning Representations (ICLR)_ , 2016.  URL https://dblp.org/rec/journals/corr/RusuCGDKPMKH15.
  * Parisotto et al. [2016] E. Parisotto, L. J. Ba, and R. Salakhutdinov.  Actor-mimic: Deep multitask and transfer reinforcement learning.  In _International Conference on Learning Representations (ICLR)_ , 2016.  URL https://dblp.org/rec/journals/corr/ParisottoBS15.
  * Hessel et al. [2019] M. Hessel, H. Soyer, L. Espeholt, W. Czarnecki, S. Schmitt, and H. Van Hasselt.  Multi-task deep reinforcement learning with popart.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, pages 3796–3803, 2019.
  * Bellemare et al. [2013] M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling.  The arcade learning environment: An evaluation platform for general agents.  _Journal of Artificial Intelligence Research_ , 47:253–279, 2013.
  * D’Eramo et al. [2020] C. D’Eramo, D. Tateo, A. Bonarini, M. Restelli, and J. Peters.  Sharing knowledge in multi-task deep reinforcement learning.  _International Conference on Learning Representations_ , 2020.
  * Yang et al. [2020] R. Yang, H. Xu, Y. WU, and X. Wang.  Multi-task reinforcement learning with soft modularization.  In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 4767–4777. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper_files/paper/2020/file/32cfdce9631d8c7906e8e9d6e68b514b-Paper.pdf.
  * Rosete-Beas et al. [2023] E. Rosete-Beas, O. Mees, G. Kalweit, J. Boedecker, and W. Burgard.  Latent plans for task-agnostic offline reinforcement learning.  In K. Liu, D. Kulic, and J. Ichnowski, editors, _Proceedings of The 6th Conference on Robot Learning_ , volume 205 of _Proceedings of Machine Learning Research_ , pages 1838–1849. PMLR, 14–18 Dec 2023.  URL https://proceedings.mlr.press/v205/rosete-beas23a.html.
  * Ahn et al. [2022] M. Ahn, A. Brohan, N. Brown, Y. Chebotar, O. Cortes, B. David, C. Finn, C. Fu, K. Gopalakrishnan, K. Hausman, et al.  Do as i can, not as i say: Grounding language in robotic affordances.  _arXiv preprint arXiv:2204.01691_ , 2022.
  * Rana et al. [2023] K. Rana, J. Haviland, S. Garg, J. Abou-Chakra, I. Reid, and N. Suenderhauf.  Sayplan: Grounding large language models using 3d scene graphs for scalable robot task planning.  In J. Tan, M. Toussaint, and K. Darvish, editors, _Proceedings of The 7th Conference on Robot Learning_ , volume 229 of _Proceedings of Machine Learning Research_ , pages 23–72. PMLR, 06–09 Nov 2023.  URL https://proceedings.mlr.press/v229/rana23a.html.
  * Wake et al. [2023] N. Wake, A. Kanehira, K. Sasabuchi, J. Takamatsu, and K. Ikeuchi.  Chatgpt empowered long-step robot control in various environments: A case application.  _IEEE Access_ , 2023.
  * Singh et al. [2023] I. Singh, V. Blukis, A. Mousavian, A. Goyal, D. Xu, J. Tremblay, D. Fox, J. Thomason, and A. Garg.  Progprompt: Generating situated robot task plans using large language models.  In _2023 IEEE International Conference on Robotics and Automation (ICRA)_ , pages 11523–11530. IEEE, 2023.
  * Du et al. [2023] Y. Du, K. Konyushkova, M. Denil, A. Raju, J. Landon, F. Hill, N. de Freitas, and S. Cabi.  Vision-language models as success detectors.  In _Conference on Lifelong Learning Agents_ , pages 120–136. PMLR, 2023.
  * Mandi et al. [2023] Z. Mandi, S. Jain, and S. Song.  Roco: Dialectic multi-robot collaboration with large language models.  _arXiv preprint arXiv:2307.04738_ , 2023.
  * Karamcheti et al. [2022] S. Karamcheti, M. Srivastava, P. Liang, and D. Sadigh.  Lila: Language-informed latent actions.  In _Conference on Robot Learning_ , pages 1379–1390. PMLR, 2022.
  * Driess et al. [2023] D. Driess, F. Xia, M. S. Sajjadi, C. Lynch, A. Chowdhery, B. Ichter, A. Wahid, J. Tompson, Q. Vuong, T. Yu, et al.  Palm-e: an embodied multimodal language model.  In _Proceedings of the 40th International Conference on Machine Learning_ , pages 8469–8488, 2023.
  * Mei et al. [2016] H. Mei, M. Bansal, and M. Walter.  Listen, attend, and walk: Neural mapping of navigational instructions to action sequences.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 30, 2016.
  * Misra et al. [2017] D. Misra, J. Langford, and Y. Artzi.  Mapping instructions and visual observations to actions with reinforcement learning.  In _Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing_ , pages 1004–1015, 2017.
  * Jiang et al. [2019] Y. Jiang, S. S. Gu, K. P. Murphy, and C. Finn.  Language as an abstraction for hierarchical deep reinforcement learning.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Stepputtis et al. [2020] S. Stepputtis, J. Campbell, M. Phielipp, S. Lee, C. Baral, and H. Ben Amor.  Language-conditioned imitation learning for robot manipulation tasks.  _Advances in Neural Information Processing Systems_ , 33:13139–13150, 2020.
  * Zitkovich et al. [2023] B. Zitkovich, T. Yu, S. Xu, P. Xu, T. Xiao, F. Xia, J. Wu, P. Wohlhart, S. Welker, A. Wahid, et al.  Rt-2: Vision-language-action models transfer web knowledge to robotic control.  In _Conference on Robot Learning_ , pages 2165–2183. PMLR, 2023.
  * Shah et al. [2023] D. Shah, B. Osiński, S. Levine, et al.  Lm-nav: Robotic navigation with large pre-trained models of language, vision, and action.  In _Conference on robot learning_ , pages 492–504. PMLR, 2023.
  * Janner et al. [2022] M. Janner, Y. Du, J. B. Tenenbaum, and S. Levine.  Planning with diffusion for flexible behavior synthesis, 2022.
  * Radford et al. [2019] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9, 2019.
  * Paszke et al. [2019] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala.  _PyTorch: An Imperative Style, High-Performance Deep Learning Library_ , page 8024–8035.  Curran Associates, Inc., 2019.  URL http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf.
  * Simonyan and Zisserman [2014] K. Simonyan and A. Zisserman.  Very deep convolutional networks for large-scale image recognition.  _CoRR_ , abs/1409.1556, 2014.  URL https://api.semanticscholar.org/CorpusID:14124313.

##  Appendix A Models and Hyperparameters

###  A.1 Tokenizer

The tokenizer is a discrete autoencoder that converts image-proprioception
pairs into tokenized representations, similar to the one proposed in [11]. We
downsample the images to a size of $64\times 64$ before feeding them to the
network. The tokenizer converts observations into $K=K_{x}+K_{\theta}$ tokens
of dimension $d_{\mathrm{enc}}$ using two separate codebooks for images and
proprioception, $C_{x}$ and $C_{\theta}$ respectively, of vocabulary size $N$.
The hyperparameters of the tokenizer are listed in table 1.

Table 1: Tokenizer hyperparameters Hyperparameter | Value
---|---
Vocabulary size ($N$) | 512
Image tokens ($K_{x}$) | 16
Proprioception tokens ($K_{\theta}$) | 1
Total tokens ($K$) | 17
Embedding dimension ($d_{enc}$) | 512
Frame dimensions ($H,W$) | $64\times 64$
Encoder/Decoder Layers | 5
Residual blocks per layer | 2
Convolution channels | 64
Self-attention layers at resolution | 8, 16
Learning rate ($\alpha^{\xi}$) | 1e-4

###  A.2 World Model

Our world model is a transformer based on the implementation of IRIS [11]. It
accepts a sequence of $H(K+2)$ tensors, including $HK$ tensors of tokenized
observations and $2H$ tensors containing the actions and embedded instructions
$a_{t},l$ for each timestep $0\leq t<H$. Using an embedding table of size
$N\times D$ for the tokens and linear projections for the actions and
instructions, we obtain a $H(K+2)\times D$ tensor which is then passed through
$L$ GPT2-like [46] transformer blocks. To predict rewards, episode ends, and
observation embeddings, 3 MLP heads follow the transformer blocks. The world
model’s hyperparameters are listed in table 2.

Table 2: World Model hyperparameters Hyperparameter | Value
---|---
Horizon length ($H$) | 8
Embedding dimension ($D$) | 512
Layers ($L$) | 10
Attention heads | 4
Embedding dropout | 0.1
Residual Dropout | 0.1
Attention Dropout | 0.1
Weight decay | 1e-4
Learning rate | 2e-5
MLP head layers | 3
Reward loss factor ($\rho$) | 10

###  A.3 Actor Critic

We implement both the actor and critic networks described in 2.5 as $n$-layer
MLPs with skip connections of stride 2. At training time, during imagination
rollouts, their input is the tokenized representation predicted by the world
model, as well as the predicted raw proprioception vectors $\hat{\theta}$. As
their dimension is quite small compared to the rest of the inputs
(7-dimensional pose) we repeat $\hat{\theta}$ $b$ times before passing it to
the MLPs. At inference time, the policy networks receive the tokenized
environment observations along with the raw proprioception information as
described above. To stabilize training, we maintain a target critic for
computing the lambda returns $V_{\lambda}$, which we periodically update with
the weights of our training critic. Table 3 lists the hyperparameters for both
networks.

Table 3: Hyperparameters for both the actor and critic MLPs Hyperparameter | Value
---|---
Hidden layers ($n$) | 9
Neurons per layer | 2048
Lambda ($\lambda$) | 0.9
Gamma ($\gamma$) | 0.65
Learning rate ($\alpha^{\phi},\alpha^{\psi}$) | 4e-5
Activation | ELU
Critic update interval (training steps) | 200
Proprioception repetitions ($b$) | 10
Entropy weight ($\eta$) | 1e-4

##  Appendix B Computational Resources

We implement LIMT and all baseline models in Python using PyTorch 2.0.1 [47].
The environments of CALVIN use PyBullet for physics simulation. Our
experiments run on a GPU cluster managed by the ClearML platform††clear.ml,
utilizing different GPU models, including NVIDIA A100, NVIDIA V100, and NVIDIA
RTX.

##  Appendix C Training details

###  C.1 Policy learning

The critic network $v_{\psi}$ regresses the value estimates $V_{\lambda}$ for
a horizon length $H$

$V^{\lambda}_{t}:=r_{t}+\gamma_{t}\begin{cases}(1-\lambda)v_{\psi}(\hat{s}_{t+1})+\lambda
V^{\lambda}_{t+1},&t<H\\\ v_{\psi}(\hat{s}_{H}),&t=H.\end{cases}$ (9)

The loss function of the actor network with parameter $\phi$ is

$L_{\phi}=-\sum_{t=\tau}^{T}V^{\lambda}_{t}-\eta\mathcal{H}(\pi(a_{t}|w,l,\theta)),$
(10)

where the last term is an entropy objective to encourage exploration. $w$ is
the tokenized observation and $\theta$ the end effector’s pose. The critic’s
objective is given by

$\min_{\psi}\mathbb{E}_{q_{\theta},a_{\phi}}\left[\sum_{\tau=t}^{t+H}\lVert
v_{\psi}(s_{\tau})-V_{\lambda}(s_{\tau})\rVert^{2}\right].$ (11)

###  C.2 Perceptual Loss

The perceptual loss introduced in Section 2.5 is similar to the one proposed
by [14]. Given a pre-trained VGG-16 CNN [48], we select a subset of layers
$M$. For each layer $j\in M$, let $\phi_{j}(x)$ be the activation of the
$j$-th layer, in our case a feature map with dimension $C_{j}\times
H_{j}\times W_{j}$. We compute the loss between a ground-truth image $x$ and a
reconstructed image $x^{\prime}$ as

$L_{\mathrm{perceptual}}(x,x^{\prime})=\sum_{j\in
M}\frac{1}{C_{j}H_{j}W_{j}}A_{j}||\phi_{j}(x)-\phi_{j}(x^{\prime})||^{2},$
(12)

where $A_{j}$ are learned affine transformations, implemented as 1x1
convolutions.

###  C.3 Online training

Once the training of our actor-critic begins, we start performing a total of
$n_{rollout}$ policy rollouts of $T$ timesteps at each epoch. We perform the
same amount of $n_{\mathrm{rollout}}/6$ rollouts for each task. We append
these online episodes to our dataset $\mathcal{D}_{\mathrm{online}}$. We limit
the size of $\mathcal{D}_{\mathrm{online}}$ to a maximum of $n_{max}$
episodes. At every epoch we compute the online sampling ratio
$p_{\mathrm{online}}$, described in 2.5 as

$p_{\mathrm{online}}=p_{\mathrm{max}}\frac{|\mathcal{D}_{\mathrm{online}}|}{n_{\mathrm{max}}},$
(13)

where we set $p_{max}$ as a maximum ratio. When sampling from $D_{online}$, as
done in [11], we prioritize later episodes: we divide $D_{online}$ into 4
quarters and we sample 50% of our episodes from the last quarter, and 25% from
the 3rd quarter. The rest of the episodes is uniformly sampled from the first
half.

All training hyperparameters described above, as well as in section 2.5 can be
found under table 4. Our overall training process is further detailed in
algorithm 1.

Table 4: General Training Hyperparameters. The (unnormalized) task weights $w$ correspond to the following tasks in order: open_drawer, close_drawer, move_slider_left, move_slider_right, turn_on_lightbulb, turn_on_led. Hyperparameter | Value
---|---
Batch size tokenizer | 64
Batch size world model | 128
Batch size actor-critic | 120
Training steps per epoch | 200
Optimizer | Adam
Adam $\beta_{1}$ | 0.9
Adam $\beta_{2}$ | 0.999
Max gradient norm | 10
Maximum online sampling ratio ($p_{max}$) | 0.5
Task weights $w$ | [1, 1, 1, 1, 2/3, 2/3]
$N_{t}$ | 200
$N_{w}$ | 100
$N_{f}$ | 50
Relabeling probability ($1-p$) | 0.15
Number of rollouts ($n_{rollout}$) | 120
Timesteps per episode ($T$) | 50
$n_{max}$ | 5000
Algorithm 1 LIMT

1:Given:

2:•pre-trained LLM $\mathcal{L}$

3:•offline datasets $\mathcal{D}_{\mathrm{offline}}$,
$\mathcal{D}_{\mathrm{filtered}}$ of episode trajectories and directives

4:•initialized discrete autoencoder $A_{\xi}$, world model $G$, actor $\phi$,
critic $\psi$, target critic $\psi_{\mathrm{target}}$

5:Encode language directives in $\mathcal{D}_{\mathrm{offline}}$ using
$\mathcal{L}$

6:$\mathcal{D}\leftarrow\mathcal{D}_{\mathrm{offline}}$

7://Training tokenizer

8:$\mathcal{D}_{\mathrm{online}}\leftarrow\\{\\}$

9:while not converged do

10: for i=1…training_steps do

11: sample $B$ states $(x^{j},\theta^{j})_{j=1}^{B}\sim\mathcal{D}$

12: $\xi\leftarrow\xi-\alpha^{\xi}\nabla L_{A}(\theta,x,\theta)$

13: end for

14: if $i>N_{t}$ then

15://Training world model

16: for i=1…training_steps do

17: sample $B$ trajectories
${\\{(x_{t}^{j},\theta_{t}^{j},a_{t}^{j},r_{t}^{j},l^{j})_{t=\tau}^{\tau+H}\\}}_{j=1}^{B}\sim\mathcal{D}$
of $H$ timesteps

18: Obtain tokens using tokenizer:
${\\{(w_{t}^{j},a_{t}^{j},l^{j})_{t=\tau}^{\tau+H}\\}}_{j=1}^{B}$

19: predict tokens, rewards and termination $\hat{w}_{t+1}^{j}$,
$\hat{r_{t}}^{j}$, $\hat{d_{t}}^{j}$ $\triangleright$ See eq. (7)

20: update $G$ using eq. (8)

21: end for

22: end if

23: if $i>N_{t}+N_{w}$ then

24: $\mathcal{D}\leftarrow\mathcal{D}_{\mathrm{filtered}}$

25: end if

26: if $i>N_{t}+N_{w}+N_{f}$ then

27://Training actor critic

28: for i=1…training_steps do

29: sample $B$ initial states and instruction
$(x^{j},\theta^{j},l^{j})_{j=1}^{B(1-p_{\mathrm{online}})}\sim\mathcal{D}$,
$(x^{j},p^{j},l^{j})_{j=1}^{Bp_{\mathrm{online}}}\sim\mathcal{D}_{\mathrm{online}}$

30: relabel language directives $l^{j}$ with probability $1-p$

31: Obtain tokens $\\{w_{0}^{j}\\}^{B}_{j=1}$ using $A_{\xi}$

32: Imagine latent trajectories $(w_{t+1},a_{t},r_{t},d_{t})_{t=1}^{H-1}$ for
each $w^{j}$ using $G$ and $\phi$

33: Compute $\lambda$-returns $\\{V^{\lambda}_{t}\\}_{t=1}^{H}$ using
$\psi_{\mathrm{target}}$

34:
$\phi\leftarrow\phi+\alpha^{\phi}\nabla\sum_{t=1}^{H}(V^{\lambda}_{t}+\eta\mathcal{H}(\pi(a_{t}|w,l,\theta)))$

35:
$\psi\leftarrow\psi-\alpha^{\psi}\nabla\sum_{t=1}^{H}||V^{\lambda}_{t}-v_{\psi}(w_{t})||^{2}$

36: end for

37: $\psi_{\mathrm{target}}\leftarrow\psi$

38://Online rollouts

39: perform $n_{\mathrm{rollout}}$ rollouts of $T$-timestep trajectories
$\\{\tau^{n}\\}_{n=1}^{n_{\mathrm{rollout}}}$ using policy $\pi_{\phi}$

40:
$\mathcal{D}_{\mathrm{online}}\leftarrow\mathcal{D}_{\mathrm{online}}\cup\\{\tau^{n}\\}_{n=1}^{n_{\mathrm{rollout}}}$

41: end if

42:end while

###  C.4 Reward Functions

The reward functions we specify for our tasks of interest are divided into two
classes:

The first one consists of tasks the completion of which can adequately be
described by a boolean variable. For example, turning on/off a lightbulb. In
these types of tasks, the robot arm has to reach and manipulate an object (e.g
a button) to bring about some change in the environment (e.g lightbulb on). A
general formula of this type of reward is given in equation 14.

$R_{b}(\theta,g,s,l)=1-\|(\theta-g)\odot
f_{s}\|_{2}+10\mathbb{I}[\mathrm{success}(s,l)]$ (14)

Here, $\theta$ contains is the proprioception vector and $g$ is the desired
pose of the end effector to reach an object of interest, such a buttons or a
switch. The first term is therefore a distance-based reward. This is employed
to guide the agent to the vicinity of the object and avoid relying on sparse
rewards, which can slow down the learning process. We multiply the difference
$(\theta-g)$ by a scaling vector $f_{s}$ to prioritize some proprioception
dimensions over others. In our case, end effector orientation is discounted
relative to end effector position. The intuition behind this is that, for
completing the tasks, it is more important for the EE to reach the specified
position of the object of interest, although the orientation also matters to a
certain extent. The reason for the latter is that, by constraining the range
of possible orientations, we make it easier for the agent to learn suitable
actions, while excluding some of the orientations that would make the
manipulation of the object significantly harder.

The second term is an indicator variable indicating completion of the task,
depending on the environment’s state $s\in S$ and the instruction $l$. We
multiply it by a constant $\beta_{b}\in\mathbb{R}$, in our case 10, to amplify
the learning signal for successes.

The second class of reward functions are designed for tasks the success of
which can best be described as a continuous, time-dependent variable. For
example, to define success for the task of opening a drawer, we need
information about its (continuous) position in the previous time steps. In
these cases, the reward function is defined as follows:

$R_{c}(p_{t},g,s_{t},s_{t-1},l,s_{g})=1-||(p_{t}-g)\odot f_{s}||_{2}+\beta
sign(s_{g}(l)-s_{t-1})\cdot(s_{t}-s_{t-1})$ (15)

where $\theta_{t}$ and $g$ are the EE’s actual and desired pose,
$s_{t},s_{t-1}$ are environment states for timesteps $t,t-1$ and $s_{g}(l)$ is
the desired environment state, determined by the instruction $l$. In the
example of opening a drawer ,$s_{g}$ would corrrespond to a fully opened
drawer. As above, the first term is a dense distance-based reward. The second
term measures progress towards completion of the task during the last timestep
and is weighed by a scalar $\beta_{c}\in\mathbb{R}$ to amplify the learning
signal. As the difference $(s_{t}-s_{t-1})$, can be quite small, we use
$\beta_{c}=50$.

Table 5: Reward Hyperparameters. The entries of the scaling factor $f_{s}$ correspond to the position, orientation and gripper width of the end effector in that order. Hyperparameter | Value
---|---
Reward scalar $\beta_{b}$ | 10
Reward scalar $\beta_{c}$ | 50
Scaling factor $f_{s}$ | [1, 1, 1, 0.5, 0.5, 0.5, 0]

##  Appendix D Imagined Trajectories

To qualitatively assess our world model, we examine some imagined trajectories
generated during policy training. Figure 5 illustrates two trajectories of
length 8 for the tasks turn_on_lightbulb (top) and close_drawer (bottom). As
our agent is trained in latent imagination and not directly on image data, we
map the latent outputs to images using the decoder $D$ of our tokenizer for
visualization purposes. As can be seen in the figure, our world model
successfully captures task-relevant visual features, like the yellow lightbulb
lighting up and the drawer transitioning from an opened to a closed position.

Figure 5: Reconstructions of imagined trajectories for the tasks
turn_on_lightbulb (top) and close_drawer (bottom). The first frame in each row
is initialized from the real environment.