Discrete Codebook World Models for Continuous Control
Authors: Aidan Scannell
Mohammadreza Nakhaeinezhadfard
Kalle Kujanpää
Yi Zhao
Kevin Sebastian Luck
Arno Solin
Joni Pajarinen
Venue: ICLR 2025 Poster (2025)
Keywords: reinforcement learning
world model
representation learning
self-supervised learning
model-based reinforcement learning
continuous control
Primary area: reinforcement learning
TL;DR: World models with discrete codebook encodings are effective for continuous control

Abstract: In reinforcement learning (RL), world models serve as internal simulators, enabling agents to predict environment dynamics and future outcomes in order to make informed decisions. While previous approaches leveraging discrete latent spaces, such as DreamerV3, have demonstrated strong performance in discrete action settings and visual control tasks, their comparative performance in state-based continuous control remains underexplored. In contrast, methods with continuous latent spaces, such as TD-MPC2, have shown notable success in state-based continuous control benchmarks. In this paper, we demonstrate that modeling discrete latent states has benefits over continuous latent states and that discrete codebook encodings are more effective representations for continuous control, compared to alternative encodings, such as one-hot and label-based encodings. Based on these insights, we introduce DCWM: Discrete Codebook World Model, a self-supervised world model with a discrete and stochastic latent space, where latent states are codes from a codebook. We combine DCWM with decision-time planning to get our model-based RL algorithm, named DC-MPC: Discrete Codebook Model Predictive Control, which performs competitively against recent state-of-the-art algorithms, including TD-MPC2 and DreamerV3, on continuous control benchmarks.

Published as a conference paper at ICLR 2025
DISCRETE CODEBOOK WORLD MODELS FOR CONTINUOUS CONTROL
Aidan Scannell
Mohammadreza Nakhaei'
Kalle Kujanpää* Aalto University
University of Edinburgh a dann scannel @ed ac.
Aalto University
Yi Zhao
Kevin Sebastian Luck Vrije Universiteit Amsterdam
Arno Solin Aalto University
Joni Pajarinen Aalto University
Aalto University
ABSTRACT
In reinforcement learning (RL), world models as internal simulators, enabling agents to predict environment dynamics and future outcomes in order to make informed decisions. While previous approaches leveraging discrete latent spaces, such as DreamerV3, have demonstrated strong performance in discrete action settings and visual control tasks, their comparative performance in state-based continuous control remains underexplored. In contrast, methods with continuous latent spaces, such as TD-MPC2, have shown notable success in state-based con- tinuous control benchmarks. In this paper, we demonstrate that modeling discrete latent states has benefits over continuous latent states and that discrete codebook encodings more effective representations for continuous control, compared alternative encodings, such as one-hot and label-based encodings. Based on these insights, we introduce DCWM: Discrete Codehook World Model, self-supervised world model with a discrete and stochastic latent space, where latent states are codes from a codebook. We combine DCWM with decision-time planning get our model-based RL algorithm, named DC-MPC: Discrete Codebook Model Predictive Control, which performs competitively against recent state-of-the-ee algorithms, including TD-MPC2 and DreamerV3, on continuous control benchmarks. See our project website www. a dadaaaaaaann com/dcmpo.
INTRODUCTION
In model-based reinforcement learning (RL), world models (Ha & Schmidhuber, 2018) have been introduced in order to simulate or predict the environment's dynamics in data-driven way. An agent equipped with world model can make predictions about its environment by "simulating" possible actions within the model and "imagining" the outcomes. This equips the agent with the ability to plan and anticipate outcomes given (leamed) reward function, and the additional ability to envision transitions and outcomes before taking them in the real world can in turn improve sample efficiency.
One of the state-of-the- art world models, DreamerV2/V3 (Hafner et al., 2022: 2023) achieves strong performance in a wide variety of tasks, by "imagining" sequences of future states within a world model and using them to improve their policies. Interestingly, Dreamer 22/V3 introduced discrete latent space, in the form of a one-hot encoding, which offered significant benefits over its predecessor, DreamerV1 (Hafner al., 2019a). This suggests that discrete latent spaces may have benefits over continuous latent spaces. It could be from the discrete latent space helping avoid compounding emrors over multi-step time horizons or enabling policy and value learning harness the benefits of discrete variable processing for efficiency and interoperability. In the context of generative modeling. discrete codebooks have been the heart of many successful approaches (Chang et al., 2023; Esser et al., 2021; Ramesh et al., 2021). However, in the context of continuous control, TD-MPC2 (Hansen et al., 2023) uses a continuous latent space and significantly outperforms DreamerV3. Whilst there multiple differences between TD-MPC2 and Dreamer 3, in this paper, we are specifically interested in exploring if discrete latent spaces can offer benefits for continuous control.
Work done while Aalto University
"Equal contribution

---

Published as a conference paper at ICLR 2025
Recently, Farebrother et al. (2024) showed that training value functions with classification may have benefits over training with regression. The benefits may arise because (i) classification considers uncertainty during training (via the cross-entropy loss), (ii) the categorical distribution is multi-modal so it can consider multiple modes during training, or (iti) learning in discrete spaces is more efficient. In the context of world models, it is natural ask, what benefits are obtained by (i) using discrete vs continuous latent spaces and (ii) modeling deterministic vs stochastic transition dynamics. Further to this, when considering stochastic latent transition dynamics, what the effect of modeling with (i) unimodal distributions (e.g. Gaussian in continuous latent spaces) vs (ii) multimodal distributions (e.g. categorical in discrete latent spaces). In this paper, we explore these ideas in the context of world models for model-based RL, i.e. does learning discrete latent space using classification have benefits over learning a continuous latent space using regression.
Contributions The main contributions are as follows:
(CI) In the context of continuous control, we show that learning discrete latent spaces with classifi- cation does have benefits over learning continuous latent spaces with regression.
(C2) We show that formulating a discrete latent state using codebook encodings has benefits over alternatives, such one-hot (like DreamerV2/V3) and label encodings.
(C3) Based on our insights, we introduce Discrete Codebook World Model (DCWM): a world model with discrete latent space where each latent state is discrete code from codebook. It obtains strong performance in the difficult locomotion tasks from DeepMind Control suite (Tassa et al., 2018) and manipulation tasks from Meta-World (Yu et al., 2019).
2 RELATED WORK
In this section, we recap world models in the context of model- based RL. We introduce two competing methods for learning latent spaces (i) those using observation reconstruction and (ii) those using latent state temporal consistency objectives. We then compare methods that learn continuous latent spaces using regression and those that learn discrete latent spaces using classification.
World models Model-based RL is often said to be more sample-efficient than model-free methods. This is because it learns a model in which it can reason about the world, instead of simply trying to learn a policy or a value function to maximize the return (Ha & Schmidhuber, 2018). The world model can be used for planning (Allen & Koomen, 1983; Basye et al., 1992). A prominent idea has been to optimize the evidence lower bound of observation and reward sequences to learn world models that operate on the latent space of learned Variational Autoencoder (VAF, Kingma & Welling (2014); Igl et al. (2018) (0. These models rely on maximizing the conditional observation likelihood p(o, zz). i.e. the reconstruction objective. The latent space of the model can then be used for both policy learning in the imagination of the world model, known as offline planning, e.g. Dreamer (Hafner et al., 20)9,), or for decision-time planning (Rubinstein, 1997; Hafner et al., 2019b; Schrittwieser et al., 2020).
Latent-state consistency Using the reconstruction loss for learning latent state representations is unreliable (Lutter et al., 2021) and can have a detrimental effect on the performance of model-based methods in various benchmarks (Kostrikov et al., 2020; Yarats et al., 2021a). To this end, TD-MPC (Hansen et al., 2022) and its successor, TD-MPC2 (Hansen et al., 2023), use a consistency loss to learn representations for planning with Model Predictive Path Integral (MPPI) control together with reward and value functions learned through temporal difference methods (Williams et al., 2015). Note that many prior works learn latent state representations using variants of a self-supervised latent-state consistency objective (Schwarzer et al., 2020; Wang et al., 2022; Ghugare et al., 2022; LeCun; Georgiev et al., 2024: Scannell et al., 2024b; Zhao et al., 2023; Scannell et al., 2024a). Given the success of learning representations without observation reconstruction in continuous control tasks, we predominantly focus on this class of methods, i.e. methods that use latent-state consistency losses.
Discrete latent spaces DreamerV1 (Hafner et al., 2019a), DreamerV2 (Hafner et al., 2022). and D3eamerV3 (Hafner et al., 2023), are world model methods which learn policies using imagined transitions from their world models. They utilize observation reconstruction when learning their world models and perform well across a wide variety of tasks. However, they are significantly outperformed by TD-MPC2 in continuous control tasks, which does not reconstruct observations. Of particular inter- est in this paper, is that DreamerV2/V3 introduced a discrete latent space, in the form of a one-hot en- coding, and trained it with classification objective, significantly improving performance. In contrast,
2

---

Published as a conference paper at ICLR 2025
O
Oc-1
0
Encoder
Encoder
Eacoder
I-QQ
Codebook
Codebook
a2
index
Cross-entropy
Cross-entropy
Latent code
Latent code G+4l
Lasent code
Dynamics
ST Gumbel-- softmax
Dytamics
ST Gumbel-
C1
Polers cc, ai)
(Rees Geal deal)
softmax
sampling
sampling
TL
Ti+1
Figure World model training DCWM is a world model with discrete latent space where each latent state is discrete code e (Mиيоl) from codebook C. Observations o are first mapped through the encoder and then quantized latent transition dynamics pa(c' a) classifier such that it captures a potentially multimodal distribution over the next state c' given the previous state c and action a. During training, multi-step predictions are made using straigh-tthouus backpropagate through time to the encoder. Given this discrete formulation, we train the latent space using classification objective, iee. cross-entropy loss. Making the latent representation stochastic and discrete with a codebook contributes to the very high sample efficiency of DC-MPC.
into one of the discrete codes. We model probabilistic
(ST) Gumbel-softmax sampling such that gradients
TD-MPC2 learns continuous latent space with mean squared error regression. In this paper, we interested learning discrete latent spaces with classification, however, contrast DreamerV2/V3, seek to avoid observation reconstruction due to its poor performance in continuous control (see Fig. 20) - and instead learn the latent space using a self-supervised latent-state consistency loss.
3 PRELIMINARIES
In this section, we recap different types of discrete encodings and compare their pros and cons. First, let assume have three discrete categories: A, B, and C.
One-hot encoding Given categories A, B, and C, a one-hot encoding would take the form
e(A) = [1, 0, e(B) j0, 1, 1). and e(C) = [0, 0, 1] respectively.
Label encoding Given categories A, B, and C, label encoding would result in e(A) = 1, respectively.
e(B) 2, and e(C)
Codebook encoding Given categories A, B, and C. codebook might encode them as e(A) 0.5, 0.5), e eBB) [o, 0]. and e(C) [0.5, 0.5] respectively.
Ordinal relationships If we have an ordinal relationship A < B < C, label and codebook encodings can ensure (c(A) c(B)] [e(A) e(C'), where e(-) is the encoding function. In this case, the global ordering is preserved along both dimensions of the codebook. It is worth noting that codebook encodings are flexible enough model ordinal relationships in multiple dimensions. For example, the following code vectors exhibit opposite ordering along their two dimensions e(E) [0.5, -0.5), e(F) [o, 0), (G) = [-0.5, 0.5), which adds a level of modeling flexibility. One-hot encoding, however, results in le(A) - -(B) = le(A) - e(C') = 22 for all distinct pairs, climinating any notion of ordering. Whilst this may be beneficial in some scenarios, e.g., when modeling distinet categories like fruits, it means that they cannot capture the inherent ordering in continuous data.
Sparsity and dimensionality Another downside of one-hot encodings is that they create sparse data (i.e., data with many zero values), which can have negative impact on neural network training. In con- trast, label and codebook encodings create dense data (i.e. many non-zero values). Finally, is worth
3

---

Published as a conference paper at ICLR 2025
noting that one-hot encodings have high dimensionality, especially when there many categories. This makes them memory memory-intensive intensive and slow to train when using a large number categories.
In this work. show that discrete codebook encodings resulting from quantization (Mentzer et al., 2024) offer benefits over both one-hot and label encodings when learning discrete latent spaces for continuous control. This is because they preserve ordinal relationships in multiple dimensions whilst being simpler, much lower-dimensionll and having less memory requirements.
4 METHOD
In this section, we detail our method, named Discrete Codebook Model Predictive Control (DC-MPC), which is model-based RL algorithm which (i) learns world model with discrete latent space, named Discrete Codebook World Model (DCWM), and then, (ii) performs decision-time planning with MPPI. The paper's main contribution is formulating discrete latent space using quantization such that latent states are codes from codebook. This allows us to train the latent representation using classification, in sulf-supervised manner. See Fig. 1 for overview of DCWM, Alg. for details of world model training and Alg. 2 for details on the MPPI planning procedure.
We consider Markov Decision Processes (MDPs, Bellman (1957)) M
A, R, 7), where agents receive observations o2 E at time step f, perform actions a₂ E A, and then obtain the next ob- servation 01+1 P(: oz, az) and reward rą R(o,, a.). The discount factor is denoted [0, 1).
4.1 WORLD MODEL
Learning world models with discrete latent spaces (e.g. DreamerV2) has proven powerful in a wide variety of domains. However, these approaches generally perform poorly in continuous control tasks when compared to algorithms like TD-MPC2 and TCRL (Zhao et al., 2023), which use continuous latent spaces. Rather than representing discrete latent space using a one-hot encoding, as was done in DreamerV2, DC-MPC aims to construct more expressive representation which is effective for continuous control. More specifically, DC-MPC represents discrete latent states as codes from a discrete codebook, obtained via finite scalar quantization (FSQ, Mentzer et al. (2024)). The world model can subsequently benefit from the advantages of discrete representations, e.g. efficiency and training with classification, whilst performing well in continuous control tasks.
Components DC-MPC has six main components:
Encoder:
a cœ(o) E I 4|xd
(1)
Latent quantization:
f(x) EC
(2)
Dynamics:
è' ~ Categorical PPcc)) with p: = -Polc' = c(0) a)
(3)
Reward:
r = Re(c,a) a) E R
(4)
Value:
ac(c, e ERN,
(5)
Policy prior:
a (c)(c)
The encoder c-(-) first maps observations to continuous latent vectors E R where the number of channels and the latent dimension d are hyperparameters. This continuous latent vector æ then quantized ((.) into one of the discrete latent codes e E C from the (fixed) codebook C, using finite scalar quantization (FSQ, Mentzer et al. (2024)). As we have a discrete latent space, we formulate the transition dynamics to model the distribution over the next latent state c' given the previous latent state c and action a. That is, we model stochastic transition dynamics in the latent space. We denote the probability of the next latent state è' taking the value of the 7th code c(1) as p: (s(c' c(0) c, a). This results in the next latent state following categorical distribution c' Categorical (pi....,
We use a standard classification setup, where we use an MLP to predict the logits = (4..... .4c = d (c. a) E RlCl Note that logits are the raw outputs from the final layer of the neural network (NN), which represent the unnormalized probabilities of the next latent state e' taking the value of each discrete code in the codebook C. The logit for the 1th code is
given by 4 [c, a)); R. We then apply softmax to obtain the probabilities [C] of the next latent state taking each discrete code in the codebook C, i.e., p; = softmax, (l). DC-MPC utilizes the discrete codes e as its latent state for future predictions and decision-making.
4

---

Published as a conference paper at ICLR 2025
Quantized latent space DCWM uses discretized latent space where world states encoded as discrete codes from a codebook C. We use latent quantization to enforce data compression and encourage organization (Hsu et al., 2023). However, implement this using finite scalar quantization (FSQ, Mentzer et al. (2024)) instead of dictionary learning (van den Oord et al., 2017). As result, our codebook is fixed and we obviate two codebook learning loss terms, which stabilizes early training. In this section, we will give an overview of our discretization method which utilizes codebooks. First, let us assume the output of the encoder is tensor' x E 1xx with d dimensions and b as the number of channels.
Each latent dimension is quantized into codebook C. That is, we have d independent codebooks, one for each latent dimension. Our first step to define the size of the codebook for each dimension, i.e. to define the ordered set of quantization levels L. Each quantization level L; corre-
{L1, L2, sponds to the i-th channel, e.g. L defines the number of discrete values in the first channel, Ly for the sec- ond and so on. In short, a quantization level of e.g. L = 11 would mean that discretize each dimen- sion in the i-th channel into 11 distinct values/sym- bols. We use integers as symbols, which would mean that the code for dimension d in channel i would be symbol from the set (-5,- -4.-. ...... 4, 53. In practice, for fast conversion from continuous values to codes we use similar discretization scheme as FSQ and apply the function
Fixed, implicit codebook
x
Figure 2: Illustration of Codebook (C) FSQ's codebook is a b-dimensional hy- percube (left). This figure illustrates b=3-dimensional codebook, where each axis of the 3-dimensional hypercube (left) corre- sponds to one dimension of the codebook (right). The 7th dimension of the hypercube is discretized into L, values, e.g., the and y-axis are discretized into Lo = L1 = 5 and the z-axis into L3 = 4. Code symbols (here integers) are normalized to the range [- 1, 1].
(1
Li 2 -tanb(a..))
f: f: L, L, round
(7)
to each channel, taking the output æ of the encoder and the channel quantization level Li- This approach results in codebook with |C| = Il L unique codes for each dimension d, each code being made of symbols, i.e. b-dimensional vector.
Intuitively, this results in a Voronoi partition of the b-dimensional space in each dimension d, where any point in space is assigned to one of the equidistantly placed centroids via Eq. (7). See Fig. 2 for visualization. In effect, this leads to an efficient and fast discretization of the latent embedding space.
In practice, Eq. (7) is not differentiable. To solve this for using standard deep learning libraries, we use the trriggh-thhoough gradient estimation (STE) approach with round soandddst(x))__)_ : x + x sg(round| (z:) 2) where the function sg(.) stops the gradient flow. Furthermore, we normalize codes to be the range [-1, after the discretization step as improved performance was reported by Mentzer et al. (2024). The hyperparameters of this approach are the number of channels b and the number of code symbols per channel Li, i.e. quantization levels. our experiments, we found the quantization levels L (5,3) (i.e. b- = 2 channels) to be sufficient.
World model training We train our world model components Eg, dg, Rg jointly using backpropaga- tion through time (BPTT) with the following objective
(()) Biio.a.'.ro..
Σ n² (CE(po čh, an), ch+1) + Rg (Ch, an) - 11
h=
Latent-state consistency
Roward prediction
with èo =(E)(0o))
Ca+11 Po(Cn41 Ch., anh) Ch (g(((ee(o)))
(8)
First latent
Stochastic dynamics
Target latent code
where H denotes the multi-step prediction horizon and 7 is the discount factor. The first predicted latent code čo is obtained by passing the observation Ou through the encoder and then quantizing the
output. At subsequent time steps, the dynamics model predicts the probability mass function over the next latent code Po Ch++1 1 ča, ax) Given this probabilistic dynamics model, we must consider how
For simplicity, we omit here the batch dimension.
5

---

Published as a conference paper at ICLR 2025
to make H H-step predictions in the latent space. In practice, we propagate uncertainty by sampling and we use the traiightthhoouge (ST) Gumbel-softuax trick (Jang et al., 2017; Maddison et al., 2017) samples to the encoder. Note that gradients must flow
so that gradients backpropagate through back to the encoder at the first time step when it was used to obtain the first latent code én, as the target codes c. are obtained by passing the next observation o' through the encoder and using the stop gradient operator sg. We then train our dynamics "classifier" using the cross-entrops (CE) loss. Finally, we note that our reward model Rg is trained jointly with the encoder eg and dynamics model Pe to ensure that the world model can accurately predict rewards in the latent space.
Policy and value learning We learn the policy ,,,(c) and action-value functions ( a) in the latent space using the actor-critic RL method TD3 (Fujimoto et al., 2018). However, we follow Yarats al. (2021b); b); Zhao et al. (2023) and augment the loss with N-step returns. The main difference to TD3 is that instead of using the original observations o, we map them through the encoder c= (()) and learn the actor/critic in the discrete latent space c. We also reduce bias in the TD target by following REDQ (Chen et al., 2021) and learning an ensemble of N₂ 5 critics, as was done in TD-MPC2. When calculating the TD target we randomly subsample two of the critics and use the minimum of these two. Let denote the indices of the two randomly subsampled critics as M. The critic is then updated by minimizing the following objective:
N
D) = of
(Dvn ([(eolo))) at) y)²
(9)
Ng
k=1
N-1 y =
7 retn + 7
min qiv (f(e(()AN)) actN), with actn = =T((C(4m) + ELtmo
KEM
CcNN
where we use policy smoothing by adding clipped Gaussian noise Cttx clip (N(O, -c, to the action Oren = TT ToyCrr4, (Cr+.)) + Ee-n- We then use the target action-value functions qú and the target policy Ty to calculate the TD target y. Note that the target networks use an exponential moving average, i.e. [47, 7 (1 [), 7 + [1b,m] n)- We follow REDQ and learn the actor by minimizing
Lx = -Eo -D
Σ gur qer (((a((o))) Ty (((o(o)))))
(10)
LM
PAEM
That is, we train the actor to maximize the average action value over two subsampled critics.
Summary Whilst this world model shares some similarities with TD-MPC2, there are some important distinctions. First, the latent space is represented discrete codebook which enables DC- MPC to train the dynamics model using the cross cross-entropy loss. Importantly, the cross-entropy loss considers a (potentially multimodal) distribution over the predicted latent codes during both training and inference. In contrast, TD-MPC2 considers deterministic dynamics and uses mean squared error regression. Interestingly, our experiments suggest that our stochastic dynamics model offers benefits in deterministic environments. Second, DC-MPC does not use value prediction when training the encoder. Instead, we follow the insight from Zhao et al. (2023) that value prediction is not necessary for obtaining good latent representation and instead, train the valiue function separately.
Importantly, our discrete latent space is parameterized as a set of discrete codes from a codebook. It is worth highlighting that our codebook encoding preserves ordinal relationships between observations. This contrasts with one-hot encodings which were used by DreamerV2 (Hafner et al., 2022). See Sec. 3 for comparison of the different discrete encodings. We hypothesize that this will offer significant improvements when representing continuous state vectors in a discrete space.
4.2 DECISION-TIME PLANNING
DC-MPC follows TD-MPC2 and leverages the world model for decision-time planning. It uses MPC to obtain a closed-loop controller and uses (modified) MPPI (Williams et al., 2015) as the underlying trajectory optimization algorithm (Betts, 1998; Scannell et al., 2021). MPPI is sampling-bai basedd based trajectory optimization method which does not require gradients. See Alg. 2 for full details. At each environment step, we estimate the parameters P Thu of a diagonal multivariate Gaussian over

---

Published as a conference paper at ICLR 2025
H step action sequence that maximizes the following objective
Mo:H Oo:H arg max E "ao.a~N(Hor) (inglordd [J(OO.H, o)]
(Ila)
Ho.s. 00:H
11 1-1
[an. 0) = Σ
? Re an) 7
Σ DLEM
Avn (C
(11b)
[M|
h=0
[C]
s.t. è = ((e(o)) and Ch+1
Pr(čn-1 = cl Cni and
(11c)
where H is the planning horizon and is discount factor. MPPI solves Eq. (1I) in an iterative manner. It starts by sampling candidate action sequences and evaluating them using the objective J J o). It then refits the sampling distribution's 's parameters μ.:H. d based on weighted average. After
several iterations, we select an action trajectory and apply its first action ag in the environment. Note that during training promote exploration by adding Gaussian noise. Importantly, Eq. (11) uses the action-value function qu (c, a) to bootstrap the planning horizon such that it estimates the full RL objective. DC-MPC follows TD-MPC2 and warm starts the planning process with N₂ action sequences originating from the prior policy Ty and we warm start by initializing Pu-H. Hh. as the solution to the previous time step shifted by one. See App. A and Alg. for further details.
Note that at planning time, we do not sample from the transition dynamics p(Cn+1 Cs. a,) because this introduces unwanted stochasticity. Instead, take the expected code, which is weighted sum over the codes in the codebook. Whilst the expected value of a discrete variable does not necessarily take valid discrete value, find it effective in our setting. This is likely because our discrete codes have an ordering such that expected values simply interpolate between the codes in the codebook.
5 EXPERIMENTS
In this section, we experimentally evaluate DC-MPC in a variety of continuous control tasks from the Deep-Mndd Control Suite (DMControl) (Tassa et al., 2018), Meta-World (Yu et al., 2019) and MyoSuite (Vittorio et al., 2022) against number of baselines and ablations. Our experiments seek to answer the following research questions:
RQ1 Does DC-MPC's discrete latent space offer benefits over continuous latent space? RQ2 What is important for learning a latent space: (i) classification loss, (ii) discrete codebook, (iii) stochastic dynamics or (iv) multimodal dynamies?
RQ3 Does DC-MPC's codebook offer benefits for dyyamics/value/polliii learning over alternative discrete encodings such as (i) one-hot encoding (similar DreamerV2) and (ii) label encoding?
RQ4 How does DC-MPC compare to statc-of-ther -f-therr model-based RL algorithms leveraging latent state embeddings, especially in the hard DMControl and World tasks?
Experimental Setup We compared DC-MPC against two sfate-of-the -rrt art model-based RL base- lines, namely D3eamerV3 (Hafner al., 2023) which utilizes discrete one-hot encoding as its latent state and TD-MPC2 (Hansen et al., 2023) using continuous latent space. We also compare against soft actor-critic (SAC) (Haarnoja et al., 2018), model-free RL baseline, and the original TD-MPC (Hansen et al., 2022). Our proposed approach utilized a latent space with d 512 dimensions and b = 2 channels, with 15 code symbols per dimension by using FSQ levels {L1 = 5, L2 = 3}.
5.1 COMPARISON OF DIFFERENT LATENT SPACES
We first evaluate how different latent dynamies formulations affect the performance. We seek to answer the following: (i) do discrete latent spaces offer benefits over continuous latent spaces? (ii) does training with classification (cross-ent (rooss-ntrroy) offer benefits over mean squared error regression? and (iii) does modeling stochastic (and potentially multimodal) transition dynamics offer benefits?
In our experiments, we consider both continuous and discrete latent spaces to investigate the impact of discretizing the latent space of the world model. In Figs. and 9, the experiments with discrete latent spaces are labelled with "Discrete" (red, green, and purple) whilst continuous latent spaces are labelled "Continuous" (orange). We also evaluate DC-MPC using the simplical normalization used
7

---

Published as a conference paper at ICLR 2025
IQM Noormalized Score
DMCotrll 10 tasks
MetaWorld tasks
IQM
Optimality Gap
Continuous-Ma5
SimNorm-MSE
Discrete-MSS
Discrete + det
Discret++CE+sstcch (ours)
0.0
0.3
0.5
0.7
0.3
0.5
0.7
500
1000
500
1000
Normalizod ized Score
Env. Steps (1c3) Env. Steps (1c3)
(a) Aggregate statistics at 500k environment steps
(b) Training curves
Figure 3: Latent space ablation Evaluation of (i) discrete (Discrete) vs continuous (Continuous) la- tent spaces, (ii) using cross-entrope (CE) vs mean squared error (MSE) for the latent-state consistency loss, and (iti) formulating deterministic (det) vs stochastic (stoch) dynamics model. Discretizing the latent space (red) improves sample efficiency over the continuous latent space (orange) and for- mulating stochastic dynamies and training with cross-entropy (purple) improves performance further.
in TD-MPC2- which bounds the latent space labelled "SimNorm" (blue) Experiments labelled with "MSE" were trained with mean squared error regression whilst those labelled "CE" were trained with the cross-entropy classification loss. The experiment labelled "Discrete+CE+det" used FSQ to get discrete latent space and trained with the cross-entropy loss, where the logits were obtained as the MSE between the dynamics prediction and each code in the codebook. This experiment enabled us if DC-MPC's performance boost resulted from training with the cross-entropy loss or from making the dynamics stochastic. In Fig. 9, experiments labelled with "log-lik." were trained by maximizing the log-likelihood, i.e. cross-entropy for "FSQ-log-lik." (purple), Gaussian log prob. for "Gaussian+log-lik (blue), and Gaussian mixture log prob. for "GMM+log-lik." (green).
Discrete vs continuous latent spaces The experiments using discrete latent spaces (red and purple) significantly outperform the ones with continuous latent spaces in terms of sample efficiency. This suggests that our discrete codebook encoding offers significant benefits over continuous latent spaces.
Classification vs regression Interestingly, training a deterministic discrete latent space using MSE regression (red) does not perform as well training stochastic discrete latent space using classification (purple). However, our experiment with the deterministic discrete latent space using classification (green) confirms that the benefit arises from the stochasticity of our latent space. This suggests that using straight- through Gumbel-softmax sampling (Jang al., 2017) when making multi-step dynamics predictions during training boosts performance. Our results extending TD-MPC2 use DC-MPC's discrete stochastic latent space in Fig. 6 support this conclusion.
Deterministic vs stochastic Given that modeling a stochastic latent space and training with maximum log- likelihokkoo is beneficial for discrete latent spaces, latent spaces. To this end, we formulate two stochastic continuous latent spaces and compare them in Fig. 9. The first models unimodal Gaussian distribution (blue) whilst the second models multimodal Gaussian mixture model (GMM) (green). Interestingly. these stochastic transition models sometimes increase sample efficiency on DMControl tasks when compared to their deterministic counterparts (orange). However, they drastically underperform on Meta- World tasks.
test if this holds in continuous
Our method (purple) has discrete latent space, trained by maximum log- seglikkeihoodd (i.e. cross- entropy), and models (potentially multimodal) distribution over the latent transition dynamics during training. These factors, combined with using ST Gumbel-softmax sampling, offer improved sample efficiency over continuous latent spaces.
5.2 IMPACT OF LATENT SPACE ENCODING
Our world model consists of NNs for the dynamics Ps(c' | c, a). reward Re (c, a), critic Qp a). and prior policy TT,(c), which all make predictions given the discrete codebook encoding In Fig. 4, we evaluate what happens when we replace the codebook encoding with (i) label encoding
Gzodes-
{0, 1310 given >! v; 1. In these experiments, we did not modify the dynamics Ps (c' c, a), that is, the dynamics continued to make predictions using the codebook encoding and did not use the one-hot or label encodings. This is because when we replaced the codebook encoding with either one-hot or label encodings, this led to the training curves (environment step episode return) flat-lining and unable solve the
Glabel = i E {1,.... |C]} and (ii) one-hot encoding Cone-hot
8

---

Published as a conference paper at ICLR 2025
Dog Run
Humanoid Walk
Dog Run
Humanoid Walk
Episode Recurn
500
500
1000
1500
1000 1500
10
30
10
Env. Steps (le3)
Env. Steps (Ie3)
Time (hours)
Time (hours)
Codes (ours)
One hat
Label
Figure 4: Discrete encodings ablation DC-MPC with its discrete codebook encoding (purple) outperforms using DC-MPC with one-hot encoding (red) and label encoding (blue), in terms of both sample efficiency (left) and computational efficiency (right). Dynamics model used codes ps a) whilst reward Rg (e, a), critic Q a) and prior policy Tr (e) used the respective encoding e.
task. This suggests that our codebook encoding is needed in our self-supervised world model setup. Nevertheless, we evaluated the performance when changing the encoding for the other components.
We evaluated the following experiment configurations: Codes (purple): All components used codes: dynamics (c'c' c, a). reward Rg (c, a). critic (c, a) and prior policy (c) Label (blue): Dynamics model used codes (c | a) whilst reward Rg (elabel, a), critic (elabel) a) and prior policy n, (etsel) used labels Glabel obtained from the code's index in the codebook. One-hot (red): Dynamics model used codes Pa(c' c, a) whilst reward Rg Conee-be a). critic Qu (Cone-hot: a) and prior policy n, (Cone-ba), used the one-hot Cone-hot representation of the label encoding.
The label encoding (blue) struggles to learn in the Humanoid Walk task and is often less sample efficient than the alternative encodings. This is likely because the label encoding is not expressive enough to model the calti-dimensional ordinal structure of our codebook. Let provide intuition via a simple example. Our codebook has 6 = 2 channels, so two different codes may take the form ecodes (A) |0.5.-0.5 0..5] and ecodes (B) [0, 0.5). As result, our codebook encoding can model ordinal structure in both of its channels, i.e., (A)) Gcodes whilst Goodes (cades()), (A)₂ < Geodes [ 2-
The corresponding label encoding would encode this ( 1 and Gjabel [B) 2, which incorrectly implies that B > A. In short, the label encoding cannot model the multi cultti-dimensionll ordinal structure of the codebook C. In contrast, the one-hot encoding (red) matches the codebook encoding in terms of sample efficiency in all tasks except Humanoid Walk. However, the one-hot encoding introduces an extremely large input dimension for the reward, value and policy networks, and this significantly slows down training. See Sec. 3 for further details on why this is the case.
5.3 PERFORMANCE OF DC-MPC
In Figs. 14, 16 and 18, we compare the aggregate performance of DC-MPC against TD-MPC2, D3eamerV3, TD-MPC, and SAC, in 30 DMControl, 45 Meta- World, and 5 MyoSuite tasks respec- tively, with 3 seeds per task. Some tasks in DMControl are particularly high-dimensionnl. For instance, the observation space of the Dog tasks is O E R223 and the action space is A E IR38 and for Humanoid, the observation space O E 67 and the action space E 2224, Fig. 13 shows that DC-MPC excels in the high dimensional Dog and Humanoid environments when compared to the baselines. We hypothesize that our discretized representations are particularly beneficial for simplifying learning the transition dynamics in -igenniieesi spaces, making DC-MPC highly sample efficient in these tasks. Similarly, we find that DC-MPC outperforms DreamerV3 in simulated manipulation tasks in the Meta- World task suite (Figs. 5, 15 and 16). We also see that DC-MPC
DM 30 tasks
Meta-World 45 tasks
MyoSuite tasks
Epioode Return
Success (9%)
Success Rate (5)
100
DC-MPC (ours)
DreamerV3
500
50
50
SAC
TD-MPC
TD-MPC2
1M
3M
1M Enc.
2M
0 BM Env. Stops
Env.
Steps
Steps
Figure 5: Aggregate training curves in DMControl, Meta-World & MyoSuite DC-MPC generally matches TD-MPC2 whilst outperforming DreamerV3, SAC and TD-MPC across all tasks. We plot the mean (solid line) and the 95% confidence intervals (shaded) across 3 seeds per task.

---

Published as a conference paper at ICLR 2025
IQM Normalized Score
DMControl tasks
Meta-World 10 tasks
1.0
IQM
Optimality Gap
TD-MPC2
0.5
TD-MPC2 DCWM
DC-MPC (ours)
0.0
0.15 0.20 0.25
500
1000
500
1000
Normalized Score
Env. Steps (1e3)
Env. Steps (1e3)
(a) Aggregate statisties at 1M environment steps
(b) Training curves
Figure 6: TD-MPC2 with DCWM Adding DC-MPC's discrete and stochastic latent space to TD-MPC2 improves performance. See Apps. B and 8.10 for more details.
generally matches the performance of TD-MPC2. Comparing the results at global level (Fig. 5), can find that our proposed method performs well across all benchmarks.
It is important note that TD-MPC2 has multiple algorithmic differences to DC-MPC which means that a straight-up comparison between them is not only affected by the latent space design. For example, (i) uses soft-actor critic (SAC) to learn the prior policy (helping with exploration in sparse reward tasks), (ii) learns the value function jointly with the world model, and (iii) the reward and value functions are formulated using discrete regression in a log-transformed space. In Fig. 6, we show that incorporating DCWM's stochastic and discrete codebook latent space into TD-MPC2 (red) offers improvements over vanilla TD-MPC2. See Apps. В and B. for more details on these experiments and App. B.9 where we tried the same experiments with DreamerV3. DreamerV3's performance is poor in the harder tasks so we do not see any benefit from using DCWM. However, identify that its poor performance stems from using observation reconstruction.
Further Experiments In Apps. B.1 and B.3, we evaluate DC-MPC's sensitivity to codebook size |C] and latent dimension d, respectively, in App. B.3, show that stochastic continuous latent spaces do not appear to offer the same benefits as stochastic discrete latent spaces, in App. B.4, we ablate FSQ and show that it either matches or outperforms vector quantization (VQ) whilst being simpler, and in App. B.5, we show the benefit of using REDQ's S ensemble of critics vs the standard double Q approach.
6 CONCLUSION
We have presented DC-MPC, a world model that learns a discrete and stochastic latent space using codebook encodings and a cross-entropy based self-supervised loss for model-based RL. DC-MPC demonstrates strong performance continuous control tasks, including Meta- Seta-World& and the complex DMControl Humanoid and Dog tasks, where it exceeds or matches the performance of SOTA baselines. Our results indicate that using straight-through Gumbel-softmax sampling when making multi-step dynamics predictions is beneficial for world model learning, both in DC-MPC and our experiments where we modified TD-MPC2's 's latent space. In summary, we have demonstrated the benefit of discrete latent space with codebook encodings embedding or classical discrete spaces such as label and one-hot encodings. These findings open up a new interesting avenue for future research into discrete embeddings for world models.
standard continuous latent
Limitations and Future Work As our goal was to evaluate latent space design, we did not prioritize making DC-MPC run with single set of hyperparameters and tuned the noise schedule and the N-step return for some tasks. In future work, it would be interesting to make DC-MPC robust to hyperparameters. For example, it would be interesting model the epistemic uncertainty associated with the latent transition dynamics arising from learning from limited data and using it to equip al.
DC-MPC with more principled exploration mechanism like Chua et al. (2018), Scannell (2024c); Daxberger et al. (2021) and remove the task-specific noise schedules. It would also be interesting to investigate if our results hold for different world model backbones (Deng et al., 2023; NVIDIA et al., 2025), such as Transformers (Vaswani et al., 2017; Robine et al., 2022; Zhang et al., 2023; Micheli et al., 2022; Bar et al., 2024) and diffusion models (Ho et al., 2020; Alonso et al., 2024). Finally, it would be interesting to investigate how well DC-MPC scales (Kaplan et al., 2020; Henighan et al., 2020; Hoffmann et al., 2022) and if it an effective setup for generalist (i.e. multi-embodiment) world modeling (Reed et al., 2022; Zhao al., 2025).
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
Aidan Scannell and Kalle Kujanpää were supported by the Research Council of Finland from the Flagship program: Finnish Center for Artificial Intelligence (FCAI), Amo Solin and Yi Zhao acknowl- edge funding from the Research Council of Finland (grant ids 339730 and 357301, respectively) and Mohammadreza Nakhaei acknowledges funding from Business Finland (BIOND4.0 - Data Driven Control for Bioprocesses). Kevin Sebastian Luck is supported by the project TeNet: Text-to-Nelwork for Fast and Enfrie--ffficientt Robol Control with file number NGF. (GF.1609.22 24009221..00 221.0111155 .015 of the research programme National Growth Fund AiNed XS Europe 24-2 which is financed by the Dutch Research Council (NWO). We acknowledge CSC IT Center for Science, Finland, for awarding this project access to the LUMI supercomputer, owned by the EuroHPC Joint Undertaking, hosted by CSC (Finland) and the LUMI consortium through CSC. We acknowledge the computational resources provided by the Aalto Science-I7 project.

---

Published as a conference paper at ICLR 2025
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Bellemare. Deep Reinforcement Learning at the Edge the Statistical Precipice. In Advances in Neural Information Processing Systems, volume 34, pp. 29304-29332 Curran Associates, Inc., 2021.
James F Allen and Johannes A Koomen. Planning using a temporal world model. In Proceedings of the Eighth international joint conference Artificial intelligence- vaalligence-Vollmme 2. pp. 741-747, 1983.
Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and François Fleuret. Diffusion for World Modeling: Visual Details Matter in Atari In The Thirty- eighth Annual Conference on Neural Information Processing Systems, November 2024.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E, Hinton. Layer Normalization. arXiv preprint 1rXiv:167.004490 2016.
Amir Bar, Gaoyue Zhou, Danny Tran. Trevor Darrell, and Yann LeCun. Navigation World Models. arXiv preprint arXiv:2412.03577 2024.
Kenneth Basye, Thomas Dean, Jak Kirman, and Moises Lejter. planning, perception, and control. IEEE Expert, 7(4):58-65, 1992.
approach to
Richard Bellman. A Markovian Decision Process. Journal of Mathematics and Mechanics, 6(5): 679-684, 1957. ISSN 0095-9057.
John T. Betts. Survey of Numerical Methods for Trajectory Optimization. Journal of Guidance, Control, and Dynamics, 21(2):193-25) 207, March 1998. doi: 10.2514/2.4231.
Huiwen Chang. Han Zhang, Jarred Barber, Aaron Maschinot, Jose Lezama, Lu Jiang, Ming- H-sun Yang, Kevin Patrick Murphy, William T. Freeman, Michael Rubinstein, Yuanzhen Li, and Dilip Krishnan. Muse: Text-To-lmage Generation via Masked Generative Transformers. In Proceedings of the 40th International Conference on Machine Learning, pp. 4055-4075. PMLR, July 2023.
Xinyue Chen, Che Wang. Zijian Zhou, and Keith Ross. Randomized Ensembled Double Q-Leaming: Learning Fast Without Model. In International Conference on Learning Representations, 2021.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep Reinforcement Learning in Handful of Trials using Probabilistic Dynamics Models. In Advances in Neural Information Processing Systems, volume 31, 2018.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace Redux Effortless Bayesian Deep Learning. In Advances in Neural Information Processing Systems, volume 34, pp. 20089-20109 Curran Associates, Inc., 2021.
Fei Deng, Junyeong Park, and Sungjin Ahn. Facing Off World Model Backbones: RNNs, Transform- ers, and S4. In Advances in Neural Information Processing Systems, volume 36, pp. 72904-72930, December 2023.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming Transformers for High-Resolution Image Synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 18873-12883 2021.
Jesse Farebrother, Jordi Orbay, Quan Vuong, Adrien Ali Taiga, Yevgen Chebotar, Ted Xiao, Alex Irpan. Sergey Levine, Pablo Samuel Castro, Aleksandra Faust, Aviral Kumar, and Rishabh Agarwal. Stop Regressing: Training Value Functions via Classification for Scalable Deep RL, March 2024.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing Function Approximation Error in Actor- Critic Methods. In Proceedings of the 35th International Conference on Machine Learning, pp. 1587-1596. PMLR, July 2018
Ignat Georgiev, Varun Giridhar, Nicklas Hansen, and Animesh Garg. PWM: Policy Learning with 2024.
Large World Models. arXiv preprint
12

---

Published as a conference paper at ICLR 2025
Raj Ghugare, Homanga Bharadhwaj, Benjamin Eysenbach, Sergey Levine, and Russ Salakhutdinov. Simplifying Model-based RL: Learning Representations, Latent-space Models, and Policies with One Objective. In The Eleventh International Conference Learning Representations, September 2022.
David Ha and Jürgen Schmidhuber. Recurrent World Models Facilitate Policy Evolution. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor. International Conference Machine Learning, pp. 1 1861-1870. PMLR, July 2018.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603 2019a.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning Latent Dynamics for Planning from Pixels. In International Conference Machine Learning, pp. 2555-22222 2565. PMLR, May 2019b.
Danijar Hafner, Timothy P. Lillierap, Mohammad Norouzi, and Jimmy Ba. Mastering Atari with Discrete World Models. International Conference on Learning Representations, February 2022.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04100 2023.
Nicklas Hansen, Hao Su, and Xiaolong Wang. TD-MPC2: Scalable, Robust World Models for Continuous Control. In The Twelfth International Conference on Learning Representations. October 2023.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Temporal Difference Learning for Model Predictive Control. In Proceedings the 39th International Conference Machine Learning, pp. 8387-8406. PMLR, June 2022.
Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B. Brown, Prafulla Dhariwal, Scott Gray, Chris Hallacy. Benjamin Mann, Alec Rad- ford, Aditya Ramesh, Nick Ryder, Daniel M. Ziegler, John Schulman, Dario Amodei, and Sam McCandlish. Scaling Laws for Autoregressive Generative Modeling, November 2020.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. In Advances in Neural Information Processing Systems, volume 33, pp. 6840-6851. Curran Associates, Inc., 2020.
Jordan Hoffimann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hemnigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training Compute-Optimal Large Language Models, March 2022.
Kyle Hsu, William Dorrell, James Whittington, Jiajun Wu, and Chelsea Finn. Disentanglement via Latent Quantization. Advances in Neural Information Processing Systems, 36:45463-45488, December 2023.
Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for pomdps. In International Conference on Machine Learning, pp. 2117-2126. PMLR, 2018.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizztioon with gumbel-softmax. In International Conference on Learning Representations, 2017.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling Laws for Neural Language Models. arXiv preprint arXiv:2001.08361, 2020.
13

---

Published as a conference paper at ICLR 2025
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Diederik Kingma and M. Welling. Auto-Encoding Variational Bayes. In International Conference on Learning Representations, 2014.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004413649, 2020.
Yann LeCun, A Path Towards Autonomous Machine Intelligence Version 0.9.2, 2022-06-27.
Michael Lutter, Leonard Hasenclever, Arunkumar Byravan, Gabriel Dulac-Amold, Piotr Trochim, Nicolas Heess, Josh Merel, and Yuval Tassa. Learning dynamics models for model predictive agents. arXiv preprint arXiv:2109. 14311, 2021.
Haoyu Ma, Jialong Wu, Ningya Feng, Chenjun Xiao, Dong Li, Jianye Hao, Jianmin Wang, and Mingsheng Long. Harmony Task Harmonization Inside World Models. In Proceedings of the 41st International Conference on Machine Learning, pp. 33983-34007. PMLR, July 2024.
Chris Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: continuous relaxation of discrete random variables. In International Conference on Learning Representations, 2017.
Fabian Mentzer, David Minnen, Eirikur Agustsson, and Michael Tschannen. Finite Scalar Quan- tization: VQ-VAE Made Simple. In International Conference on Learning Representations, 2024.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are Sample-Efiicient World Models. In The Eleventh International Conference on Learning Representations, September 2022.
Diganta Misra. Mish: A self regularized non-monotonic activation function. arXiv preprint arXiv: 08681, 2019.
NVIDIA, :, Niket Agarwal, Arslan Ali, Maciej Bala, Yogesh Balaji, Erik Barker, Tiffany Cai, Prithvijit Chattopadhyay, Yongxin Chen, Yin Cui, Yifan Ding, Daniel Dworakowski, Jiaojiao Fan, Michele Fenzi, Francesco Ferroni, Sanja Fidler, Dieter Fox, Songwei Ge, Yunhao Ge, Jinwei Gu, Siddharth Gururani, Ethan He, Jiahui Huang, Jacob Huffman, Pooya Jannaty, Jingyi Jin, Seung Wook Kim, Gergely Klár, Grace Lam, Shiyi Lan, Laura Leal-Taixe, Anqi Li, Zhaoshuo Li, Chen-Hsuan Lin, Tsung- Yi Lin, Huan Ling, Ming-Yo Yu Liu, Xian Liu, Alice Luo, Qianli Ma. Hanzi Mao, Kaichun Mo. Arsalan Mousavian, Seungjun Nah, Sriharsha Niverty, David Page. Despoina Paschalidou, Zeeshan Patel, Lindsey Pavao, Morteza Ramezanali, Fitsum Reda, Xiaowei Ren, Vasanth Rao Naik Sabavat. Ed Schmerling, Stella Shi, Bartosz Stefaniak. Shitao Tang, Lyne Tchapmi, Przemek Tredak, Wei-Cheng Tseng, Jibin Varghese, Hao Wang, Haoxiang Wang, Heng Wang, Fing-Chun Wang, Fangyin Wei, Xinyue Wei, Jay Zhangjie Wu, Jiashu Xu, Wei Yang, Lin Yen-Chen, Xiaohui Zeng, Yu Zeng. Jing Zhang, Qinsheng Zhang, Yuxuan Zhang, Qingqing Zhao, and Artur Zolkowski. Cosmos world foundation model platform for physical ai. arXiv preprint arXiv:2501.03522 2025.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, High-Performance Deep Learning Library. In Advances Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and llya Sutskever. Zero-Shot Text-to-Image Generation. In Proceedings of the 38th International Conference on Machine Learning, pp. 8821-8831. PMLR, July 2021.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay. Jost Tobias Springenberg, et al. Generalist Agent. Transactions Machine Learning Research (TMLRR, 2022.
14

---

Published as a conference paper at ICLR 2025
Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-base World Models Are Happy With 100k Interactions. In The Eleventh International Conference on Learning Representations, September 2022.
Reuven Y Rubinstein. Optimization of computer simulation models with rare events. European Journal of Operational Research, 99(1):89-112 112. 1997.
Aidan Scannell, Carl Henrik Ek, and Arthur Richards. Trajectory Optimisation in Learned Multimodal Dynamical Systems Via Latent-ODE Collocation. In Proceedings of the IEEE International Conference Robotics and Automation. IEEE, 2021.
Aidan Scannell, Kalle Kujanpää, Yi Zhao, Mohammadreza Nakhaeinezhadfard, Arno Solin, and Joni Pajarinen. Quantized Representations Prevent Dimensional Collapse in Self-predictive RL. In ICML 2024 Workshop: Aligning Reinforcement Learming Experimentalists and Theorists, July 2024a.
Aidan Scannell, Kalle Kujanpää, Yi Zhao, Mohammadreza Nakhaei, Arno Solin, and Joni Pajarinen. iQRL- - Implicitly Quantized Representations for Sample-efficient Reinforcement Learning. arXiv preprint arXiv:2406.0200 02696, 2024b.
Aidan Scannell, Riccardo Mereu, Paul Edmund Chang, Ella Tamir, Joni Pajarinen, and Arno Solin. Function space parameterization of neural networks for sequential learning. In The Twelfth International Conference on Learning Representations, 2024c.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Gracpel, Timothy Lillicrap, and David Silver. Mastering Atari, Go. Chess and Shogi by Planning with a Learned Model. Nature, 588(7839):604 609, December 2020. ISSN 0028-0836, 1476-4687. doi: 10. s41586-020-003551-4
Max Schwarzer, Ankesh Anand, Rishab Goel, R. Devon Hjelm, Aaron Courville, and Philip Bachman. Data-Efficien Reinforcement Learning with Self- Peff-Predictive Representations. In International Conference on Learning Representations, October 2020.
R.S. Sutton and A.G. Barto. Reinforcement Learning, Second Edition: An Introduction. Adaptive Computation and Machine Learning Series. MIT Press, 2018 ISBN 978-0-262-35270-3
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv:1807.000 00690, 2018
Aaron van den Oord, Oriol Vinyals, and koray kavukcuoglu. Neural Discrete Representation Learning. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Caggiano Vittorio, Wang Huawei, Durandau Guillaume, Sartori Massimo, and Kumar Vikash. Myosuite - a contact-rich simulation suite for musculoskeletal motor control. arXiv preprint 1:Xi: 2022.
Tongzhou Wang, Simon Du, Antonio Torralba, Phillip Isola, Amy Zhang, and Yuandong Tian. Denoised MDPs: Learning World Models Better Than the World Itself. In Proceedings of the 39th International Conference on Machine Learning, pp.
PMLR, June 2022.
Grady Williams, Andrew Aldrich, and Evangelos Theodorou. Model predictive path integral control using covariance variable importance sampling. arXiv preprini arXiv:1509.00 1509.01 149, 2015.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2 2021a.
15

---

Published as a conference paper at ICLR 2025
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering Visual Continuous Control: Improved Data-Augmetted Reinforcement Learning. In International Conference on Learning Representations, October 2021b.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning (CoRL). 2019.
Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang. STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning. In Advances in Neural Information Processing Systems, volume 36, pp. 27147-27166, December 2023.
Yi Zhao, Wenshuai Zhao, Rinu Boney, Juho Kannala, and Joni Pajarinen. Simplified Temporal Consistency Reinforcement Learning. In Proceedings of the 40th International Conference Machine Learning, pp. 42227-42246. PMLR, July 2023.
Yi Zhao, Aidan Scannell, Yuxin Hou, Tianyu Cui, Le Chen, Amo Solin, Juho Kannala, and Joni Pajarinen. Generalist world model pre-training for efficient reinforcement learning. arxiv preprint arXiv: 2025.
16

---

Published as a conference paper at ICLR 2025
APPENDICES
This appendix organized as follows. In App. we provide further details on our method. App. B provides further experimental results, including evaluating DC-MPC's sensitivity to the codebook size in App. B.I, I, its sensitivity to latent dimension in App. B.2, further details on the latent space ablation App. B.3, comparison of DC-MPC using VQ instead of FSQ App. B.4, a comparison of DC-MPC's ensemble REDQ critic approach vs the standard double Q approach in App. B.5, full DeepMind control suite results in App. B.6, Meta-Words World results App. B.7, MyoSuite results in App. B.8, evaluation of DreamerV3 using DC-MPC's latent space in App. B.9 and evaluation of TD-MPC2 using DC-MPC's latent space App. B. In App. C, we provide further implementation details, including default hyperparameters, hardware, etc. In App. D, provide further details of the baselines and in App. E we detail the different DeepMind control, Meta- World and MyoSuite tasks used throughout the paper.
METHOD DETAILS
18
FURTHER RESULTS
B.1 SENSITIVITY TO CODEBOOK SIZE |C]
21
B.2 SENSITIVITY TO LATENT DIMENSION d
B.3 ABLATION OF LATENT SPACE
23
B.4 ABLATION FSQ VECTOR QUANTIZATION (VQ)
24
B.5 ABLATION OF REDQ CRITIC vs STANDARD DOUBLE Q APPROACI
25
B.6 DEEPMIND CONTROL RESULTS
26
B.7 META-WORLD MANIPULATION RESULTS
28
B.8 MYOSUITE MUSCULOSKELETAL RESULTS
30
B.9 DOES DCWM IMPROVE DREAMERV3?
B. 10 IMPROVING TD-MPC2 WITH DC-MPC
32
C IMPLEMENTATION DETAILS
33
BASELINES
E TASKS
36
appendiccs continue on next page-
17

---

Published as a conference paper at ICLR 2025
A
METHOD DETAILS
Alg. outlines DC-MPC's training procedure.
Algorithm DC-MPC's training
N, Input: Encoder cø, dynamics dg, reward Rg, critics {ga, h policy Thy learning rate a, target network update rate T, episode length T', replay buffer D =
for 1: -random episodes do D < DU (o, at, Oc+1, To To
> Collect data using random policy
end for
for 1: Nepisodes do
D < DU (o,, at, 01+1, Ti vi)T-o
> Collect data using DC-MPC
for i 1 to T do
00, db [] + 00 d.
+ aV (C(0, 60 (:
D. Update world model, Eq. (8)
v <- + oV (L.g th: D))
Update critic, Eq. (9)
if % == = 0 then
1) - 1) + a (C, D))
Update actor less frequently than critic, Eq. (10)
end if
(2b π] < (1 - +) r + - -| nl
Update target networks
end for end for
Alg. outlines how we perform trajectory optimization using MPPI (Williams et al., 2015), closely following the formulation of MPPI by Hansen et al. (2022), with two key modifications. First, during each rollout, we use the expected next latent state, i.e. a weighted sum over the codes in the world model training where we sample from the transition
codebook. Note that this contrasts dynamics p(Ct+1 (C,, Cn). This approach reduces the variance state transitions, which results in more stable trajectory evaluations. Second, we do add noise sampled from the standard deviation oū returned from MPPI. Instead, we promote exploration by adding noise sampled from a separate noise schedule. This method, inspired by TD3 (Fujimoto et al., 2018), strikes a better balance between exploration and exploitation, leading to more stable training performance.
is worth noting that MPPI resembles the CEM-based planner in Chua et al. (2018), however, instead of simply fitting a Gaussian to the top K action samples at each iteration, MPPI uses weighted importance sampling, which weights all samples by their empirical return estimates. However, we follow Hansen et al. (2022) and use hybrid approach, which selects the top K action samples (like CEM) but then use weighted importance sampling (like MPPI). At each iteration, we calculate the mean and variance of the action trajectory follows,
2
Pa:8 fit_mean
(((a )))))
K
Σ 1=1
K Ω
Gu.H
(5)
(12)
importance weight
d fit var (I(a))) )0))) ΣK ()) (i)
(13)
ΣK 2(
where 20( is the exponentiated normalized empirical return estimate given by 2 exp (TMPPI ($()) max (f ...., ((N,+NN))))) Note that TMPPI is the (inverse) temperature parameter and (0)) denotes the return estimate for the 7th action trajectory an After J (default 6) iterations, we sample one of the top K action sequences {ao. bieftg....ft) where each action
(8)
sequence weighted by its empirical return estimate [((()) ( We then apply the first action a the environment.
(i*)
appendices continue on next page-
18

---

Published as a conference paper at ICLR 2025
Algorithm 2 DC-MPC's inference (modified MPPI)
Input: current observation o, planning horizon H, iterations J, population size Np. prior popula- tion size N, number of elites K, exploration noise std noise Co 4 (e)(())
D Encode observation into discrete code Initialize Mb: (o0.1)) )o with the solution from the last time step shifted by one.
for each iteration ......JJJ do
Sample N₂ action trajectories of length H/ from {an (N ()))))
D Sample
action candidates
Sample N₂ action trajectories of length H using Th and dø
Prior policy samples
for all N2 + N₂ action sequences {-() = [bb" ..... H 1=1
M 1
do P Trajectory evaluation
< ((1))
for step / = 0...., 1 do
(1)) (1) + 7 Re( a (6)
Compute immediate reward
Eh+1 = k=1 Pr(cn+1 ckk) ch a (8) c((k)
>> Compute next state
end for
((1) < (1)) + + 1
Ex 2ora (CH, 0H (1)
D. Bootstrap with ensemble of Q-functions
N.
end for ((1) (ik) topk(( ( (b{,+++a1))
P. Get top-K elite scores Update mean of action dist. > Update variance of action dist.
No. + fit_ fit_mean [(a x1()){
iE (e
d.u + fit_var var [( (())) iE ((i
end for
i' ~ Categorical (softmax({ ({ ..... p((k))
Sample action index according to scores Final output with exploration noise
return ag + with e~ NN(
appendices continue on next page-
19

---

Published as a conference paper at ICLR 2025
B FURTHER RESULTS
In this section, we include further results and ablations.
Aggregate metrics In Figs. 14, 16 and 18, compare the aggregate performance of DC-MPC against TD-MPC, TD-MPC2, DreamerV3, and SAC, in 30 DMControl tasks, 45 Meta- World tasks, and MyoSuite tasks respectively, with 3 seeds per task. Following Agarwal al. (2021), we report the median, interquartile mean (IQM), mean, and optimality gap at 1M environment steps, with emor bars representing 95% stratified bootstrap confidence intervals. For DMControl, we use min-max normalization as the maximum possible return in an episode is 1000 whilst the minimum is 0, i.e. Normalized Return Return/ (1000 0). For Meta- World, we report the success rate which does not require normalization as it is already between 0 and 1.
In Figs. and we report aggregate metrics over 10 DMControl and 10 Meta-World tasks. The tasks are follows:
DMControl 10: Acrobot Swingup. Dog Run, Dog Walk, Dog Stand, Dog Trot, Humanoid Stand, Humanoid Walk, Humanoid Run, Reacher Hard, Walker Walk. Aetaa-World 10: Button Press, Door Open, Drawer Close, Drawer Open, Peg Insert Side, Pick Place, Push, Reach, Window Open, Window Close.
appendices continue on next page-
20

---

Published as a conference paper at ICLR 2025
B.1 SENSITIVITY TO CODEBOOK SIZE [C]
In this section, we evaluate how the size of the codebook [C] influences training. We indirectly configure different codebook sizes via the FSQ levels L [L,,-. .......... La} hyperparameter. This is because the codebook size is given by |C| - L. The top row Fig. compares the training curves for different codebook sizes. The algorithm's performance is not particularly sensitive to the codebook size. A codebook that too large can result slower learning. The best codebook size varies between environments.
Given that a codebook has a particular size, we can gain insights into how quickly DC-MPC's encoder starts activate all of the codebook. The connection between the codebook size and the activeness the codebook is intuitive: the bottom row of Fig. shows that the smaller the codebook, the larger the active proportion.
Dog Run
Humanoid Walk
Reacher Hard
Walker Walk
1000
Codebook Active Percent (%) Episode Return
500
0
100
90
80
500
1000
1000
500
1000
200
Env. Steps (1e3)
Env. Steps (1e3)
Env. Steps (1e3)
Env. Steps (1e3)
L
[5
88
8,8)
Figure 7: Sensitivity codebook size We compare how the codebook size affects the performance of DC-MPC (top) and the percentage of the codebook that is active during training (bottom). In general, smaller codebooks become fully active faster than larger codebooks. We plot the mean and the 95% confidence intervals (shaded) across 3 random seeds for all environments.
appendiccs continue on next page-
21

---

Published as a conference paper at ICLR 2025
B.2 SENSITIVITY TO LATENT DIMENSION d
This section investigates how the latent dimension d affects the behavior and performance of DC- MPC in four different environments. In the top row of Fig. 8, our algorithm is robust to the latent dimension d. although latent dimension too small can result in inferior performance, especially in the more difficult environments. The bottom row of Fig. demonstrates that DC-MPC learns to use the complete codebook irrespective of the latent dimension.
that the performance of
Dog Run
Humanoid Walk
Reacher Hard
Walker Walk
1000
Codebook Active Percent (%) Episode Return
500
100
90
80
500
1000
500
1000
500
1000
250
500
Env. Steps (1e3)
Env. Steps (1e3)
Env. Steps (1e3)
Env. Steps (1e3)
128
d 256
d = 512
1024
Figure Sensitivity to latent dim d We compare how the latent dimension affects the performance of DC-MPC (top) and the percentage of the codebook that is active during training (bottom). In general, our algorithm robust to the latent dimension of the representation, although in more difficult environments, such as Humanoid Walk, too small can harm the agent's performance. We plot the mean and the 95% confidence intervals (shaded) across random seeds for all environments.
appendices continue on next page-
22

---

Published as a conference paper at ICLR 2025
B.3 ABLATION OF LATENT SPACE
In this section, we provide further details on the comparison of different latent spaces experiments in Sec. 5.1. To validate our method, we test the importance of quantizing the latent space and training the world model with classification instead of regression. In Fig. 9, we compare DC-MPC to world models with different latent spaces formulations, which we now detail.
Dog
Dog Walk
Humanoid Walk
Assembly
Epsode Retur
Epsode Return
1000
Epsode Return
Success Rate (%)
100
500
500
1000
1000
1000
1000
Env. Steps (Ie3)
Env. Steps (1e3)
Env. Steps (1e3)
Env. Steps (1e3)
MSE
FSQ+MSE.
Gaussian-log-lk
GMM-+log-lik
FSQ+log-lik (ours)
Figure 9: Latent space comparison Comparison of different latent space formulations. Continuous and deterministic latent space trained with MSE regression (orange), deterministic and discrete trained with MSE (red), continuous and unimodal Gaussian latent space trained with maximum log-likelihood (blue), continuous and multimodal GMM trained with maximum log-likelihood (green), and discrete trained with classification (purple). Discretizing the latent space with FSQ (red) improves sample efficiency and making the dynamics stochastic and training with classification (purple) improves performance further.
MSE (orange) First, we consider continuous latent space with deterministic transition dynamics trained by minimizing the mean squared error between predicted next latent states and target next latent states.
FSQ+MSE (red) Next, we consider quantization of the latent space and training based on mean squared error regression. This experiment allows us to analyze the importance of quantization.
Gaussian+log-lik. (blue) consider stochastic continuous dynamics, we configure the transition dynamics to model a Gaussian distribution over predictions of the next state. During training, sample from the Gaussian distribution using the reparameterization trick. The world model is then trained to maximize the log- ligliikl of the next latent state targets. This allows to investigate if modeling stochastic transition dynamics offers benefits when using continuous latent spaces.
GMM+log-lik. (green) To consider continuous multimodal transitions, consider Gaussian mixture with three components. During training. we sample a Gaussian from the mixture with the ST Gumbel-softma trick and then sample from the selected Gaussian using the reparameterization trick. The world model is then trained to maximize the log-likelihood of next latent state targets.
appendices continue on next page-
23

---

Published as a conference paper at ICLR 2025
B.4 ABLATION OF FSQ VS VECTOR QUANTIZATION (VQ)
To understand how the choice of using FSQ for discretization contributes to the performance of our algorithm, tried replacing the FSQ layer with standard Vector Quantization layer. We evaluated the methods in Walker Walk, Dog Run, Humanoid Walk, and Reacher Hard. We used standard hyperparameters.
0.25, and an EMA-updated codebook with a size of 256 and either 256 (dog) or 128 (other tasks) channels per dimension. We did not change other hyperparameters from DC-MPC. However, we found that to approach the performance of standard FSQ. VQ needs environment-depennden adjusting of the planning procedure. In Humanoid Walk, the performance of FSQ aligns closely with the VQ with a weighted sum over the codes in the codebook for planning (expected code) but significantly outperforms sampled VQ. Conversely, standard sampling is superior in Reacher Hard, which unsurprising, as the discrete codes VQ have not been ordered like in FSQ. The necessary environment-specife adjustments for VQ undermine its general applicability compared to FSQ.
Dog Run
Humanoid Walk
Reacher Hard
Walker Walk
1000
Episode Return
500
1000 Env. Steps (1e3)
1000 Env. Steps (1e3)
1000
200
Env. Steps (1e3)
Env. Steps (1e3)
FSQ
VQ (Sample)
VQ (Weighted Avg)
Figure 10: Ablation of FSQ VQ FSQ does not require the extra loss terms required by VQ and it generally performs equal to or better and VQ.
appendices continue on next page-
24

---

Published as a conference paper at ICLR 2025
B.5 ABLATION OF REDQ CRITIC VS STANDARD DOUBLE Q APPROACH
In this section, we compare the ensemble of Q-functions approach, used by DC-MPC, REDQ (Chen al., 2021) and TD-MPC2 (Hansen al., 2023), the standard double Q approach (Fujimoto et al., 2018). In Fig. 11, we evaluate how our default ensemble size of N₂ = 5 (purple) compares with the standard double Q approach, which is obtained by setting the ensemble size to NN. = (blue). Note that we always sample two critics so the Nº = 2 result reduces to the standard double Q approach. Fig. 11 shows that DC-MPC works fairly well with both approaches but the ensemble approach offers benefits in the harder Dog Run and Humanoid Walk tasks.
Dog Run
Humanoid Walk
Reacher Hard
Walker Walk
1000
Episode Return
500
500
1000
500
1000
500
1000
200 Env. Steps (1e3)
Env. Steps (1e3)
Env. Steps (1e3)
Env. Steps (1e3)
N₂ 2
N₂
Figure 11: Ablation REDQ critic vs standard double Q DC-MPC uses a Q ensemble, similar to REDQ, of size Ng 5 (purple) and sub samples two critics when calculating the mean minimum Q-value. We compare this approach to the standard double Q approach by setting N₂ = 2 (blue) and see that the ensemble approach offers slight benefit in the harder Dog Run and Humanoid Walk.
appendices continue on next page-
25

---

Published as a conference paper at ICLR 2025
B.6 DEEPMIND CONTROL RESULTS
Acrobot Swingup
Cartpole Balance
Cartpole Balance Sparse
Cartpole Swingup
Cartpole Swingup Sparse
Episods Return
Cheetah Run
Cup Catch
Cup Spin
Dog Run
Dog Stand
1000
Episade Return
Dog Trot
Dog Walk
Finger Spin
Finger Turn Easy
Finger Turn Hard
1000
Episodl Return
500
Fish Swim
Hopper Hop
Hopper Stand
Humanoid Run
Humanoid Stand
Epicade Return
Humanoid Walk
Pendulum Spin
Pendulum Swingup
Quadruped Run
Quadruped Walk
1000
Epicode Return
Reacher Easy
Reacher Hard
Walker Run
Walker Stand
Walker Walk
1000
Episode Return
500
1M
2M
1M
2M
2M
Env. Steps
Env. Steps
Env. Steps
Env. Steps
Env. Steps
DC MPC (aurs)
DreamerV3
SAC
TD MPC
TD MPC2
Figure 12: DeepMind Control results. DC-MPC performs well across variety of DMC tasks. We plot the mean (solid line) and the 93% confidence intervals (shaded) across seeds (DC-MPC) or 3 seeds episodes.
where each seed averages over 10 evaluation
appendiccss continue on next page-
26

---

Published as a conference paper at ICLR 2025
Dog Run
Dog Trot
Dog Walk
Humanoid Run
Humanoid Walk
1000)
Episade Return
1M
3M
0
2M
Env. Steps
Env. Steps
Env.
Steps
Env.
Steps
Env.
Steps
DC-MPC (ours)
SAC
TD-MPC2
TD-MPC
DreamerV3
Figure 13: High-dimension locomotion DC- MPC (purple) significantly outperforms TD-MPC2 (blue) and DreamerV3 (red) in the complex, high-dimensional locomotion tasks from DMControl.
DMControl @ 1M Env. Steps (30 tasks, seeds per task)
Median
IQM
Mean
Optimality Gap
TD-MPC SAC DreamerV3 TD-MPC2
DC-MPC (ours)
0.6
0.7
0.8
0.9
0.60
0.75
0.90
0.5
0.6
0.7
0.8
0.2
0.4
0.5
Normalized Return
Figure 14: DMControl aggregate results DC-MPC generally outperforms TD-MPC2 and Dream- crV3 in DMControl tasks. This due to DC-MPC's strong performance the hard Dog and Humanoid tasks. Error bars represent 95% stratified bootstrap confidence intervals.
appendices continue on next page-
27

---

Published as a conference paper at ICLR 2025
B.7
META-WORLES MANIPULATION RESULTS
Success Rate (86) Success Rate (96) Success Rate (56) Success Rate (56 Success Rate (56) Success Rate (16) Success Rate (18) Success Rate (5) Success Rate (%)
100
Assembly
Basketball
Bin Picking
Box Close
Button Press
Button Press Td
Button Press Td Wall
Button Press Wall
Coffee Button
Coffee Pull
100
Coffee Push
Dial Turn
Disassemble
Door Close
Door Lock
100
Door Open
Door Unlock
Drawer Close
Faucet Close
Faucet Open
100
Hammer
Hand Insert
Handle Press
Handle Press Side
Handle Pull
100
Handle Pull Side
Lever Pull
Peg Insert Side
Peg Unplug Side
Pick Out Of Hole
100
Pick Place
Plate Slide
Plate Slide Back
Plate Slide Back Side
Plate Slide Side
100
Push
Push Wall
Reach Wall
Soccer
Stick Pull
100
Stick Push
Sweep
Sweep Into
Window Close
Window Open
100
1M
2M
1M
2M
1M
2M
1M
2M
1M
2M
Env.
Steps
Env. Steps
Env.
Steps
Env. Steps
Env. Steps
DC-MPC (ours)
DreamerV3
SAC
TD-MPC
TD-MPC2
Figure 15: Meta- World manipulation results DC-MPC performs well across Meta-World tasks. We plot the mean (solid line) and the 95% confidence intervals (shaded) across seeds, where each seed averages over 10 evaluation episodes.
-appendices continue on next page-
28

---

Published as a conference paper at ICLR 2025
World @ 1M Env. Steps (45 tasks, seeds per task)
Median
IQM
Mean
Optimality Gap
TD-MPC SAC DreamerV3 TD-MPC2
DC-MPC (ours)
0.45 0.60 0.75 0.90
0.6
0.45 0.60 0.75
0.90
0.15 0.30 0.45
0.60
Episode Success
Figure 16: Meta-World results DC-MPC performs well in Meta- worad, generally matching TD- MPC2, whilst significantly outperforming DreamerV3 and SAC. Error bars represent 95% stratified bootstrap confidence intervals.
appendices continue on next page-
29

---

Published as a conference paper at ICLR 2025
B.8 MYOSCITE MUSCULOSKELETAL RESULTS
In this section, we evaluate DC-MPC in five musculoskeletal tasks from MyoSuite.
In these experiments, followed Hafner al. (2023); Hansen et al. (2023) and scaled the rewards using symlog()),
symlog(#) sign(x) (x)]][[=] 1).
(14)
This compresses large and small rewards whilst preserving the imput sign as it is symmetric function. Note that we simply transform the rewards with symlog and learn both the reward function and Q-functions using these transformed rewards. We N 1-step returns in Hand Key Turn, Hand Obj Hold and Hand Pen Twirl and we use = 5-step returns Hand Pose and Hand Reach. Hand Pose we also had adjust the temperature from 0.5 0.2. In future work, would be interesting investigate if using X-returns - which uses a weighted-sum of N-step returns can make DC-MPC robust the N-step hyperparameter. Further to this, it would be interesting explore methods for dynamically tuming the MPPI (inverse) temperature 7MPPI-
In Fig. show the training curves for the individual tasks, Fig. 18 then reports aggregate metries at IM environment steps over three random seeds in the five tasks. On average, DC-MPC performs well, generally matching TD-MPC2 1M environment steps and outperforming the other baselines.
Hand Key Turn
Hand Obj Hold
Hand Pen Twirl
Hand Pose
Hand Reach
Success Rate (%)
1.0
0.5
300
1000
50X)
1000
1000
1000
1000
Env. Steps (1e3)
Env.
Steps (1+3)
Env. Steps (1e3)
Env. Steps (1e3)
Env. Steps (1e3)
DC-MPC (ours)
DreamerV3
SAC
TD-MPC
TD-MPC2
Figure 17: MyoSuite training curves We plot the mean (solid line) and the 95% confidence intervals (shaded) across 3 seeds, where each seed averages over 10 evaluation episodes.
MyoSuite @ Env. Steps (5 tasks, seeds per task)
Median
IQM
Mean
Optimality Gap
TD-MPC
SAC
DreamerV3
TD-MPC2
DC-MPC (ours)
0.00 0.25 0.50 0.75 1.00 0.00 0.25 0.50 0.75 1.00 Episode Success
0.25 0.50 0.75 1.00 Episode Success
0.00 0.25 0.50 0.75 Episode Success
Episode Success
Figure 18: MyoSuite results DC-MPC performs similarly to TD-MPC2 in MyoSuite. Error bars represent 95% stratified bootstrap confidence intervals.
appendiccs continue on next page-
30

---

Published as a conference paper at ICLR 2025
B.9 DOES DCWM IMPROVE DREAMERV3?
In this section, we seek to evaluate what happens when we replace DreamerV3's one-hot discrete encoding with the codebook encoding used in DC-MPC. Fig. 19 shows that in the easy Reacher Hard and Walker Walk environments, FSQ (blue) and one-hot (orange) perform similarly. However, in the difficult Dog Run and Humanoid Walk tasks, no discrete encoding can enable DreamerV3 to perform as well as DC-MPC (purple). We hypothesize that DreamerV3's poor performance in the Dog Run and Humanoid Walk tasks results from its decoder struggling to reconstruct the observations.
Dog Run
Humanoid Walk
Reacher Hard
Walker Walk
1000
Episode Return
500
0-
0
1000
1000 Env. Steps (1c3)
1000 Env. Steps (1e3)
200 Env. Steps (1e3)
Env. Steps (1e3)
DC-MPC (ours)
DreamerV3 w/ FSQ
DreamerV3 w/ one-hot
Figure 19: DreamerV3 with FSQ Replacing DreamerV3's one-hot encoding (orange) with DC- MPC's codebook encoding (blue) does not improve performance. Moreover, DreamerV3 is not able to learn in the hard Dog Run and Humanoid Walk tasks and is significantly outperformed by DC-MPC (purple).
Learning to minimize the observation reconstruction error has been widely applied in model-based RL (Sutton & Barto, 2018; Ha & Schmidhuber, 2018; Hafner et al., 2019b), and an observation decoder has been component of many of the most successful RI. algorithms date (Hafner et al., 2023). However, recent work in representation learning for RL (Zhao et al., 2023) and model-baser RL (Hansen et al., 2022) has shown that incorporating a reconstruction term into the representation loss can hurt the performance, as learning to reconstruct the observations is inefficient due to the observations containing irrelevant details that are uncontrollable by the agent and do not affect the task.
To provide thorough analysis of DC-MPC, we include results where we add a reconstruction term to our world model loss in Eq. (8): C. Eo,-DDlle 0.1 à. = (c,),,,
(15)
where hw is learned observation decoder that takes the latent code as the input and outputs the recon- structed observation. The decoder h, a standard MLP. We perform reconstruction at each time step in the horizon. The results in Fig. 20 show that in no environments does reconstruction aid learning, and in some tasks, such as the difficult Dog Run and Humanoid Walk tasks, including the recon- struction term has a significant detrimental effect the performance, and can even prevent learning completely. Our results support the observations of Zhao et al. (2023) and Hansen et al. (2022) about the lack of need for reconstruction loss in continuous control tasks. However, it is worth noting that we weighted all loss terms equally whilst the results in Ma et al. (2024) suggest that the observation reconstruction, temporal consistency, and reward prediction loss terms need to be carefully balanced.
Dog Run
Humanoid Walk
Reacher Hard
Walker Walk
1000
Episode Return
500
0
1000
1000
1000
200 Env. Steps (1e3)
Env. Steps (1e3)
Env. Steps (1e3)
Env. Steps (1e3)
w/o reconstruction (ours)
w/ reconstruction
Figure 20: Reconstruction harms performance Adding observation reconstruction to DC-MPC (blue) harms the performance of DC-MPC across mixture of easy and hard DMControl tasks.
31

---

Published as a conference paper at ICLR 2025
B.10 IMPROVING TD- TD-MPC2 WITH DC-MPC
In this section, we investigate using DC-MPC's latent space inside TD-MPC2, Note that TD-MPC2's latent space is continuous and trained with MSE regression. also uses simplical normalization (SimNorm) to make its latent space bounded. In these experiments, we removed SimNorm and replaced with our discrete and stochastic latent space, and then trained using cross-entropy for the consistency loss. In particular, we made the following changes to the TD-MPC2 codebase: (i) removed SimNorm, (ii) added FSQ to the encoder, (iii) modified the dynamics predict the logits instead of the next latent state, (iv) modified the dynamics to use ST Gumbel-softma sampling for multi-step predictions during training and our weighted average approach during planning, and (v) changed the world model's loss coefficients for consistency, value, and, reward, to all be 1.
In Fig. we report aggregate metrics over 3 random seeds 10 DMControl tasks and 10 Meta-World tasks. Fig. 6 (left) shows the IQM and optimality gap 1M environment steps over the 20 tasks. shows that adding DC-MPC's discrete and stochastic latent space to TD-MPC2 offers some improvement. Fig. 6 (right) shows the aggregate training curves (IQM over 10 tasks) for DMControl and Meta- World, respectively. The results show that using DCWM M inside TD-MPC2 offers some benefits in the 10 DMControl tasks, whilst in the 10 seems about equal. This suggests that, in the context of continuous control, discrete and stochastic latent spaces advantageous for world models. This is interesting result which we believe motivates further research into discrete and stochastic latent spaces for world models.
tasks, the performance of all methods
appendices continue on next page-
32

---

Published as a conference paper at ICLR 2025
C IMPLEMENTATION DETAILS
Architecture We implemented DC-MPC with PyTorch (Paszke et al., 2019) and used the AdamW optimizer (Kingma & Ba, 2015) for training the models. All components (encoder, dynamics, reward, actor and critic) are implemented as MLPs. Following Hansen et al. (2023) we let all intermediate layers be linear layers followed by LayerNorm (Ba al., 2016). We Mish activation functions throughout. Below we summarize the DC-MPC architecture for our base model.
DCMPC
(model): WorldModel
[_sg] FSC Levell=((5, 3])
L ModuleDict
(statel: Sequentiel(
20t_ act*Mishi (0): (in out aeeerure-llattttt..mm.a..
(1): Linear
(trrans:: Sequential
(0) NormedLineat Seatures (latent_dimmnaat cimmmaa_chammells fact_dimg ous_teatures=5.2, act-Mist)
(1): NermedLinear (in iattuuae out fet_ffestres==12, act-Mish)
(2): Linear (in_ar(in_fettee (in_tetrrsss, out_featuressiatee direessokiiooeeook......
rewaard Sequentiall
chanaels) act_dim, ou: fes_ act-Mist) (1): Normedti near Hin settuurrer-,,, out features=b12, act-Mish)
10) NermedLinear iin__eatures-eeeeeeeeeeeeeeeeeee (lsteett_ittttttt_diiineeeee chammll
(2): Linear (el__ussrees=512, fet_feeuu))==-)
(pt): Sequentiall
Hormedtineat tnnuuueesrssll www.sseennndiiinuumnmm Sinsenum_channee fealuurese,,-5,2, act-Mishl
(1): NormedLineaz featuueetre-51, out fet__feerrres511, act-Mish)
(2): Linear (LA___ettttttt out out__eatuueeeeet__dr
Les): vectorized ModuleList(
(0-4): 5 Soquential
10) Normedbineaat iin "eaturee (Iatteeeeelaaaateee (latent_cinne
fact_dim, dim, out_teatures-5:2, act-Maat) NozmedLineat Sic. feat act-Mish)
fact
(2): (inear fet_ueeaue)ss
(pi_ter): Sequentiall
Normedf/ near (In_ (in_
fin shannelll, featureese, act
(1): Normedbinead Tettueeees,ss1,2, out fes_urrsr,, act-Miat)
(La_feturress yut_ffeauurenact_din
(0a_ : Vectari ModuleMiat
(0-4): 5 Sequentiel(
(0) Normeddiineariie faaturrrs chanmm__caaals)aaa tact tact_dim, out cor_feotture=5512, acteMish)
NormedLinear tin.etietnntnnnntututteeaea rit_reatu,,,--,,, reatures act-Mish)
(2) Linear Linear(in_Ee (en__eeeture=511, out (et__eeaures=")
where obs_ is the dimensionality of the observation space, act_dim is the dimension- ality of the action space, latent_din is the number of the latent dimensions d (default 512), num_channels is the number of channels per latent dimension b (default 2), and codebook_size is the codebook size |C] (default 15).
Statistical significance We used five seeds for DC-MPC and three seeds for TD-MPC2/Dream- erV3/SAC/TD-MPCC in the main figures, at least three seeds for all ablations, and plotted the 95 % confidence intervals as the shaded area, which corresponds to approximately two standard errors of the mean. However, in Figs. 3 and 6 follow Agarwal et al. (2021) and plot the interquartile mean (IQM) with the shaded area representing 95% stratified bootstrap confidence intervals.
Hardware We used NVIDIA A100s and AMD Instinct MI250X GPUs to our experiments have been run on single GPU with a single-digit number of CPU workers.
experiments. All
Open-source code For full details of the implementation, model architectures, and training, please check the code, which is available the submitted supplementary material and available on github at https: / //g/ thub. com/a a danscannell htw.neeell//mm...eee/
appendices continue on next page-
33

---

Published as a conference paper at ICLR 2025
Hyperparameters Table 1 lists all of the hyperparameters for training DC-MPC which used for the main experiments and the ablations.
Table 1: DC-MPC hyperparameters We kept most hyperparameters fixed across tasks. However, task-specific exploration noise schedules and Nsstep returns.
HYPERPARAMETER
VALUE
DESCRIPTION
TRAINING
REPEAT
(1 IN MYOSUITE)
MAX EPISODE LENGTH (T)
500 IN DMCONTROL
ACTION REPEAT MAKES THIS 1000 ACTION REPEAT MAKES THIS 200
100 IN META META-WORLD
100 IN MYOSUITE
NUM. EPISODES
RANDOM EPISODES (NNANDON (rssones)
10
NUM. RANDOM EPISODES AT START
MPPI PLANNING
PLANNING HORIZON
PLANNING
(J)
POPULATION SIZE (N₂)
512
PRIOR POPULATION SIZE (N.)
24
NUM. POLICY SAMPLES TO WARM START
NUMBER ELITES (AC)
MINIMUM
0.05
MAXIMUM STD
2
(INVERSE) TEMPERATURE (ТмРР))
0.5
ACTOR FREQ.
UPDATE
BATCH SIZE
512
BUPFER SIZE
10
DISCOUNT FACTOR (7)
0.99
EXPLORATION
Linear(L.0, 0.1, (EASY)
DMCONTROL
Linear(1. [1.0, 0.1, 150) (MEDIUM)
DMCONTROL.
Linear( [.0,, 0..0, 0.1, 500) (HARD)
DMCONTROL
Linear(1.0, 0.1,
META WORLD MYOSUITE
LEARNING RATE
10
MLP DIMS
(512,512)
FOR ACTOR/CRTTTCOYYYMICCEEEEE
MOMENTUM COEF, (7)
0.005
NUM. Q- FINCTIION (N₂)
NUM. Q-FUNCTIONS
NOISE CLIP (c)
0.3
N-STEP TD
OR 3 IN DMCONTROL
IN META- WORLD
5 MYOSUITE
POLICY NOISE
0.2
URDATE-TO-DATA (UTD) RATIO
WORLD MODEL
DISCOUNT FACTOR (7)
ENCODER
RATE
10-
ENCODER MI.P DIMS
[256]
FSQ LEVELS
[.
GIVES |C| = 15 R 2
HORIZON (H)
FOR WORLD MODEL TRAINING
LATENT
512
1024 (HUMANOID/DOG)
appendiccs continue on next page-
34

---

Published as a conference paper at ICLR 2025
D BASELINES
In this section, we provide further details of the baselines we compare against.
DreamerV3 (Hafner et al., 2023) is reinforcement learning algorithm that uses world model to predict outcomes, a critic to judge their value, and an actor to choose actions to maximize value. It uses symlog loss for training and operates on model states from imagination data. The critic is a categorical distribution with exponentially spaced bins, and the actor is trained with entropy regularization and return normalization. The world model is only used for training and there is no decision-time planning. In contrast, DC-MPC learns a deterministic encoder with a discrete latent space and stochastic dynamics in the world model. We report the results of DreamerV3 from the TD-MPC2 official repository 2.
Temporal Difference Model Predictive Control (TD-MPC2, Hansen et al. (2023)) is a decoder-free model-based reinforcement learning algorithm with focus on scalability and sample efficiency. It includes an encoder, latent transition dynamics, a reward predictor, a terminal value (critic), and a policy prior (actor). In contrast to DreamerV3, it utilizes deterministic encoder and transition dynamics implemented with MIPs, layer normalization (Ba et al., 2016) and Mish (Misra, 2019) activations. To avoid exploding gradients and representation collapse, the latent space is normalized with projection followed by softmax operation. All components except the policy prior are trained jointly based on predicting the latent embedding, reward prediction, and value prediction, while reward and value predictions are based on discrete regression log-tran taggsrrasforeed space. Similarly, we use a deterministic encoder, but we train the transition dynamics with a cross-entropy loss function, which considers multi-modality and uncertainties, and we decouple representation learning from value learning. We report the results from the TD-MPC2 official repository
Temporal Difference Model Predictive Control (TD-MPC, Hansen et al. (2022)) is the first version of TD-MPC2. It is also decoder-free model-based RL algorithm consisting of an encoder, latent transition dynamics, reward predictor, terminal value (critic), and policy prior (actor). In contrast to TD-MPC2, it does not apply simplical normalization (SimNorm) to its latent state, it trains the reward and value prediction using the MSE loss instead of the cross- entropy loss, and it uses SAC as the underlying RL algorithm. We refer the reader to the TD- MPC paper for further details. We report the results from the TD-MPC2 official repository Soft Actor-Critic (SAC, Haarnoja et al. (2018) is an off- off-policy model-free RL algorithm based on the maximum entropy RL framework. That is, it attempts to succeed at the task whilst acting as randomly as possible. worth highlighting that TD-MPC2 uses SAC as it's under- lying model-free RI. algorithm. We report the results from the TD-MPC2 official repository 5,
appendices continue on next page-
35

---

Published as a conference paper at ICLR 2025
E TASKS
We evaluate our method in 30 tasks from the DeepMind Control suite (Tassa et al., 2018), 45 tasks from Keea-World (Yu et al., 2019) and tasks from MyoSuite (Vittorio al., 2022). Tables 2 provide details of the environments we used, including the dimensionality of the observation and action spaces.
Table DMControl We consider a total of 30 continuous control tasks from DMControl.
TASK
OBSERVATION DIM
ACTION DIM SPARSE?
ACROBOT SWINGUP
N
CARTPOLE BALANCE
5
N
CARPOLE BALANCE SPARSE
5
Y
CARTPOLE SWINGUP
CARTPOLE SWINGUP SPARSE
CHEETAH RUN
17
N
CUP CATCH
8
Y
CUP SPIN
8
N
DOG RUN
223
DOG STAND
223
38
N
DOG TROT
223
38
N
DOG WALK
223
38
N
FINGER SPIN
FINGER TURN EASY
FINGER TURN HARD
12
Y
FISH SWIM
24
N
HOPPER HOP
15
N
HOPPER STAND
HUMANOID RUN
67
N
HUMANOID STAND
67
24
N
HUMANOID WALK
67
24
N
PENDULUM SPIN
3
PENDULUM SWINGUP
QUADRUPED RUN
78
12
N
QUADRUPED WALK
78
12
N
REACHER EASY
6
Y
Y
REACHER HARD
WALKER RUN
24
N
WALKER STAND
24
N
WALKER WALK
24
6
N
appendiccs continue on next page-
36

---

Published as a conference paper at ICLR 2025
Table 3: Meta-Worid We consider total of 45 continuous control tasks from Meta-World. This benchmark is designed for multitask research so all tasks share similar embodiment, observation space, and action space.
TASK
OBSERVATION DIM
ACTION DIM
SPARSE?
ASSEMBLY
BASKETBALL
39
4
BIN PICKING
39
4
N
BOX CLOSE
39
4
N
BUTTON PRESS
39
4
N
BUTTON PRESS TOPDOWN
BUTTON PRESS TOPDOWN WALL
39
4
N
BUTTON PRESS WALL
39
4
N
COFFEE BUTTON
39
4
COFFEE PUSH
COFFEE PULL
39
N
DIAL TURN
39
4
N
DISASSEMBLE
39
4
DOOR CLOSE
39
4
N
DOOR LOCK
DOOR OPEN
39
4
N
DOOR UNLOCK
39
4
DRAWER CLOSE
39
4
FAUCET CLOSE
39
FAUCET OPEN
HAMMER
39
4
N
HAND INSERT
39
4
N
HANDLE PRESS
39
4
HANDLE PRESS SIDE
HANDLE PULL
39
4
N
HANDLE PULL SIDE
39
4
LEVER PULL
39
4
PEG INSERT SIDE
39
4
PEG UNPLUG SIDE
N
PICK OUT OF HOLE
39
4
N
PICK PLACE
39
4
PLATE SLIDE
39
4
N
PLATE SLIDE BACK
39
N
PLATE SLIDE BACK SIDE
N
PLATE SLIDE SIDE
39
4
N
PUSH
39
4
INN
PUSH WALL
39
4
N
REACH WALL
N
SOCCER
39
4
N
STICK PULL
39
4
N
STICK PUSH
39
4
SWEEP
N
SWEEP INTO
WINDOW CLOSE
39
4
N
WINDOW OPEN
39
4
N
Table 4: MyoSuite We consider a total of continuous control tasks from MyoSuite. This benchmark is designed for digh-dimensional muscoloskeletal motor control which involves complex object manipulation with dexterous hand.
TASK
OBSERVATION DIM
ACTION DIM
SPARSE?
KEY TURN
93
39
N
OBJECT HOLD
91
39
N
PEN TWIRL
83
39
N
POSE
108
N
REACH
115
37

---

Published as a conference paper at ICLR 2025
Figure 21: Tasks visualizations Visualization of the DMControl, Meta-Workke World, and MyoSuite tasks used throughout the paper.
38


Reviews:
[{"review_id": "mQBX8LvsIm", "reviewer": "Reviewer_RvJB", "summary": "This paper explores the use of discrete latent spaces for continuous control with world models. The authors introduce DCWM: Discrete Codebook World Model, a model-based RL method which surpasses recent state-of-the-art algorithms. it’s a well written paper in general.", "questions": "I think a comparison of the method with other approaches that are not based on the same embeddings ideas can be added.", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors demonstrate that quantized discrete codebook encodings are more effective representations for continuous control, compared to alternative encodings, such as one-hot and label-based encodings. \n\n2. The authors  introduce DCWM: Discrete Codebook World Model, a model-based RL method which surpasses recent state-of-the-art algorithms, including TD-MPC2 and DreamerV3, on continuous control benchmarks.", "weaknesses": "What it misses I think  is  a comparison of the method with other approaches that are not based on the same embeddings ideas. For example I think MAMBA and Hungry hungry hippos (H3) apply to similar scenarios.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mDu7ri59H0", "reviewer": "Reviewer_BJj9", "summary": "The authors propose a new architecture for world modeling for continuous control that utilizes discrete representations. Rather than use a learned codebook or one-hot encodings, they utilize finite scalar quantization (FSQ) which simplifies the loss function and stabilizes the learning according to their claims. They build upon the TD-MPC2 algorithm which utilizes model predictive path integral (MPPI) to search the world model for high value actions.\n\nThey test their algorithm in a number of high dimensional continuous control tasks including Meta-World manipulation and DeepMind Dog and Humanoid tasks, showing that their algorithm achieves higher success and learns with fewer samples compared to baselines. In addition to comparison against baselines, they also provide experiments supporting their choice of discrete encodings, showing that they can achieve similar performance to one-hot encodings while increasing training efficiency.", "questions": "What does the algorithm performance look like without using REDQ?\n\nWhat are the computational costs or benefits to this method over TD-MPC in terms of runtime?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well-written and explains all necessary background material in sufficient detail to understand the algorithm. While none of the individual components are novel, the experimental evaluation provides evidence that the proposed algorithmic details are critical for the success of model-based learning.", "weaknesses": "The one question that I have relates to the use of REDQ to reduce bias in the TD learning. Since the baseline does not seem to use this method, it begs the question of how much of the demonstrated performance is due to this component. It would be nice to see how DCWM compares to a version that uses the standard one or two Q functions to make sure the performance gap is in fact due to the choice of latent spaces.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Sn4UjBolsV", "reviewer": "Reviewer_Nctn", "summary": "The authors contribute a novel method, called DCWM, for model-based RL on continuous control tasks. Specifically, building on the works of Dreamer V3 and TD-MPC2 (prior model-based continuous control methods), the paper sheds light on the impact of how these methods encode latent state (different ways of discrete encoding vs. continuous), how that state is trained (a categorical loss or a regressive loss), and the impact of stochasticity in modeling dynamics. The paper conducts experiments and ablations on Deepmind Control Suite and Meta World to identify which components are most impactful on performance, and combines them to present the aforementioned method, while comparing to baselines from literature. The authors conclude that learning a discrete latent state (specifically one encoded by codebook quantization over other methods), and training with a cross entropy objective results in improved performance.", "questions": "Question:\n- I imagine that taking the expected code might sometimes result in an invalid state, as discussed here:\n\n>  Whilst the expected value of\na discrete variable does not necessarily take a valid discrete value, we find it effective in our setting.\n\nIs that simply then fed into the dynamics model (after normalization etc.) and the reward model MLPs?\n- Relatedly, when you normalize the code to [-1, 1], you claim it improves performance. Is there a reference for this or was this empirically observed? Would you be able to speak to why this is required? (One thought I had was it normalizes the different quantization levels $L_i$, i.e. the two channels would otherwise have values ranging from [-5, 5] and [-3, 3]).\n\nNits (not affecting score):\n- line 217, formatting for `i-th` and `d-th` is inconsistent, and should use latex\n- line 352 typo in 'warms starts'\n- line 465, missing 'and' before models in 'entropy), **and** models'", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The authors present a strong paper, with excellent methodology and solid experimental results. They compare with two SOTA baselines, Dreamer V3 and TD-MPC2, and demonstrate that their method indeed outperforms or matches the baseline performance. They conduct additional experiments (sections 4.3, 4.4 and 4.5) to understand how different variables affect performance, which enhances the paper's contributions to bettering understanding latent spaces in world model learning. The paper's clarity is also very good, with a great presentation of background and method. The authors provide sufficient detail, such as the discussion around the fixed FSQ codebook illustration in Figure 2 or the extensive documentation of various tricks used to improve performance (lines 309, 352, 355, etc.). Additionally, the supplementary material is strong, with detailed extra information on method, environment, architecture, etc, and addition ablation experiments. The results contributed are significant, not only advancing the state of the art for model-based continuous control (to my knowledge), but also contributing to the field's understanding of the impacts of latent space encodings, losses, and stochasticity in modeling transition dynamics. I would argue that the method merits originality as well, since it cleverly combines prior work (the consistency approach from TD-MPC2 and discrete encodings/cross entropy objective from Dreamer V3) and introduces novel elements as well (codebook quantization). \n\nOverall, this is a strong contribution to the field, and I recommend acceptance.", "weaknesses": "The paper is free of major flaws: it has thoroughly conducted experiments and is clear and well-written. However, I have identified some nits and clarifying questions listed in the section below that would be minor improvements.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hnOl24DWZY", "reviewer": "Reviewer_hXNL", "summary": "This work proposes a model-based RL method for continuous control tasks. Compared with TD-MPC2, the latent space is represented as a discrete codebook. The proposed modifications lead to some improvements over TD-MPC2.", "questions": "1. The authors kept most hyperparameters fixed across all tasks, what parameters will affect the performance of some special tasks?\n2. Is the codebook a direct application of quantization-aware training (QAT)?\n3. How to determine the weights of the expected code? Is this a hyperparameter of the algorithm?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well organized and easy to follow. \n2. The description of the contribution to the article is simple and concise. \n3. The ablation experiment is sufficient and persuasive.", "weaknesses": "1. The test for multi-tasking learning is missing. It would be great if the authors could consider extend their studies to multi-task learning and study the differences in codebooks between different control tasks. \n2. In some Meta-World tasks, the performance of the proposed method is comparable to TD-MPC2, and the variance is larger. \n3. The stochastic transition dynamics in the latent space could be better described and explained.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kLYEcLJnW1", "reviewer": "Reviewer_HZL5", "summary": "This paper proposes a discrete latent space approach for world models in continuous control, utilizing FSQ within a MBRL algorithm, DCWM, following TD-MPC2. The experiments, conducted on DMC and Metaworld, demonstrate improvements over state-of-the-art methods, including TD-MPC2 and DreamerV3. Additionally, the authors analyze the effectiveness of discrete vs. continuous latents, classification vs. regression objectives, and deterministic vs. stochastic transitions.", "questions": "On clarifcation (major):\n\n1. The claim in the abstract that discrete latent spaces typically underperform in continuous control tasks needs specification. First of all, to my knowledge, TD-MPC2 only conducts experiments on state-based tasks, lacking results on pixel-based continuous control, while DreamerV3 does include these. Second, it is unfair to compare methods with (TD-MPC) and without inference-time planning (Dreamer), as in the MuZero paper we can see that this makes a great performance boost. I think this work can naturally motivate itself from TD-MPC2 alone instead of comparing DreamerV3 and TD-MPC2 without head-to-head experiments.\n2. A preliminary section clarifying different discrete encodings is necessary; I find the distinction between label encoding and one-hot encoding unclear for the first-time reading. The methods section extensively describes FSQ, which is not originally proposed here.\n3. Also, the methods section contains many elements that directly follow TD-MPC2, which may confuse non-expert readers regarding what is novel. Clear differentiation is needed. For example, TD-MPC2 exactly uses an ensemble of five critics with randomly sampled two for bootstrap, but this is not explicitly mentioned by the authors in Line 277.\n4. Algorithm 1 appears to be a standard procedure in model-based RL and could be moved to the appendix.\n\nOn clarifcation (minor):\n\n1. Can you elaborate on what you mean by \"codebook encoding captures similarity between observations\" and how this contrasts with one-hot encodings in Line 248?\n2. I did not understand why \"we did not conﬁgure the transition model to accept the label and one-hot encodings as this resulted in the agent being unable to learn.\" Can you explain more?\n3. I did not understand the sentence \"our codebook is conﬁgured with b = 2 channels and the label encoding incorrectly assumes an ordinal structure through both.\" How does the hyperparameter $b$ in your proposed codebook encoding affect label encoding?\n4. In Line 514, why is $|C|=2^4$ for $L=[5,3]$, when it is claimed $|C|=\\prod L_i$ in Line 224.\n\nOn experiments: \n\n1. Could you provide additional experiments on more tasks, such as ManiSkill2/Myosuite, which TD-MPC2 used? It would also be beneficial to include experiments on more base algorithms, such as DreamerV3 (see question 2).\n2. The most direct evidence that the proposed discrete latents outperform one-hot and continuous latents would be to replace the latent space in DreamerV3 and TD-MPC2 with your proposed approach. Why wasn't this head-to-head comparison conducted?\n3. Have you compared FSQ and VQ (referred to as dictionary learning in Line 206)?\n4. Codebook encoding only outperforms one-hot encoding in one task in Figure 6, yet it is claimed to be more computationally efficient in Line 485. Is there any numerical analysis to support this claim?\n\nOverall, I find this paper interesting. If the authors adequately address my concerns, I would be willing to significantly improve my rating.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The investigation into latent space design is valuable for the model-based RL community.\n2. The experiments are informative, offering insights across multiple dimensions of the design.", "weaknesses": "1. A significant concern is the need for clarification; some claims may be overstated. See questions below. \n2. For an impactful empirical study on algorithmic design, the scale of experiments—particularly in terms of the number of environments and tasks—seems insufficient. The experimental design is also somewhat confusing; see questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5l9qMrD8KR", "reviewer": "Reviewer_iQtr", "summary": "This paper explores the application of discrete latent spaces for continuous control in reinforcement learning (RL), an innovative perspective other than the typical use of continuous latent spaces. The authors introduce the Discrete Codebook World Model (DCWM), which utilizes quantized discrete codebook encodings to represent latent states. The DCWM is demonstrated to surpass recent state-of-the-art RL algorithms, including TD-MPC2 and DreamerV3, especially in those tasks with high-dimensional action spaces.", "questions": "1. What's your idea about if FSQ-based discretization could be used for action space discretization? (like action sampling) If yes, how will you do that?\n2. Although it could be more computationally challenging, could you please provide several visual input results? You can ignore the most difficult tasks like humanoid walk/run.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. Although similar ideas like the token-based world model were recently published, the proposed method seems computationally efficient, benefiting from the 'Finite Scalar Quantization' paper.\n2. The paper first made the discrete latent space method work in continuous control tasks.\n3. The paper writing is clear and easy to understand.", "weaknesses": "1. This paper lacks comparisons with the more advanced baselines like EfficientZero-V2.\n2. The authors didn't test the proposed method with visual inputs, which made this paper not that strong.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)