Constrained Latent Action Policies for Model-Based Offline Reinforcement Learning
Authors: Marvin Alles
Philip Becker-Ehmck
Patrick van der Smagt
Maximilian Karl
Venue: NeurIPS 2024 poster (2024)
Keywords: Offline Reinforcement Learning
Model-Based Reinforcement Learning
Latent Action
Constrained Policy
Generative Model
Applied Reinforcement Learning
Primary area: reinforcement_learning

Abstract: In offline reinforcement learning, a policy is learned using a static dataset in the absence of costly feedback from the environment. In contrast to the online setting, only using static datasets poses additional challenges, such as policies generating out-of-distribution samples. Model-based offline reinforcement learning methods try to overcome these by learning a model of the underlying dynamics of the environment and using it to guide policy search. It is beneficial but, with limited datasets, errors in the model and the issue of value overestimation among out-of-distribution states can worsen performance. Current model-based methods apply some notion of conservatism to the Bellman update, often implemented using uncertainty estimation derived from model ensembles. In this paper, we propose Constrained Latent Action Policies (C-LAP) which learns a generative model of the joint distribution of observations and actions. We cast policy learning as a constrained objective to always stay within the support of the latent action distribution, and use the generative capabilities of the model to impose an implicit constraint on the generated actions. Thereby eliminating the need to use additional uncertainty penalties on the Bellman update and significantly decreasing the number of gradient steps required to learn a policy. We empirically evaluate C-LAP on the D4RL and V-D4RL benchmark, and show that C-LAP is competitive to state-of-the-art methods, especially outperforming on datasets with visual observations.

Constrained Latent Action Policies for
Model-Based Offline Reinforcement Learning
Marvin Alles!
Philip Becker-Ehmck' Patrick van der Smagt Maximilian Karl
Machine Learning Research Lab, Volkswagen Group "Technical University of Munich
Eötvös Loránd University Budapest
fmarvin.all mhilip.backer-ehmck. maximi lan karl)
Abstract
In offline reinforcement learning, a policy learned using static dataset in the absence of costly feedback from the environment. In contrast to the online setting, only using static datasets poses additional challenges, such as policies generating out-of-distribution samples. Model-based offline reinforcement learning methods try to overcome these by learning a model of the underlying dynamics of the environment and using guide policy search. It is beneficial but, with limited datasets, errors the model and the issue of value overestimation among out-of-distribution states can worsen performance. Current model-based methods apply some notion of conservatism to the Bellman update, often implemented using uncertainty estimation derived from model ensembles. In this paper, we propose Constrained Latent Action Policies (C-LAP) which learns generative model of the joint distribution of observations and actions. We cast policy learning as constrained objective to always stay within the support of the latent action distribution, and use the generative capabilities of the model to impose implicit constraint on the generated actions. Thereby eliminating the need to use additional uncertainty penalties on the Bellman update and significantly decreasing the number of gradient steps required to learn policy. We empirically evaluate C-LAP on the D4RL and V-D4RL benchmark, and show that C-LAP is competitive state-of- the-art methods, especially outperforming datasets with visual observations
Introduction
Deep-learning methods are widely used in applications around computer vision and natural language processing, related the fact that datasets are abundant. But when used for control of physical systems, in particular with reinforcement learning, obtaining data involves interaction with an environment. Learning through trial-and-error and extensive exploration of an environment can be done in simulation, but hard to achieve in real world scenarios [1, 2 S Offline reinforcement learning tries to solve this by using pre-collected datasets eliminating costly and unsafe training in real-world environments 4 6).
Using online reinforcement learning methods in an offline setting often fails. A key issue is the distributional shift: the state-action distribution of the offline dataset, driven by the behavior policy, differs from the distribution generated by a learned policy. This leads to actions being inferred for states outside the training distribution. Therefore, value-based methods are prone to overestimating values due to evaluating policies on distri--iitiiittionn states. This leads to poor performance and unstable training because of bootstrapping 66] 8 Offline reinforcement learning methods address
"Corresponding author.
'Code is available at: httpp:///gttbu...co//vvvvaallll
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

12
2
V3
T
91
01
09
02
93
03
91
01
(a) Offline Model training
(b) Policy training
(c) Environment interaction
Figure I: Overview of C-LAP. (a) The model is trained offline. It is encoding observations or and actions a, to latent states (gray circle) and latent actions u, (green circle), and decoding them thereafter. Furthermore it is predicting rewards År. (b) The policy is learned in the latent action space, but constrained the support of the latent action prior, and uses the generative capabilities of the action decoder. Gradients are computed by back-proy paok-propggating estimated values V; and rewards Ae through the imagined trajectories. (c) The policy is used in the real world, again using the generative action decoder
this issue with different approaches and can be categorized into model-free and model-based methods, similar to online reinforcement learning.
Model-free offline reinforcement learning usually follows one of the following paradigms: constrain the learned policy to the behavior policy [0, 71: or introduce some kind of conservatism to the Bellman update [12
3 14 8]. Model- reinforcement learning methods transform the offline to an online learning setting: They approximate the system dynamics and try to resolve the evaluation of out-of-distributiin states by using the generalization capabilities of the model and generating additional samples. But as the training distribution fixed, the estimation capabilities of the model are limited. Therefore, these model-based methods also rely conservative modification to the Bellman update as a measure to counteract value overestimation which is mostly achieved through uncertainty penalties [15 16
20 21]. Apart from the typical approach of using an auto-regressive model to estimate the dynamics, other model-hased methods treat the objective as trajectory modeling (22 23 24 These methods aim to combine decision making and dynamics modeling into one objective. Instead of learning policy, they sample from the learned trajectory model for planning. We will refer to the first kind of methods, which learn a dynamics model train a policy, as model-based reinforcement learning.
17
18
19
We aim to solve the problem of value overestimation in model- reinforcement learning by jointly modeling action and state distributions, without the need for uncertainty penalties or changes to the Bellman update. Instead of learning conditional dynamics model p(s a), we estimate the joint state-action distribution p(s, a). This is similar to methods that frame offline reinforcement learning as trajectory modeling, but the objective as a generative model of the joint distribution states and actions, we create an implicit constraint on the generated actions, similar to 10 25). The goal of this approach is address the shift in the entire distribution, rather than looking at out-of-distribution actions and states separately. We achieve this using a recurrent state-space model with a latent action space, which we call the recurrent latent action state-space model. Using a latent action space allows us to learn a policy that uses the latent action prior as an inductive bias. This approach keeps the policy close the original data while allowing it to change when needed, which makes learning the policy much faster. To achieve this, we treat policy optimization as constrained optimization problem, similar to enforcing support constraint [9] We provide high level overview of our method in T
an auto-regressive model and still learn policy. By formulating
'Emoji graphics used in Figure 1 Mree licensed under CC-BY 4.0 by Twemoji.

---

Overall, we summarize our contribution as follows:
We introduce latent action state-space models for model-based offline reinforcement learning, treating it auto-regressive generative modeling of the joint distribution of states and actions.
We formulate policy optimization as a constrained optimization problem, using the latent action space to generate actions within the support of the dataset's action distribution and jump-start policy learning by using the generative action decoder.
We evaluate our approach on one benchmark with image observations (V-D4RI. [20] and another one with low- oow-dimensional feature observations (D4RL (26)).
We evaluate the effect of our approach on value overestimation.
2 Preliminaries
We consider a partial observable Markov decision process (POMDP) defined by M (S,A R, Ω 7) with S as state space, A action space, O as observation space, 8 S as state, E A as action, o E O as observation, : S x A + S as transition function, R: + I as reward function, f2 find a policy x : O + A that maximizes the expected discounted sum of rewards EI [27]
O emission function and y (0, 1] as discount factor. The goal is to
In online reinforcement learning, an agent iteratively interacts with the environment M and optimizes its policy æ. offline reinforcement learning, however, the agent cannot interact with the environment and must refine the policy using fixed dataset D {(OL.T a1.7, ri.r))==) Therefore, the agent must understand the environment using limited data to ensure the policy maximizes the expected discounted sum of rewards when deployed (6. Auto-regressive model-based offline reinforcement learning tries to learn a parametric function to estimate the transition dynamics T. The transition dynamics model is then used to generate additional trajectories which can be used to train policy. The majority of these approaches learn a dynamics model directly in observation space (o: 0t-1, Qt-1)
[16
18
19 28 while others latent dynamics model To(s; 8t-1, a,-1)
20 21.
3
Constrained Latent Action Policies
A main issue in offline reinforcement learning is value overestimation, which we address by ensuring the actions generated by the policy stay within the dataset's action distribution. Unlike previous model- based methods, we formulate the learning objective states and actions. We do this by combining latent action space with a latent dynamics model. Next, the generative properties of the action space to constrain the policy to the dataset's action distribution. A general outline of our method, Constrained Latent Action Policies (C-LAP) ), is shown in Appendix B It starts with learning generative model, followed by actor-critic agent training imagined trajectories, similar to the methods in 29 30, 20 21.
generative model of the joint distribution of
ha
ha
hg
83
Figure 2: Recurrent latent action state-space model. The generative process is shown by solid lines and inference by dashed lines. Stochastic variables are denoted by circles and deterministic variables by rectangles.
Generative model Model-based offline reinforcement learning requires learning a model that is accurate in areas with low data coverage but also generalizes well. Therefore, it's crucial to balance

---

staying within the dataset's distribution and generalizing to unseen states. We propose generative model
p(o1:T (1:T) p(OLT a1. $1.7, u1.7 (p($1:) u1r)) ds du.
(1)
that jointly models the observation and action distribution of static dataset D {(11:, 01:T, >=rr) by using latent states along with latent actions ut. Unlike other model- based offline reinforcement learning methods that learn conditional model p(o1: (1:r) and rely on ensemble based uncertainty penalties on the Bellman update to generate trajectories within the data distribution 221 18. 16. [7], our approach uses latent action space to impose an additional implicit constraint. By implementing a policy in the latent action space, generated actions will stay within the dataset's action distribution, thus enabling generalization within the limits of the learned model [TO]. We empirically validate this claim in Appendix (F)
To obtain a state space model with Markovian assumptions on the latent states 9 we impose the following structure:
p(orTr, ai. S1T, U1:T) II plor st)p(ae s2, uc),
(2)
T p(81TT (1r) = II p(ue ())p((())) 44-1)
(3)
We implement the probabilistic model modifying the design of recurrent state-space model 31.. Thus, the latent dynamics model p(s; st-1. is based on the deterministic transi- tion f(ht-1, st-1, az-1) using the latent action decoder po 8t-1, Ut-1) to generate actions. In the following, we mostly omit deterministic states hy for notational brevity. The resulting recurrent latent action state-space model is shown in Figure 2 and consists of the following components, specifically
latent state prior latent action prior
pe(st Sc-1, uc-1)
Po ss),
observation decoder and action decoder
pe(O:
st),
po(a;
St, ut).
The latent state prior predicts the next latent state 8t given the previous latent state 8:-1 and action U4-1 using the deterministic transition and the action decoder. The latent action prior predicts latent actions u4 given latent state $y. Latent states well and as latent actions decoded using their respective decoder. Similar to [25 actions are reconstructed given latent state and latent action.
Directly maximizing the marginal likelihood is intractable, hence we maximize the evidence lower bound (ELBO) on the log-likelihood log p(O1.T, 01:r) instead. approximate the true posterior, we introduce
latent state posterior
g0(8e St-1,
and latent action posterior
(o(u az))
as inference models. The latent state posterior encodes observations o, to latent states 8t by using the deterministic transition. The latent action posterior encodes actions a; to latent actions u, conditioned on latent states 8t. All parameters of the generative model are indicated by and parameters of the inference model by Φ.
We derive the ELBO
T
log Dogp(o1r,a111
Es Bongo
log 8t)
+ log po (1: 88, uz)
observation reconstraction
action recomstruction
(L(go(51 St-1, 0t-1, || Pe(st 1t-1, U2-1))
(4)
observation consistency
-KL(go(u (t) az) |l (o ))] -LELBO(01T,44:))
action consistency

---

which can be organized into individual terms for reconstruction and consistency of actions and observations. The derivation can be found in Appendix
Maximizing the objective enables us to learn a model which can generate trajectories close to the data distribution DD by sampling from both priors. As want to use the model to learn policy via latent imagination, we add a reward pr(rx st) and termination po(tz st) model. Hence, the complete model training objective
LOOLT ai.)) LELRO(OUT,
Esg B Log 8c)) + (pg sx)]]. (5)
Prior po(u, sx)
Constrained Policy g(xy(u, po(u, sx))
Policy xyy(u, sx)
E07g
Figure 3: Policy constraint through explicit parametrization by using linear transformation g the latent action prior pe pu(us (t) and the bounded policy (2 st) E [-1. 1]. The generated actions
a, p(a, st.g (u, st), Po(ue st)) are implicitly constrained to the data distribution.
Constrained latent action policy We use the sequence model to generate imagined trajectories and use actor-critic approach train the policy. To predict state values we learn value model vg ($x) alongside the policy. Therefore we use the n-step return of a state
& (#t) = E Bee SeNTR
Σ nr + h (e(())
with = min(7 + k 1 + H)
(6)
Polices trained on trajectories generated by model are prone to end up with degrading performance if the model only has access to limited data distribution, as in the case of offline reinforcement learning. Compounding modeling errors and value overestimation of edge-of-reach states [19] are reasons for the decline. Since we train a generative action model, generated actions are implicitly constrained to the datasets action distribution by sampling from the action decoder d; po [a St, u.). Hence, states outside the datasets observation distribution are hard to reach and our approach is resilient to value overestimation of edge-of-reach states. Compounding modeling eπrоrs are still source of diminishing performance, but be counteracted by increasing the representation power the model or generating only short trajectories. To leverage the generative action model, we learn policy To (ut | st) in the latent action space similar to [10 25). But, as both, the latent action prior po(u 8t) and the policy T (Me se) are flexible, it is not ensured that they share the same support. Thus, we formulate policy optimization constrained optimization problem
as regression target for rig (s:)
max Ppo Pe
Σ
(++) (8+)
S.t. Ees-pedde (ūo
sc)] Z
similar to a support constraint 9 We implement the constraint explicitly through parametrization to stay within the support, but do not impose any restrictions inside the supported limits (Figure 3. This is different to using a divergence measure which on one hand does not strictly ensure support limits and on the other hand is more restrictive as it also imposes constraint on the shape of a distribution. Here and in the following ùz stands for a latent action sampled from the policy Ty (wz | st).
The policy is trained to maximize the n-step return V (s7) while staying in support of the latent action prior. Since the latent action prior po po(u, (u, st) is normally distributed as N (H)(s))), oo(s))) we can express the constraint as
po(ae st) 2 == (pu Eog st)
(8)
5

---

with as a parameter setting support as multiples of dg centered around Naa From the properties of a normal distributed variable follows that
μ + Eog 2 úe 2 μ48 EO
We implement the constraint explicitly by parameterizing the policy as linear function g dependent on To (u st) and po(ut 8t):
μe + Eσg 2 (( (up | po(e s:)) 2 μi Edg.
(10)
The support of the policy distribution is chosen to be bounded
üt (us 8t),
ue [-1,1]
(11)
and ( (Me st), (u st)) as linear combination of the latent action predicted by the policy àe and the distribution parameters dg and μs of the latent action prior:
g ūt, po. σa) μe + uq Gg.
(12)
4 Experiments
In the next section, evaluate the effectiveness of our approach. It is divided into three parts: first, we assess the performance using standard benchmarks; then, we study how different design choices affect value overestimation; lastly, we analyze the influence of the support constraint parameter. We additionally provide the final performances in two tables in Appendix
E
We limit our benchmark evaluation to the most relevant state-of-thee tfrhe-rr offline reinforcement learning methods to answer the following questions: 1) How do latent action state-space models compare to state-space models? 2) How comparable are model-free methods focusing on latent action spaces to latent action state models? 3) Does C-LAP suffer from value overestimation? 4) How does the support constraint affect the performance? 5) How does the performance differ between visual observations and observations with low- dimensiooo features? To focus on the latter, we separately evaluate the performance on fow-dimensional feature observations using the D4RL benchmark (26). and on image observations using the V-D4RL benchmark 20
4.1 Benchmark results
D4RL Since most offline model-based reinforcement learning methods are designed for obser- vations with low-dimensional feature observations, there exist many options for comparison. We make selection to include the most relevant methods focusing latent actions and auto-regressive model-based reinforcement learning. Therefore, we include the following methods: PLAS, which model-free method using a latent action space [IO]. MOPO, a probabilistic ensemble-based offline model- based reinforcement learning method using modification to the Bellman update to penalize high variance in next state predictions [66.. And MOBILE, which is similar to MOPO, but penalizes high variance in value estimates instead [18]. We compare the algorithms on three different locomotion environments, namely halfcheetah, walker2d and hopper, with four datasets (medium-replay. medium, medium-expert, expert) each and the antmaze navigation environment with four datasets (umaze, umaze- -ivere,, medium-play medium- diverse). The results, shown in Figure[4] display the mean and standard deviation of normalized returns over four seeds during the phase of policy training, with steps denoting gradient steps. The dashed lines indicate the asymptotic performance for MOPO and MOBILE. A detailed summary of all implementation details is provided in the Appendix D
When comparing C-LAP to PLAS, we find that learning a joint generative model of actions and observations outperforms generative model of only actions when used with actor-critic reinforcement learning. Both methods can use the generative nature of the model to speed up policy learning, which becomes especially clear in the results on all locomotion expert and medium-expert datasets. Compared to MOPO and MOBILE, C-LAP shows superior or comparable performance on all saffcheetah-mediing
datasets except halfchhetall-mediu--rgggeee the antmaze environment, where MOPO and MOBILE fail to solve the task for any of the considered datasets. The asymptotic performance of MOBILE on locomotion environments sometimes exceeds the results of C-LAP, but needs three times many gradient steps.
Especially outperforming -ed-im-rreplla--222222222--
and hopper-medium-v2.

---

beltcheedal-meedimm
Noraalood Retum
75
75
50
50
50
50-
25
25
25
25
0--
walkerdd 0alla22mmdiun-mpplayyy
waller2dmmddmm22
walker2daxpaet-22
125
125
125
125
Normal ardd Return
1.00
100
100
100
75-
75
75
75
50-
hoppat-mdiin--eppappa
hoger--adium-22
hopper-expertv2
Mormalized Retur
1.00
75-
75
75
75
50-
50
50
50
25
25-
25
25
>-ttmaze-umaae-22
antmaze-nedium-poyyo
antmaze-meding d-a.er---e-e222222
125
125
125
125
Normelized Return
1.00
100
100-
100-
0.0
1.0
0.0
1.0
0.0
0.5 Steps
1.0
0.0
0.5 Steps
1.0
Steps
le6
Steps
le6
1e6
1e6
C-LAP
HOTO
MOBILE
PLAS
Figure 4: Evaluation dow-dimensional feature observations using D4RL benchmark datasets. We plot mean and standard deviation of normalized returns over 4 seeds.
Overall the results indicate that latent action state-space models with constrained latent action polices only match the state-of-the-art observations with low-dimensional features as observations, but also jump-stan start policy learning by using the action decoder to sample actions that lead to high rewards already after the first gradient steps: If the dataset is narrow, generated actions when sampling from the latent action prior will fall into the same narrow distribution. For instance, in an expert dataset, sampled actions will also expert-level actions. During policy training, instead of sampling from this prior, we restrict the support of the policy dependent on the latent action prior. Thus, sampled latent actions from the policy will always be decoded fall into the dataset's action distribution. So even randomly initialized policy in the beginning of the training can generate a high reward by using the latent action decoder. This effect is especially prominent narrow datasets such as expert datasets.
V-DARL There are currently few auto-regressive model- reinforcement learning methods that specifically target visual observations, with none emphasizing latent actions. In our evaluation, we include LOMPO [21] and Offline DV2 (20]. Both methods use latent state space model and an uncertainty penalized reward. However the specifics of the penalty calculations are different: while LOMPO uses standard deviation of log probabilities penalty, Offline DV2 uses mean disagreement. Additionaly, LOMPO trains an agent on a mix of real and imagined trajectories with an off-policy actor-critic approach, whereas Offline DV2 exclusively trains imagined trajectories and back- propagates gradients through the dynamics model. Further implementation details are included in Appendix [
C-LAP demonstrates superior performance across all datasets, especially significant on cheetah- run uun-medium, expert, walker-walk malde-uual--
eeppertreeepertttt expert and walker-walk eepper. Datasets with large
7

---

chretaah-runemeeiim. 100
cheetah--n-mmddim
dheetah-ru--medimm expprrt
cheesalllum-eeeett
100
100
100
Nommltee Recurn
00
80
80
80-
60-
60-
00
60-
40-
40-
40-
40-
20-
20
20
20
wwlker-wak-mediimn.eeo 100
replay
walkerwalk-medium
ww..er-all--meiimm expert
walker-wall-exper
100
100
100
Normalzed Return
80-
80-
80-
80-
60-
60
60-
60-
40-
40
40
40-
20
20
20-
0.0
2.5
5.0
0.0
2.5 Stops
5.0
0.0
2.5 Steps
5.0
0.0
2.5
5.0
Steps
led
lel
Steps
1ed
CLAP
Ofline DV2
LOMPO
Figure 5: Evaluation visual observations using V-DARL benchmark datasets. We plot mean and standard deviation of normalized returns over 4 seeds.
diversity of actions, such medium-replay datasets, exhibit weaker inductive bias for generative action model. Hence, they require more additional policy steps, as can be seen for both the D4RL and V-D4RL benchmarks.
4.2 Value overestimation
Limiting value overestimation plays a central role in offline reinforcement learning. To evaluate the effectiveness of C-LAP, we report value estimates alongside normalized returns on all walker2d datasets in Figure [ A similar analysis for all considered baselines is provided in Appendix G To further analyze the influence of different action space design choices, we include the following ablations: variant no constraint, which does not formulate policy optimization objective, but uses a Gaussian policy distribution to potentially cover whole Gaussian latent action space; and a variant no latent action, which does not emphasize latent actions, but uses regular state-space model as in Dreamer [29] Besides that, we added dashed lines to indicate the dataset's s average return and average maximum value estimate. The no latent action variant fails to learn
constrained
malle:22dmadiumm-pppay222
walcer25mmedumm222
waller26-modummoxxpttt
walker2d-oxpert-22
Normsizee Retur
300
100
100
100
75
75-
75
50
50-
50-
50-
25
25-
25
2alko/2d conodiumrrpppay ly
www.mmmaallmmooom walle-2dmmcium.coo 800
medium-2
walker2d
walker2d axperttw2
Vilue Estiaae
800
400
400
400
200
200
200
200
0-
0
1.0
106
Epochs
Spochs
Epochs
1c6
CLAP
no constroint
no latent action
dataset
Figure 6: Ablation study, comparing C-LAP to the following variants: no constraint, C-LAP without enforcing the policy constraint dependent on the action prior; no latent action, C-LAP without a latent action space similar to Dreamer [29]. We plot mean and standard deviation of normalized returns and value estimates over seeds. Moreover we add the dataset's s average return and average maximum value estimate indicated by dashed lines.
8

---

an effective policy: normalized returns are almost zero and the dataset's reference returns remain unattained; value estimates are significantly exceeding the dataset's reference values, indicating value overestimation. The no constraint variant can use the generative action decoder to limit generated actions to the dataset's action distribution, but the Gaussian policy is free to move to regions which are unlikely under the action prior. Thus, nullifying the implicit constraint imposed by the action decoder, resulting in collapsing returns and value overestimation. Only C-LAP achieves a high return and generates value estimates which are close to the dataset's reference. The value estimates on malker2d-medium- -epllaay-- are higher than the dataset's s reference, as the agent's ' performance is also exceeding the reference performance. The results confirm the importance of limiting value overestimation in offline reinforcement learning, and demonstrate that constraining latent action policies can be an effective measure for achieving this.
4.3 Support constraint parameter
To evaluate the influence of the support constraint parameter غ on the performance C-LAP, we perform sensitivity analysis across all walker2d datasets (Figure 7 Except for the more diverse medium-replay-v2 dataset, adjusting é from 0.5 to 3.0 only has a minor impact on the achieved return. However, when choosing an unreasonable large value such constraint altogether (Figure(6), (6. we observe a collapse during training. This highlights key insight: constraining the policy the support of the latent action prior is essential. And in many cases, using a smaller support region closer to the mean (small é) proves sufficient.
10.0 removing the
wolker2d-meeium--g
wolker2d
Norratzed Retumn
200
100
100
100
75-
75-
75
75-
50
50
23
0.0
0.5
1.0
0.0
0.5
1.0
0.0
Epochs
105
lee
Spochs
106
Epochs
106
Figure 7: Sensitivity analysis of the support constraint parameter for the considered D4RL walker2d datasets. We plot mean and standard deviation of normalized returns over 4 seeds.
5 Related Work
Offline reinforcement learning methods fall into two groups: model-free and model-based. Both types aim to tackle problems like distribution shift and value overestimation. This happens because the methods use a fixed dataset rather than learning by interacting with the environment.
Model-free Current model-free methods typically work by limiting the learned policy by regularizing value estimates. TD3+BC [T]] adds a behavior cloning regularization term to the policy update objective enforce actions generated by the policy to be close to the dataset's action distribution. Similarly, SPOT [ includes a regularization term in the policy update, derived from a support constraint perspective. It also uses conditional variational auto-encoder (CVAE) to estimate the behavior distribution. Following a comparable intention, BEAR 8 constraints the policy to the support of the behavior policy via maximum mean discrepancy. BCQ 7 and PLAS (10) use CVAE similarly but don't 't use a regularization term. Instead, they constrain the policy implicitly by making the generative model part of the policy. Beside these methods, many other approaches exist, with CQL [12] and IQL [14] being some of the most well- CQL uses a conservative policy update by setting a lower bound on value estimates to prevent overestimation, while IQL avoids out-of-distribun values by using modified SARSA-lik- objective in combination with expectile regression to only use state-action tuples contained in the dataset.
Model-based Model-based offline reinforcement learning methods learn dynamics model to generate samples for policy training. This basically converts offline learning to an online learning problem. Model-based methods mainly address model errors and value overestimation by using

---

probabilistic ensemble and adding an uncertainty penalty in the Bellman update. MOPO [[6] uses a probabilistic ensemble as in [32] and adds the maximum standard deviation of all ensemble predictions as uncertainty penalty. Similar to that, MORel. [7 adheres to the same methodology, but uses pairwise maximum difference of the ensemble predictions as penalty instead. Analogously, MOBILE [8] estimates the values for all by the ensemble predicted states and uses the standard deviation of value estimates as penalty. Edge of reach [99] comes to the conclusion that value estimation edge of reach states are the overarching issue compared to model errors. In the end, they come up with a comparable solution to MOBILE, but use an ensemble of value networks they
alongside the ensemble of dynamic models. COMBO [28] pursues different approach, integrate the conservatism of CQL into value function updates, removing the need for uncertainty penalties. Besides that, some methods model in observation space, they use latent state-space models to make predictions on latent states. Among these methods is LOMPO (21), which builds up on Dreamer (2]], but integrates an ensemble to predict stochastic states and use the standard deviation of the log probability of the ensemble predictions as uncertainty penalty similar previous methods. The policy is trained on mix of imagined and real world samples, hence they use an off-policy actor-critic style approach for policy learning. Offline DV2 [20] uses a similar model, but is based on a different penalty. Namely, they use the difference between the individual ensemble mean predictions and mean over all ensembles as uncertainty penalty. Furthermore, the policy trained only on imagined trajectories with gradients calculated by propagatng through the dynamics model. Overall, Offline DV2 is the method most comparable our approach, but still different in many ways as we propose latent action state-space model compared to a usual state-space model, and frame policy learning as constrained optimization. So far all discussed models operate in an auto-regressive fashion, but another class of methods exists, which casts offline model-based reinforcement learning as trajectory modeling. Instead of learning a policy, these kind of approaches integrate decision making and modeling of the underlying dynamics into single objective and use the model for planning. Among them are Diffuser (24), which employs guided diffusion for planning; TT 222, which builds on advances in transformers; and TAP [23]. which uses a VQ- VAE with a transformer-baser architecture to create discrete latent action space for planning.
a different class of models: instead of learning predictive
6 Conclusion
We present C-LAP, a model-based offline reinforcement learning method. To tackle the issue of value overestimation, we first propose an auto-regressive latent-action state space model learn generative model of the joint distribution of observations and actions. Second, we propose a method for policy training to stay within the dataset's 's action distribution. We explicitly parameterize the policy depending on the latent action prior and formulate policy learning as constrained objective similar to support constraint. We find that C-LAP significantly up policy learning, is competitive on the D4RL benchmark and especially outperforms on the V-D4RL benchmark, raising the best average score across all dataset's from previously 31.5 to 58.8.
Limitations Depending on the dataset and environment the effectiveness of C-LAP differs: Datasets which only contain random actions are challenging for learning generative action model, thus we do not include them in our evaluation. The effect of jump-starting policy learning with the latent action decoder to already achieve high rewards in the beginning of policy training is prominent in narrow datasets, but less effective for diverse datasets. While training the model of C-LAP does not require additional gradient steps, it still takes more time compared to LOMPO (21) and Offline DV2 [20] as the latent action sate-space model is more complex than a usual latent state-space model.
References
[I] Gabriel Dulac-Aroodd, Daniel J. Mankowitz, and Todd Hester. Challenges of real-world reinforcement learning. CoRR, 0bss/1904.12900 2019.
[2] Jens Kober, J. Andrew Bagnell, and Jan Peters. Reinforcement learning in robotics: A survey. The International Journal of Robotics Research, 32: 1238 1274, 2013.
10

---

[3] Wenshuai Zhao, Jorge Peña Queralta, and Tomi Westerlund. Sim-to-real transfer in deep reinforcement learning for robotics: survey. 2020 IEEE Symposium Series on Computational Intelligence (SSCI), pages 737-744, 2020.
[4] Damien Emsst Pierre Geurts, and Louis Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research, 6:503-556, 556, 04 2005.
[5] Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning, Reinforce- ment Learning: State of the Art, 01 2012.
[6] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives open problems. CoRR, abs/2005.01643, 2020.
[7] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. CoRR, abs/1812.02991 2018.
[8] Aviral Kumar, Justin Fu, George Tucker, and Sergey Levine. Stabilizing off off-policy q-learning via bootstrapping error reduction. CoRR, abss/196.00090 990600999,, 2019.
[9] Jialong Wu, Haixu Wu, Zihan Qiu, Jianmin Wang, and Mingsheng Long. Supported policy optimization for offline reinforcement learning, 2022.
[10] Wenxuan Zhou, Sujay Bajracharya, and David Held. PLAS: latent action space for offline reinforcement learning. CoRR, abs/2011.072111 2020.
[11] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach offline reinforcement learning. CoRR, abs/2106.06860, 2021.
[12] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. CoRR, abs/2006.04779 2020.
[13] Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-based offline reinforcement learning with diversified q-ensemble. CoRR, abs/2110.01548, 2021.
[14] Ily Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. CoRR, 10.06169, 2021.
[15] Cong Lu, Philip J. Ball, Jack Parker- Holder, Michael A. Osboree, and Stephen J. Roberts. Revisiting design choices in model-based offline reinforcement learning. CoRR, 0bb//110.04135, 2021.
[16] Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: model-based offline policy optimization. CoRR, 132399, 2020.
[17] Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. CoRR, abs/2005.05951, 2020.
[18] Yihao Sun, Jiajin Zhang, Chengxing Jia, Hao-Chu Lin, Junyin Ye, and Yangze Yu. Model- bellman inconsistency for model-based offline reinforcement learning. In International Confer- ence on Machine Learming, 2023.
[19] Anya Sims, Cong Lu, and Yee Whye Teh. The edge-of-reach problem in offline model-based reinforcement learning, 2024.
[20] Cong Lu, Philip Ball, Tim G. J. Rudner, Jack Parker-Holder, Michael A. Osborne, and Yee Whye Teh. Challenges and opportunities in offline reinforcement learning from visual observations, 2023.
[21] Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Offline reinforcement learning from images with latent space models. CoRR, abs/2012.11547,, 2020.
[22] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In Advances in Neural Information Processing Systems, 2021.
11

---

[23] Zhengyao Jiang, Tianjun Zhang, Michael Janner, Yueying Li, Tim Rocktäschel, Edward Grefen- stette, and Yuandong Tian. Efficient planning in compact latent action space, 2023.
[24] Michael Janner, Yilun Du, Joshua B. Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis, 2022.
[25] Arthur Allshire, Roberto Martín-Marthn Charles Lin, Shawn Manuel, Silvio Savarese, and Animesh Garg. LASER: learning a latent action space for efficient reinforcement learning.
CoRR, abs/2103.15793, 2021.
[26] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR, abs/2004.07219, 2020.
[27] Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018.
[28] Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. COMBO: conservative offline model-based policy optimization. CoRR, abs/2102.08363, 2021.
[29] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. CoRR, abs/1912.01664 2019.
[30] Philip Ehhck, Maximilian Karl, Jan Peters, and Patrick van der Smagt. Learning to fly via deep model-based reinforcement learning. CoRR, 0bs/2003.08877 2020.
[31] Danijar Hafner, Timothy P. Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. CoRR, abs/1811. 11.0551,, 2018.
[32] Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. CoRR, abs/ abss/805, 1805. 121 14, 2018.
[33] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018.
[34] Patrick Kidger and Cristian Garcia. Equinox: neural networks JAX via callable PyTrees and filtered transformations. Differentiable Programming workshop at Neural Information Processing Systems 2021, 2021.
12

---

A ELBO derivation
We start with Equation (T) jointly model the dataset's observation and action distribution
/ p(orer,arr
p(o.,, (IrT) = =log log
81.7 (irr)(ssrs u.r) ds du,
(13)
and the defintions in Equation 3 namely:
T p BIT, uI.T) = II p(or 8t)p(ar 8t, me),
(14)
((SIT,UITT) = II
p(ue st)p(s; I st-1, ue-1).
(15)
t=1
We introduce the inference model
(St,, Ut St-1, at de, 0.) = ((5 $ (1-, or) se: at)
(16)
and by replacing p(o1:r, 01:T
S1:7, u1:1) and p($1:T, u1:r) in Equation 13) we get:
logp(or, a1r) = log / II
T
p(o,
st)p(a; | st, up)p st)p(s; 2.-1,4-11 ds du.
(17)
t=1
We include the inference model and resolve using Jensen's Inequality:
f II p(or
T
((.r, Ue
St-1. 01-1,0-11,00 a2 0z)
log (1.)) = log
se)p(a
st, ue)p(ue st)p(s; sc-1, u1-1)
St-1, a2 oz) du
( Ue
t=1
(18)
plor st)p(ar I 8t, wr)p(us se)p(se 88- 1,-11
= log E.
(19)
((s,, us 88-1, 01-1, az 0z)
p(s: 8t-1, uz-1)
p(ue 8t)
log E.
p(oz st)p(as 8c, us) ((st 8t-1, 0,-1, 0z) q(1 81, at)
(20)
t=1 l
log plst ( + log
p(u, sc)
logplo; s;)p(a, I st, uz) (.)4e)
q(s s1-1,
( u
8t, a.)
(21)
E, [logp(o st)p(a: sc, uz]
(22)
(L((( 18-1, || p(se St-1,4411 KL(g(u st,aa) 11 po(u st]]]
(23)
(24)
-LLELBO(01TTA117
(25)
13

---

B Algorithm
We provide the general algorithm of C-LAP.
Algorithm 1: C-LAP
Given: Dataset D, learning rates Og, Otds Oigs Org, sequence length K, dream rollout length H
// Model Training
Initialize model parameters 0 and & for epoch in model training epochs do for update step r 1..T do
Sample batch of trajectories {(o), 0t, (x,r)t)))) ltkk ~ Compute model training objective L(01:r, 01:T) via Equation 5
Update model parameters 9 + 0 + as VoL(on. air), e t o VoLLonn
end for end for
Agent Training
Initialize policy and critic parameters & and for epoch in agent training epochs do for update step == 1.TT do
Sample batch of trajectories {(or, (t rts (1) k+)
Compute latent states 8, 20(8t 8t-1, (t-1,
Sample random starting state Sinit from each trajectory
Create dream rollouts {(T,, ur: ar) 1 from Sinit using po(st st-1, u1-1) and
~ p(a, st, (TT (Mt | (2)), Pu ss)
Predict rewards r ~ pa(r= 87), termination 1, pll 87) and values
'gg ~ (u- 87)
Compute value estimates V (s-) via Equation (
Update policy parameters ti + + a V D (s+)
Update critic parameters + E + ag V rut (s,) V (s+)1
end for end for
C Computational Resources
We different hardware setup for experiments with visual observations and experiments with low-dimensional features as observations. We run all experiments on a shared local cluster. C-LAP experiments with visual observations take around 10 hours RTX8000 GPU and experiments with low-dimension feature observations around 11 hours on A 1100 GPU. We aim to execute most of our code GPU and parallelize our implementation. Environment evaluations represent bottleneck as they require execution on CPU, Overall, it takes around 70 combined GPU days to reproduce the benchmark results of all methods. This does include the compute required for the evaluation of preliminary implementations or hyper-paramene settings.
D Implementation Details
We implement all methods in JAX 33 using Equinox 34 We provide the hyper-parameters of C- LAP in Table and the constraint values used for the D4RL benchmark Table for the V-D4RL benchmark in Table 4 In general consider constraint values in the range E (0.5, 1.0, 2.0, 3.0).
We implement MOPO and MOBILE following [18 and use respective hyper-parameters provided in the publication. For MOPO, we select the max-aleatoric version. As no hyper-parameters are provided for the expert datasets of the D4RL. benchmark we sweep through the range specified in Table and use the one selected in the table. For the antmaze environment we additionally evaluate
14

---

the discount factors of {0.99, 0.995, 0.997} to compensate for the sparse reward, but do not find hyper-parameter to solve the task.
We implement PLAS following 10 and use the default hyper-parameters for the expert datasets.
We implement LOMPO and Offline DV2 using the continuous state implementation from Dreamer [2)]. which is different from the original implementation using discrete states. This change keeps the model closer to C-LAP, which is also using continuous states. For LOMPO and Offline DV2, we take hyper-parameters from [20].
As all our implementations might differ from the original one, we include the original scores from the paper alongside our results Table and Table 6
MОPO
MOBILE
penalty coefficient é (0.5,2.6,5.011000
10.,)
penalty coefficient 2.55 rollout sleps: (1,5)
Hyper-parametd range
rollout steps: è [1.5}
datasct ratio: {0.05, 0.5, 0.8)
datasct ratio: [0.05, 0.5, 0.8)
penalty coellicient: 1.0
penalty enefficient: 2.5
setrhectah-expert.oe
millaut steps: 5
rillout steps:
dataset ratio: 0.5
dataset ratio: 0.5
penalty coefficient: 10.0
penalty coefficient: 2.5
walker2d expert-v2 .22
rollout steps: 1
rollout steps:
dataset ratio: 0.8
dataset ratio: 0.5
penalty coefficient: 5.0
penalty coefficient: 3.5
hopper-expert-v2
rollout steps:
rollout sleps;
dataset ratio: 0.05
dataset ratio: 0.5
Table 1: MOPO and MOBILE haper-parameters for the expert datasets in the D4RL benchmark.
Environment
Constraint
medium-replay
3.0
Halfcheetah
medium
3.0
medium-expen
2.0
2.0
Environment
Dataset
Constraint
medium-replay
3.0
medium-replay
3.0
Walker2d
medium
1.0
Cheetab-run
medium
3.0
mecpium-expert
3.0
redium--xpeet
3.0
expert
2.0
expert
3.0
medium-replay
2.0
medium-replayy
3.0
Hopper
medium
3.0
Walker-walk
medium
3.0
medium-expert
medium-experi
expert
0.5
expert
0.5
umaze
1.0
umaze-diverse
1.0
Table 4: C-LAP constraint values for the V-D4RL
Antmaze
medium play
benchmark
mediun-diverse
1.0
Table 3: C-LAP constraint values for the D4RL benchmark
15

---

Model
Stochastic
64 (antmace)
30
fnn-dimessional features)
Deterministic
antmaze)
30 Jimenssonal (euluures)
Observation embralding size
64 (animazz)
1024 (vasual)
stent action size
12
activation
MLP
unirs: 128
layers:
State encoder
CNN
channels: 256]
Mride: 14, 4,4.44
MLP
units: 128
layers: distibution: Gaussian
channels: [126, 64,32,3
kemols: 6,
strde: 2
distrbution: Goussian
layers: distribution: Coussian
units: 512
Lasent action decoder
layers: distributin:: Rota
prior
layers: distributio: Gaussian
Learning rate
3e 4
Batch size
64
Window longth
50
Agent
Hidbee activation
selu
units: 266
Policy
layers:
distribution: TanhCGnussian
Value
layers:
number netwarks:
Leurning rale
Se-5
Legrobbentreyy regulariser scaling
0.01
rollout
Discount factor
0.907
(antmaze umaze, intmaze medium-diveree)
0.906
(intmaap-unaee antisne-metiom--lee
Table 2: C-LAP syper-parameters
16

---

E D4RL and V-D4RL benchmark results
We provide the benchmark results in the usual offline reinforcement learning table format. As our implementations might differ from the original one, we include the original scores from the paper alongside our results in tables. We also include the average over all datasets except the expert datasets to make comparable to the average scores provided in the respective publication.
Environment
Dataset
PLAAS
(paper)
10
PLAS
MOPO
(paper)
18
MOPO
MOBILE
(paper) [88
MOBILE
C-LAP
medium-replay
41.8
72.1
06.3++5.3 5.3
71.7 71.7+122
63.7 +
56.5 +
Halfchectah
medium
39.4
44.8 0.5
72.4
57.5 30.4
74.8
1.2
77.3 2.0
56.2
mediam-cxpert
96.6
82. 0.5
83.6
104.4+227
108.2 2.5
103.2
96 0.3
expest
25.7
median-replay
30.2
67.7 11.9
85.2
51.61:20.5
89.9 1.5
85.1 10.6
86.0 20.1
Walker22d
medium
44.6
79.4 1.8
84.1
78.0 22.5
57.7 1.1
81.0 0.5
82.5
medium--cxpert
$9.6
109.1
105.3
105.6 = 7.1
113.2 0.7
113.5
111.8
expert
109.1 0.4
113.4 1: 0.5
15.0 :: 17.3
111.7 0.3
replay
Hopper
medium
32.9
57.0 4.3
62.8
76.8 41.3
106.6 0.6
106.2 + 0.2
80.3 18.6
mediant-cxpert
111.0
56.7 = 6.4
74.9
43.5 28.2
112.6 0.2
111.8 + 1.8
106.0 = 7.7
expert
07.1
26.2 5.1
78.6 - 38.4
Locemotion average without expert
37.3
66-3
81.5
74.9
96.7
93.1
83.6
Locomotion average
74.0
92.7
89.3
umaze
0.7
46.7 =: 7.6
0.0
0.0
81.3 1. 4,8
umaze-diverse
0.5
33.3 14.4
0.0
76.0 20.4
mediam-play
0.2
0.0
0.0
77.5 9.6
mediam--iverse
0.0
0.0
0.0
0.0
46.0 1:0 37.0
Navigation average
0.35
20.0
0.0
69.7
Table 5: Results on the D4RL benchmark. Showing normalized returns and standard deviations at the end of policy training.
Offline DV2
LOMPO
Environment Dataset
Offline DV2
LOMPO
C-LAP
(paper) (20
(paper) 220
medium-replay
61.6 à: 1.0
54.6 à: 6.5
36.3 à: 13.6
42.2 3.5
52.5 = 1.3
medium
17.2 1: 3.5
16.8 8.7
1.3
52.3=3.8
57.4 1 0.9
Cheetah-run
medium-expert
10.4
9.4 6.3
11.9199
5.9
73.2 = 8.5
expert
10.9 3.2
5.
1.8
8.4
36.6
11.3
medium-replay
56.6
42.5 14.2
34.7
11.4
5.8
34.5 5.5
medium
34.1
19.7
40.6 10.8
43.4
11.1
39.0
5.7
54.8
1.2
Walker-walk
medium-expert
43.9 1. 34.4
41.0 1 20.0
39.2 1 19.5
58.0
5.3
93.2
2.2
expert
4.8 0.6
6.1 5.6
5.3 1 7.7
25.6
17.6
68.1
14.0
Average
29.9
27.0
25.2
31.5
58.8
Table Results on the V-D4RL benchmark. Showing normalized returns and standard deviations the end of policy training.
Action distribution
To evaluate the claim that latent actions generated by the latent action prior are close to the dataset's action distribution we use the following approach: We randomly select one trajectory from the hopper-expert-v2 dataset and employ k-nearest neighbors to identify the 20 nearest observations and their corresponding actions within the whole dataset for each step. We then sample from the action prior and decoder to generate actions based on the nearest observations at each step. Thereafter, we split the selected trajectory into sections from leaving the ground to touching the ground and fit normal distributions to the aggregated dataset actions and reconstructed prior actions. Thus, we end up with an approximation of the dataset's 's action distribution and an approximation of the action distribution generated by the prior for each step in the trajectory. Figure & shows the aggregated distributions (from leaving the ground touching the ground) for one exemplary trajectory of the hopper-expert-vV v2 dataset.
17

---

Action distribution hopper-expertv22
Action
latent action prior
Action
Action 2
Stop
Figure 8: Comparison of the dataset's action distribution to the distribution of actions sampled from the latent action prior and latent action decoder. The figure shows the aggregated distributions (from leaving the ground to touching the ground) for one exemplary trajectory the hopper-expert-v2 dataset
wolkerZd-mediam--plaauu
waller22dmediumm-2
wwtler2--medummmppprr-
wollerr2dexxper--2
Nomalted Return
100
100
100
100
50
50-
23-
25-
25
25
Walker2d 72d-mmedium--ppayyy.
waller24-mediumm-2
walker23-axpertx22
800
500-
800
800
600
800
0.0
0.5
1.0
0.0
0.5 Ecochs
1.0
0.0
0.5 Epochs
0.0
1.0
Epoche
Epochs
C-APP
FLAS
MOPO
MOBILE
dataset
Figure 9: Evaluation of value overestimation for all baselines. We plot mean and standard deviation of normalized returns and value estimates over seeds. Moreover, we add the dataset's 's average return and average maximum value estimate indicated by dashed lines.
G Value overestimation for the considered baselines
We further report value estimates alongside normalized returns all walker2d datasets for the considered baselines (Figure(9). As they estimate Q-values, we calculate the corresponding value estimates by averaging over actions sampled from their respective policies. MOPO and MOBILE have a low value estimate, which can be attributed to the incorporated uncertainty penalties. PLAS's value estimates are only stable for walker2d walker2d-expers -xpprrtt-22 datasets, but collapse for the other considered datasets. Overall, it seems that value overestimation is not the cause for degrading performance for these methods.
18

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We claim that C-LAP competitive to state-of-the-arl methods and especially outperforming datasets with visual observations. These claims are based evaluations on the D4RL and V-D4RL benchmark and the comparison to other relevant methods. We claim that aim solve value overestimation with our approach, which we empirically show on the walker2d datasets.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We include the limitations in a separate paragraph in the conclusion.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
19

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We outline the derivation of the ELBO in the main section (Sec. 3) and provide a detailed derivation in the appendix A. We explain the derivation of the constrained latent action policy step-by-step in the main section (Sec. 3).
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We include the explanations of your approach in the main section (Sec. 3). We report the used hyper-paramene parameters and an outline of the algorithm in the appendix (Appendix B. D). We provide the code. We use publicly available benchmarks.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
20

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We provide the code and include configuration files to reproduce results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ //nips httpps..coppbllc/guu.eeee
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permirted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: Yes, shortly explain the experimental setting in the experiments section (Sec. 4). We provide a general algorithm, a section on implementation details and the syper-parameten the appendix (Appendix A, B, D). Furthermore, we share the code.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We provide the number used seeds and report mean and standard deviation in our evaluations in the experiments section (Sec. 4).
Guidelines:
21

---

The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi-
dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative етror rates).
If error bars are reported in tables or plots. The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We explain the required computational resources appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http::/s/////ers.....ooo http://wwwwp..co//ll/l/ccccheeea
Answer: [Yes]
Justification: We give access to the complete method details and thoroughly share details of the experiments.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
22

---

Justification: We proposed method falls in the category of foundational research in the field offline reinforcement learning and the method is not tied an application. Of course, applications exist where reinforcement learning might have negative societal impact, but our method only has a limited influence on that.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative socictal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeett of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: There is no high risk of misuse of our method.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite the original papers which published the benchmarks which we use. We cite the authors of JAX and Equinox, as our implementation is based on their framework. We give credit Twemoji in figure 1.
Guidelines:
23

---

The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We provide code and include configuration files to reproduce experiments.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/mooel as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No crowdsourcing research with human subjects.

---

Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
25


Reviews:
[{"review_id": "4tdJNUjkR1", "reviewer": "Reviewer_UrFS", "summary": "The work introduce C-LAP (Constrained Latent Action Policies), a novel approach to model-based offline reinforcement learning in POMDPs. Notably, C-LAP does not employ explicit reward penalty terms regarding the action space.\n\nThe paper proposes a methodology for learning latent variables for both states and actions based on the state-space models (SSMs). It then frames policy optimization as a constrained optimization problem, aiming to maximize N-step returns while ensuring that latent actions remain within the support of the latent action prior.\n\nThis approach allows the policy to maximize returns predicted by the model without incorporating reward penalty terms, such as uncertainty penalization. Empirically, C-LAP demonstrates superior performance compared to baseline methods that incorporate uncertainty penalties in their reward structure.", "questions": "Q1. The manuscript suggests potential applicability of the C-LAP framework beyond the traditional models explored. Have the authors considered leveraging foundation models, such as transformers or diffusion models, within the latent space formulation of C-LAP? If so, could you elaborate on how C-LAP might integrate with these advanced model architectures?\n\nQ2. The proposed C-LAP method appears to be promising for environments characterized by high-dimensional action spaces, including those analogous to language action spaces. Could the authors provide insights into the adaptability of C-LAP to such high-dimensional scenarios? What modifications, if any, would be necessary to optimize its performance in these contexts?", "limitations": "The authors included their limitation in Section 4.2. They address societal impacts of their work in the checklist #10. Broader Impacts.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The manuscript presents a novel and effective model-based method for offline reinforcement learning in POMDPs. Prior approaches in the domain of model-based offline RL often incorporate penalty terms to mitigate OOD actions and prevent overestimation of the value function. This penalization, however, can restrict the maximization of value functions related to possibly promising actions.\n\nIn this work, the authors extend a standard state-space model (SSM), exemplified by model-based approaches such as Dreamer, to include latent action variables. By constraining the latent action space to adhere to an action prior, the generation of OOD actions is effectively precluded. \nThis modification is both simple and effective, addressing a critical challenge in the field.\n\nThe effectiveness of the proposed method is validated across both MDP and POMDP settings using standard benchmarks such as D4RL and V-D4RL. Moreover, the authors provide a thorough ablation study to identify the critical components contributing to performance improvements. As described in Figure 6, omitting the constraint on the latent action prior is the most crucial to performance, underscoring its significance as a key contribution of this research.", "weaknesses": "Despite the novelty of this work, the motivation for adopting latent action spaces is insufficiently articulated. \nThe paper does not clearly address why constraining with respect to the latent action space is advantageous over the actual action space. \nSpecifically, it remains unclear whether C-LAP offers any benefits compared to SSMs that include a support constraint term directly on the actual action space, as expressed by the constraint  $\\mathbb{E}_{s \\sim p, a\\sim \\mu(s)}[\\pi (a|s) ] \\ge \\epsilon$\n(where $\\mu$ denotes the behavior policy) for an instance.\nI believe that a more thorough explanation for the motivation could significantly strengthen the contributions and arguments of this work.\n\n(Minor comment) In Eq. (7), a constraint term would be a typo. It would be more accurate to revise this as $\\mathbb{E}_{p,\\pi}[p(h_t|s_t)]\\ge \\epsilon$, given that $h_t$ is sampled from $\\pi(s_t)$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**The paper does not clearly address why constraining with respect to the latent action space is advantageous over the actual action space. Specifically, it remains unclear whether C-LAP offers any benefits compared to SSMs that include a support constraint term directly on the actual action space, as expressed by the constraint  (where  denotes the behavior policy) for an instance. I believe that a more thorough explanation for the motivation could significantly strengthen the contributions and arguments of this work.**\n\nTo implement a support constraint alongside a state-space model, it is necessary to estimate the behavior density. This can be achieved by training an additional VAE, as demonstrated in SPOT [1]. Conversely, C-LAP integrates the objectives of state-space modeling and behavior density estimation into a single generative model of observations and actions.\nUsing a regular state-space model combined with a behavior density estimator, one could implement a support constraint directly on the actual action space. However, this approach presents the following challenges:\n- If the policy generates actions directly in the environment’s action space, the generative action decoder cannot be used to jump-start policy learning. This drawback becomes especially clear in the \"expert\" datasets shown in Figures 4 and 5, as discussed further in the general rebuttal section. Another approach is to learn a latent action policy and use the generative action decoder, similar to our C-LAP and PLAS [2], while implementing the constraint on the actual action space.   \n- However, both of these variants require the introduction of a Lagrange multiplier in the policy objective. A Lagrange multiplier does not guarantee compliance with the constraint and is difficult to tune and analyze. C-LAP avoids altering the policy objective and uses an explicit parameterization of the latent action space through the latent action prior. The parametrization has a clear interpretation (multiples of $\\sigma_{\\theta}$ centered around $\\mu_{\\theta}$) and is not very sensitive to the choice of $\\epsilon$ (Rebuttal Figure 4b). \n- Implementing a support constraint via a Lagrange multiplier $-\\lambda \\log \\pi_{\\beta}$ not only restricts the support, but also influences the whole shape of the distribution. Explicit parameterization on the other hand only scales the distribution to ensure the boundaries.\n- The behavior density on the actual action space needs to be approximated with multiple samples from the latent action prior and subsequent decoding, while a support constraint on the latent action space does not require sampling from the latent action prior at all.\n\n**Q1. The manuscript suggests potential applicability of the C-LAP framework beyond the traditional models explored. Have the authors considered leveraging foundation models, such as transformers or diffusion models, within the latent space formulation of C-LAP? If so, could you elaborate on how C-LAP might integrate with these advanced model architectures?**\n\nYes, as some foundation models are optimizing a similar objective as $p(o_{1:T}, a_{1:T})$, leveraging these models is an interesting direction for future research. C-LAP explicitly defines both a latent state space and a latent action space through the latent action state-space model formulation. This formulation is applicable to other model architectures. Therefore, the encoder-prior-decoder structure of states and actions in C-LAP can be replaced with a transformer or a diffusion model to learn a foundation model with an explicit internal structure. Moreover, given an already pretrained foundation model, a couple of interesting research directions come to our mind: \n- C-LAP can be used to adapt the foundation model’s action space. For example, given a generalist robot policy such as Octo [3], it is feasible to use C-LAP to learn a new action head while ensuring state space consistency. \n- C-LAP can be used to distill knowledge from a foundation model into a smaller model by using the foundation model as a differentiable environment. \n- Foundation models optimizing $p(o_{1:T} \\mid a_{1:T})$ can replace the latent state prior, latent state posterior and observation decoder. By fine-tuning with C-LAP on a small dataset, we can use the predictive power of the foundation model and only learn the additional latent action representation. \n\n\n**Q2. The proposed C-LAP method appears to be promising for environments characterized by high-dimensional action spaces, including those analogous to language action spaces. Could the authors provide insights into the adaptability of C-LAP to such high-dimensional scenarios? What modifications, if any, would be necessary to optimize its performance in these contexts?**\n\nWhen adapting C-LAP to high-dimensional action spaces, two aspects should be considered: First, the size of the latent action space must be adjusted to suit the environment. Second, it can be beneficial to condition the deterministic transition $f(h_{t-1}, s_{t-1}, a_{t-1})$ directly on latent actions $u_t$​ instead of decoded actions $a_t$​. This approach simplifies the learned latent state prior $p_{\\theta}(s_{t} \\mid s_{t-1}, u_{t-1})$ by reducing the dimensionality of the considered actions. \nTo adjust C-LAP for action spaces informed by language, we propose to extend the generative model to $p(o_{1:T}, a_{1:T} | L_{1:T})$ with $L_{1:T}$ as a language condition. Deriving a state-space formulation, the latent action posterior, latent action prior and action decoder could be extended with an additional language condition similar to “LILA: Language-Informed Latent Actions” [4]\n\n---\n\n[1] Wu et al., Support Policy Optimization for Offline Reinforcement Learning, NeurIPS 2022\n\n[2] Zhou et al., PLAS: Latent Action Space for Offline Reinforcement Learning, CoRL 2020\n\n[3] Ghosh et al., Octo: An Open-Source Generalist Robot Policy, 2024\n\n[4] Karamcheti et al., LILA: Language-Informed Latent Actions, CoRL 2021", "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": "KEX2SeidUy", "reviewer": "Reviewer_932j", "summary": "The paper approaches the problem of model-based policy learning from static datasets. It firstly identifies that this offline MBRL setup inherits two major concerns from its components: the problem of value overestimation from out-of-distribution actions (common in offline RL) and the bias originated by learning policies from model predictions. In order to address these issues, it proposes C-LAP, a MBRL method that besides learning a latent state space model - as common in the literature - also learns a latent action space from the static dataset. This generative model of states AND actions is further leveraged by the policy optimization objective to constrain the policy actions to share support with the latent action prior, in order to prevent out-of-distribution samples. The paper provides experiments in two offline RL benchmarks, with feature and pixel-based observations, and showing good improvements in the pixel-based observation setting.", "questions": "Please refer to the weaknesses section. Each Major Concern contains questions to be addressed during the rebuttal.", "limitations": "- The main limitation in my perspective is that it is not clear why the method presents much stronger results in the visual benchmark (which is supposed to be harder) than in the feature-based benchmark. This opens questions about whether the introduced latent modeling is indeed doing what is hypothesized in the paper.\n\n- There are some limitations in how the results are presented and discussed in the work (please refer to first two bullet points in the Weaknesses section).\n\n\n**Summary of the Review**:\n\nOverall, the paper is well-written, clearly describes the proposed methodology and offers strong results in the visual benchmarks. The work is based upon a very interesting insight on what latent variables could potentially offer to restrict out-of-distribution actions in offline RL settings. The paper does present some concerns, most related to better explaining the presented results and adjusting claims accordingly. There are also other minor concerns/further suggestions that do not prevent acceptance but would also enrich the paper (which I would appreciate with a higher score). That being said, I believe the paper is already above the acceptance threshold and I am willing to increase my score if my concerns (mostly questions) are properly discussed in the rebuttal phase.\n\n\n\n\n\n\n\n# Post-Rebuttal\n\nThank you, authors, for properly addressing my questions and concerns during the rebuttal. After rebuttal, I am increasing my score.", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "- The paper addresses the open problem of offline model-based reinforcement learning, which is very relevant to the RL community nowadays. The paper is well-written, very clear and organized: the methodology is carefully and formally described; the baselines are clearly and concisely described (and justified).\n\n- Despite the incremental nature of the paper (while mixing up different building blocks from previous literature), the work is really clever on identifying that a latent modeling of actions would implicitly constrain the action space to the data distribution, which would naturally provide a regularization effect on learning policies with static datasets. Therefore, the general methodology is very well motivated and applied to the problem.\n\n- The asymptotic results in the V-D4RL are really strong and the ablations in Section 4.2 presents very clearly the effect of the proposed method in the problem of value overestimation. These results give solid evidence that the method is indeed addressing the previously raised problems.", "weaknesses": "- The claim in L200 regarding a “significant speed up in policy learning” is questionable, since comparing the number of gradient steps to achieve a certain performance is not fair. Some methods can leverage less computation in the updates than others. Also, one can carefully tune the hyperparameters (such as the learning rate, batch size) to minimize the number of gradient steps required, at the cost of more computation. Therefore, showing the gradient steps in the x-axis is not ideal and does not seem to provide useful information. It would be interesting to bring the computational cost for claiming speed up. Optionally, the work could remove such a claim.\n\n- There is an unclear trend in the results of the proposed method: in many cases the starting policy already performs close to the best, while other methods start from performance 0. For instance, refer to the expert cases in D4RL. Why does this happen? \n    - The most interesting case is the walker 2d-expert-v2 in Figure 6, where one ablation starts optimal and then goes down. It would be interesting to justify or at least hypothesize what is happening.\n\n- The bound of Eq. 11 is not justified in the paper and looks arbitrary. From my understanding no prior is placed in the policy so this looks like an assumption. It would be interesting to elaborate better about this.\n\n- The last concern is perhaps the most general and crucial to be cleared out during rebuttal. It is clear that the method worked much better for visual domains than feature domains (the results in the D4RL dataset do not show any asymptotic improvement – it is actually worse than MOBILE as per Table 5). In both cases, the action space is the same. It is also low-dimensional. This raises the question: is the model really implementing latent variables for actions or is it actually implementing a hierarchical latent model for the state space? The work explicitly defines “u_t” as latent actions, but it could also be interpreted as a higher-level latent state variable. Given that the gradients from the action decoder backpropagates through both s_t and u_t, I couldn’t find an explanation on why this is not a valid explanation of the method, and this would better justify the discrepancy in the results.\n\n**Minor Concerns/Further Suggestions**\n\n- It would be great to show the value overestimation plots (Figure 6) for the considered baselines, in order to compare the difference among them. It is unclear if they perform worse than the proposed method because of value overestimation or other reason.\n\n- It would also be nice to show a sensitivity analysis of the policy constraint hyperparameter epsilon across a few environments, to help understanding how the proposed method behaves and how much it is environment specific.\n\n- In L105, the paper contrasts the proposed C-LAP to previous works stating that they “rely on uncertainty estimates to generate trajectories within the data distribution”, but I would consider rephrasing that since C-LAP can also be understood as leveraging the uncertainty of the latent action prior to generate trajectories. Perhaps the distinction is that previous methods explicitly use uncertainty estimates, while C-LAP relies on a more implicit constraint.\n\n**Typo:**\n\n- A missing “)” in Eq. 12.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**The claim in L200 regarding a “significant speed up in policy learning” is questionable, ...**\n\nAfter discussing pros and cons, we revise the claim and keep gradient steps as we would need to rerun all experiments again to compare computation time (not comparable because of variations in cluster utilizations and nodes) or computational costs in FLOPS (not logged).\n\nTherefore, we revise our claim as follows:\n“...jump-start policy learning by using the action decoder to sample actions that lead to high rewards already after the first gradient steps. This effect is especially prominent in narrow datasets such as expert datasets.”\n\n**There is an unclear trend in the results of the proposed method: in many cases the starting policy already performs close to the best, while other methods start from performance 0. For instance, refer to the expert cases in D4RL. Why does this happen?**\n\nPlease take a look at our general reply.\n\n**The most interesting case is the walker 2d-expert-v2 in Figure 6, where one ablation starts optimal and then goes down. It would be interesting to justify or at least hypothesize what is happening.**\n\nIn Figure 6, the initial close to best performance is achieved for the same reason as pointed out in the general reply. The ablation “no constraint” degrades in performance thereafter, as the policy is free to generate actions which are not contained in the dataset’s action distribution. Opposing to that, C-LAP restricts the support to the latent action prior with multiples of $\\sigma_{\\theta}$ centered around $\\mu_{\\theta}$ (Equation 9), limiting the possibility to sample out-of-distribution actions and thereby restricting value overestimation.\n\n**The bound of Eq. 11 is not justified in the paper and looks arbitrary. From my understanding no prior is placed in the policy so this looks like an assumption. It would be interesting to elaborate better about this.**\n\nWe noticed a mistake in text and notation that might have caused confusion. The support of the distribution is chosen to be bounded, not the probability. Thus, we change line 163 to “The support of the policy distribution…” and Equation 11 to:\n$\\hat \\pi_{\\phi}(u_t \\mid s_t)$ with  $u_t \\in [-1, 1] $\nTo explicitly implement the proposed constraint, the support of the policy needs to be bounded between [-1, 1]. Without this bound, using the linear combination in Equation 12 and setting the support with $\\hat{\\epsilon}$ as multiples of $\\sigma_{\\theta}$ centered around $\\mu_{\\theta}$ would not be feasible. Thus, we implement the policy as a TanhNormal distribution.\n\n**Concern: Is the model really implementing latent variables for actions or is it actually implementing a hierarchical latent model for the state space? Why are C-LAP’s results more impressive on images?**\n\nThe latent action posterior, prior and decoder can shape the latent state by back-propagating through $z_t$. However, $u_t$ has no effect on state predictions as the latent state prior and latent state posterior are based on the deterministic transition function $f(h_{t-1}, s_{t-1}, a_{t-1})$, which is conditioned on actions $a_{t-1}$. Hence, the only objectives acting on $u_t$ are the action consistency (KL) and action reconstruction loss. Still, some information of $z_t$ could be in $u_t$ as latent action posterior, prior and decoder are conditioned on $z_t$. But to incentivize some kind of hierarchical state in $u_t$, the observation consistency and observation reconstruction loss would need to act on $u_t$, which is not the case.\n\nComing to the initial concern, there are several possible explanations for why C-LAP’s results are more impressive on images (V-D4RL) compared to full state information (D4RL):\n- The datasets, collected using different policies, and the environments, namely dm-control and gym, differ significantly. For example, the results on halfcheetah-expert-v2 (D4RL) appear less impressive, while the results on cheetah-run-expert (V-D4RL) look great in comparison with the baselines, even though C-LAP achieves 97.1 on halfcheetah-expert-v2 and 36.6 cheetah-run-expert. This makes direct comparisons between D4RL and V-D4RL challenging.\n- The Dreamer architecture commonly performs better on image observations than full state information.\n- The baselines on V-D4RL and D4RL are not the same. With D4RL being the benchmark which attracted more research in the past.\n\n**Value overestimation plots (Figure 6) for the considered baselines**\n\nWe created a separate plot to analyze value overestimation for the considered baselines (Rebuttal Figure 3). For PLAS, MOPO, and MOBILE, which estimate Q-values, we calculate the corresponding value estimates by averaging over 10 actions sampled from their respective policies. MOPO and MOBILE have a low value estimate, which can be attributed to the incorporated uncertainty penalties. PLAS’s value estimates are only stable for walker2d-expert-v2 datasets, but collapse for the other considered datasets. Overall, it seems that value overestimation is not the cause for degrading performance for these methods.\n\n**Sensitivity Analysis of the policy constraint hyperparameter epsilon**\n\nWe performed a sensitivity analysis (Rebuttal figure 4b) of the policy constraint across all walker2d datasets. Except for the more diverse medium-replay-v2 dataset, adjusting the constraint within the specified range had only a minor impact on the achieved return. Nevertheless, the absence of a constraint leads to a collapse during training (Figure 6). \n\n**Rephrasing L105 “rely on uncertainty estimates to generate trajectories within the data distribution”**\n\nWe will rephrase it as follows to highlight the distinction:\n“Unlike other model-based offline reinforcement learning methods that learn a conditional model $p(o_{1:T} |a_{1:T})$ and rely on ensemble based uncertainty penalties on the Bellman update to generate trajectories within the data distribution, …”", "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": "kqIM9iPYnV", "reviewer": "Reviewer_6cNB", "summary": "This paper proposes C-LAP, a model-based offline RL method that mitigates the distributional shift problem without using uncertainty penalties or modifying the Bellman update. It learns a joint distribution of latent states and latent actions and constrains the latent action policy to the dataset distribution by using a linear transformation. Experiments on D4RL datasets and V-D4RL datasets report performance comparable to other strong offline RL baselines.", "questions": "Please see the weaknesses section.", "limitations": "This paper discusses the limitations, but that is not enough.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. This paper addresses the issue of distributional shift in offline RL by incorporating latent actions to enhance policy learning. The empirical results demonstrate that the proposed method outperforms other offline RL baselines by requiring fewer gradient steps to learn a policy.\n2. The idea of combining a latent action prior and a bounded policy to obtain the constrained policy is novel and intuitive.", "weaknesses": "1. Figure 1 requires a clearer presentation of the model diagram. It would be better to simultaneously show the prior and posterior in Figure 1(a), and the latent action prior $p_\\theta(u_t|s_t)$ and the policy $\\pi_\\psi(u_t|s_t)$ in Figure 1(b) to improve understanding. During the policy training phase, are the imagined trajectories rolled out from the latent state prior $p_\\theta(s_t|s_{t−1}, u_{t−1})$ or the latent state posterior $q_\\phi(s_t|s_{t−1}, a_{t−1}, o_t)$? If it is the latent state prior, why is the output of the action decoder used? If it is the latent state posterior, $o_t$ is unavailable.\n2. Could the authors provide some distribution results to support the claim that the latent actions generated by the action prior are close to the dataset's action distribution? Does the predictive power of using latent actions degrade in performance compared to models using real actions? Some quantitative or qualitative results can be provided.\n3. It would be better to highlight the differences between the proposed method and Dreamer in terms of model designs and algorithms.\n4. The experiments mainly focus on Mujoco tasks. How does the proposed method perform on navigation tasks such as AntMaze?\n5. Why do some experiments in Figure 4, 5 show high returns at the beginning of policy learning (zero gradient step)?\n6. In table 5 and table 6, the performance of C-LAP is inferior to some baseline models in medium or medium-relay datasets. This results seem to indicate that the effectiveness of the proposed method depends on the quality of the offline dataset.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**1a) Improvements regarding Figure 1**\n\nWe updated Figure 1 (see Rebuttal Figure 1) to include the latent action prior and highlight the support constraint. \n\n**1b) During the policy training phase, are the imagined trajectories rolled out from the latent state prior  or the latent state posterior ?**\n\nThe policy is trained on imagined trajectories generated by the latent state prior. The latent state prior is based on the deterministic transition $f(h_{t-1}, s_{t-1}, a_{t-1})$, which requires actions $a_{t-1}$ in the environment’s action space. Thus, we use the latent action decoder to generate actions from latent actions sampled from the policy.\n\n**2a) Claim that the latent actions generated by the action prior are close to the dataset's action distribution**\nTo evaluate the claim we use the following approach: \nFor each step in every trajectory within the dataset, we employ k-nearest neighbors to identify the 20 nearest observations and their corresponding actions. We then fit a normal distribution to these actions for each step. Additionally, we sample from the action prior and decoder to generate actions and fit a normal distribution to these for each step as well. Thus, we end up with an approximation of the dataset’s action distribution and an approximation of the action distribution generated by the prior for each step in every trajectory.\nTo retrieve a single metric for the whole dataset, we calculate the KL divergence for each step and average over the action dimensions, all steps and trajectories. We conduct the evaluation for two different datasets, resulting in an averaged KL divergence of 0.22 for hopper-expert-v2 and 0.27 for walker2d-medium-replay-v2.\nMoreover, we create an additional figure (Rebuttal Figure 2) to show the aggregated distribution (from leaving the ground to touching the ground) for one exemplary trajectory from the hopper-expert-v2 dataset. The selected trajectory has an averaged KL divergence of 0.33.\n\n**2b) Does the predictive power of using latent actions degrade in performance compared to models using real actions?**\n\nThe latent state prior and latent state posterior depend on the deterministic transition $f(h_{t-1}, s_{t-1}, a_{t-1})$. Hence, we use actions from the dataset during model training. This is the same as in a regular RSSM/Dreamer model. Only during policy search, latent actions are decoded and fed into the transition. As a result, the predictive power of our model is comparable.  \n\n**3) Differences between the proposed method and Dreamer.**\n\nC-LAP learns a generative model of states and actions $p(o_{1:T},  a_{1:T})$, while Dreamer learns a conditional model $p(o_{1:T} \\mid a_{1:T})$. Both learn a policy on imagined trajectories. While Dreamer learns a policy directly in the environment’s action space, C-LAP learns a policy in the latent action space and uses the latent action decoder to generate actions. Dreamer works well for online learning, but fails in the setting of offline reinforcement learning. This can be seen in Figure 6, where “no latent action” corresponds to Dreamer. Dreamer does not restrict the policy to generate actions close to the dataset’s action distribution, hence suffers from value overestimation.\n\n**4) Performance on navigation tasks such as AntMaze?**\n\nWe conducted an additional evaluation on D4RL antmaze datasets (Rebuttal Figure 4a). Initially, we did not test our method on these datasets because none of the methods we compare against report results on them. Therefore, we performed a hyperparameter search for MOBILE, varying the penalty coefficient between [0.5, 1.5, 2.5, 3.5] and the rollout steps in [1, 5],  but none of these hyperparameters made MOBILE work on these environments. PLAS is mediocre in the umaze environments but does not work for medium environments at all. For C-LAP we perform a hyperparameter search of $\\epsilon$ in [0.5, 1.0, 2.0, 3.0]. Overall, a more extensive hyperparameter search might improve the results for all methods. \n\nIn the more challenging diverse datasets and large environments, none of the methods proved effective.These datasets contain numerous suboptimal trajectories, are collected with non-markovian policies and are known to require strong Q-learning, such as Max Q Backup [1] in model-free methods, to be successful. We believe this points to an interesting research direction, as datasets with these characteristics have not been thoroughly explored for model-based approaches.\n\n\n[1] Kumar et al., Conservative Q-Learning for Offline Reinforcement Learning, 2020\n\n\n**5) Why do some experiments in Figure 4, 5 show high returns at the beginning of policy learning (zero gradient step)?**\n\nPlease take a look at our general reply.\n\n**6) In table 5 and table 6, the performance of C-LAP is inferior to some baseline models in medium or medium-relay datasets. This results seem to indicate that the effectiveness of the proposed method depends on the quality of the offline dataset.**\n\nYes, especially the effect of jump-starting policy learning by leveraging the latent action decoder to already achieve high rewards in the beginning of policy training clearly depends on the dataset. We will also add this to the limitations.\n\n**This paper discusses the limitations, but that is not enough.**\n\nTo further discuss the limitations of C-LAP, we intend to move lines 214-217 to a separate paragraph in the conclusion. Moreover, we add the following limitations:\n- Depending on the dataset and environment the effectiveness of C-LAP differs.\n- The effect of jump-starting policy learning with the latent action decoder to already achieve high rewards in the beginning of policy training is prominent in narrow datasets, but less effective for diverse datasets.\n- Datasets containing suboptimal trajectories produced by non-Markovian policies, such as those in navigation tasks, pose a challenge for C-LAP\n\nDo you have any further limitations in mind, which we should address?", "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:
First and foremost, we want to thank all reviewers for their time, effort, and insightful feedback! This clearly helped us to improve our paper!

**We conducted further experiments and added the following figures:**

- Updated version of Figure 1 to include the latent action prior and highlight that the policy is generating actions within the support of the latent action space. (Rebuttal Figure 1)
- Comparison of the dataset's action distribution to the distribution of actions sampled from the latent action prior and decoder. (Rebuttal Figure 2)
- Analysis of value-overestimation for all baselines (Rebuttal Figure 3)
- Evaluation of AntMaze environments (Rebuttal Figure 4a)
- Sensitivity analysis of the support constraint parameter epsilon on the walker2d datasets (Rebuttal Figure 4b)

Moreover, we noticed multiple reviewers wondering why C-LAP’s performance is sometimes close to the highest reward already in the beginning of policy training (Figure 4 and Figure 5). 

**Why is C-LAP’s performance sometimes close to the highest reward already in the beginning of policy training?**

The model is trained to generate actions contained in the dataset’s action distribution. If the dataset is narrow, generated actions when sampling from the latent action prior will fall into the same narrow distribution. For instance, in an expert dataset, sampled actions will also be expert-level actions. During policy training, instead of sampling  from this prior, we  restrict the support of the policy dependent on the latent action prior. Thus, sampled latent actions from the policy will always be decoded to fall into the dataset’s action distribution. So even a randomly initialized policy in the beginning of the training can generate a high reward by using the latent action decoder. This effect clearly diminishes if the dataset’s action distribution is less narrow, which can be seen in the extremes, e.g., when comparing medium-replay to expert datasets in Figure 4. Thus, policy training is more essential for medium-replay datasets to only sample actions leading to high rewards.

Meta-review: This paper introduces Constrained Latent Action Policies (C-LAP), an approach for model-based offline RL designed to handle value overestimation. C-LAP employs a latent action state-space model to generate actions that are constrained within the support of a learned action distribution, mitigating the risk of generating out-of-distribution actions. By formulating the policy learning as a constrained optimization problem, C-LAP eliminates the need for additional uncertainty penalties, thereby speeding up the learning. The approach is tested on benchmarks with both low-dimensional and visual observations, demonstrating competitive performance, particularly in visual tasks. In summary, C-LAP demonstrates a robust solution to the challenges of offline reinforcement learning by maintaining close distances to the original data distribution while enabling efficient policy training.
Decision: Accept (poster)