Scalable Decision-Making in Stochastic Environments through Learned Temporal Abstraction
Authors: Baiting Luo
Ava Pettet
Aron Laszka
Abhishek Dubey
Ayan Mukhopadhyay
Venue: ICLR 2025 Spotlight (2025)
Keywords: Sequential Decision-Making
Monte Carlo Tree Search
Temporal Abstraction
Planning
Model-based Reinforcement Learning
Offline Reinforcement Learning
Primary area: reinforcement learning
TL;DR: A scalable approach for sequential decision-making in high-dimensional continuous action spaces by learning macro actions and using MCTS.

Abstract: Sequential decision-making in high-dimensional continuous action spaces, particularly in stochastic environments, faces significant computational challenges. We explore this challenge in the traditional offline RL setting, where an agent must learn how to make decisions based on data collected through a stochastic behavior policy. We present \textit{Latent Macro Action Planner} (L-MAP), which addresses this challenge by learning a set of temporally extended macro-actions through a state-conditional Vector Quantized Variational Autoencoder (VQ-VAE), effectively reducing action dimensionality. L-MAP employs a (separate) learned prior model that acts as a latent transition model and allows efficient sampling of plausible actions. During planning, our approach accounts for stochasticity in both the environment and the behavior policy by using Monte Carlo tree search (MCTS). In offline RL settings, including stochastic continuous control tasks, L-MAP efficiently searches over discrete latent actions to yield high expected returns.
Empirical results demonstrate that L-MAP maintains low decision latency despite increased action dimensionality. Notably, across tasks ranging from continuous control with inherently stochastic dynamics to high-dimensional robotic hand manipulation, L-MAP significantly outperforms existing model-based methods and performs on par with strong model-free actor-critic baselines, highlighting the effectiveness of the proposed approach in planning in complex and stochastic environments with high-dimensional action spaces.

Published as a conference paper at ICLR 2025
SCALABLE DECISION-MAKING IN STOCHASTIC EN- VIRONMENTS THROUGH LEARNED TEMPORAL AB- STRACTION
Baiting Luo', Ava Pettet', Aron Laszka", Abhishek Dubey", Ayan Mukhopadhyay' Vanderbilt University, ?Pennsylvania State University, Nissan Advanced Technology Center Spaiting.luo abhishek. dubey, ayan wwt..mmwwwwwaaaayaaamma.a htwe.....oooooha//////aaaaa..... heteraaayyvaadebb....e... alaszka@psu, edu, ava. ava.petratinnisan--use com
ABSTRACT
Sequential decision-making in high-dimensional continuous action spaces, partic- ularly in stochastic environments, faces significant computational challenges. We explore this challenge in the traditional offline RL setting, where an agent must learn how to make decisions based on data collected through a stochastic behav- ior policy. We present Latent Macro Action Planner (L-MAP), which addresses this challenge by learning set of temporally extended macro-actions through a state-condinn contecondiiional Vector Quantized Variational Autoencoder (VQ-VAE), effec- tively reducing action dimensionality. L-MAP employs a (separate) learned prior model that acts as latent transition model and allows efficient sampling of plau- sible actions. During planning, our approach accounts for stochasticity in both the environment and the behavior policy by using Monte Carlo tree scarch (MCTS). In offline RL settings, including stochastic continuous control tasks, L-MAP ef- ficiently searches over discrete latent actions to yield high expected returns. Em- pirical results demonstrate that L-MAP maintains low decision latency despite increased action dimensionality. Notably, across tasks ranging from continuous control with inherently stochastic dynamics to high-dimensional robotic hand ma- nipulation, L-MAP significantly outperforms existing model- methods and performs pa-par with strong model-free actor-critic baselines, highlighting the ef- fectiveness of the proposed approach in planning complex and stochastic envi- ronments with high-dimensional action spaces.
INTRODUCTION
Planning-based reinforcement learning (RL) has achieved remarkable success in domains with dis- crete, low-dimensional action spaces, such as board games and video games (Silver et al. 2017] Schrittwieser et al. 2020 Ye et al.) [22T), and continuous control tasks (Hubert et al. 2021 Schrit- Schritt twieser et al. 2021) However, extending these methods to high-dimensional continuous action spaces, especially in stochastic environments, presents significant challenges. Many environments inherently stochastic appear stochastic agents with limited capacity to model complex dy- namics. For example, in autonomous driving, the behavior of other vehicles and pedestrians intro- duces substantial uncertainty that must be processed in real time (Carvalho et al.) 2014 Luo al. 2023a). Recent planning-base based offline RL approaches like Trajectory Transformer (Janner et al, 2021 face significant latency issues when trying model and respond to these stochastic behav- 1ors (i et al. 2023). Similarly, in robotic manipulation, sensor noise introduces randomness that requires fast adaptive responses (Yang al. 2023). In such cases, deterministic models often fail to capture the necessary randomness and intricacies (Antonoglou et al., (202)). Moreover, the vast and uncountable nature of continuous action spaces makes traditional planning approaches inefficient when operating directly in the raw action space (Jiang et al. 2023). These inefficiencies are further exacerbated in stochastic settings, leading to "large and long planning problems where agents must manage numerous continuous variables over extended time horizons. This results in the curse of dimensionality and the curse of history, significantly hindering effective decision making (Hubert al. 2021)

---

Published as a conference paper at ICLR 2025
MCTS
L-MAP
Higher Performance Lower Latency
68-
Performance (96)
64
60-
56-
Lower Performance Higher Latency
Decader
0.05
0.10
0.15
0,20
0.25
0.30
Decision Latency (s)
(a) Planning with Pre-constructing search space
(b) Decision Latency vs. Performance
Figure I: (a) Overview of planning over the pre-constructed search space. (b) As the number of MCTS iterations increases (10, 50, 100 from left to right), using pre-constructed search space with MCTS achieves better performance with lower decision latency.
In this paper, we posit that planning in such challenging settings could greatly benefit from temporal abstractions, i.e., representations of multi-step primitive behaviors such as macro actions (Dietterich, 2000] [Sutton al. 1999 Barto Mahadevan 2003).. By leveraging these abstractions, planners can navigate high-dimensional continuous action spaces more efficiently, potentially mitigating the curse of dimensionality and reducing decision-makee latency in stochastic environments. This pa- per considers the standard setting where agent access set of trajectories (i.e., sequence of state, action, and reward traces) collected through a fixed behavior policy. Given this setting, we propose the Latent Macro Action Planner (L-MAP), which constructs a lower dimensional repre- sentation of temporally extended primitive actions by using a state-conditioned Vector Quantized Variational AutoEncoder (VQ-VAE) (van den Oord al., 2017). The encoder integrates the current state and macro-action to generate a discrete latent code. Subsequently, a sequential model (in our case, Transformer) employed autoregressively model the distribution of these latent codes, conditioned on the current state (and the behavior policy). This Transformer facilitates two-step inference process: initially, given state, it enables the sampling of probable latent macro-actions under the behavior policy, effectively acting as a prior policy. Subsequently, conditioned on both the state and the sampled macro-action, generates subsequent latent codes that encapsulate infor- mation about expected returns and potential next states. This dual functionality of the Transformer enables efficient exploration promising action trajectories while forming a compact representation of the plausible trajectories during planning.
As shown in Fig Ta leveraging these models, we build latent search space that serves as a struc- tured initialization for planning, encapsulating likely trajectories based the learned environment dynamics. To address stuchasticity and optimize decision-making we integrate Monte Carlo Tree Search (MCTS) with progressive widening efficiently navigate this latent space. Initially, the search concentrates on the prebuilt latent space, facilitating rapid decision-makinn grounded in learned abstractions. If additional computation time becomes available, we progressively widen the search tree to extend the search beyond the prebuilt latent space incrementally. This dynamic ex- pansion strategy enables our method balance rapid planning using learned abstractions with more exhaustive exploration when computational resources permit. As shown in Fig 1b] this strategy achieves better performance with lower decision latency compared planning with vanilla MCTS. Upon selecting latent macro-action, we operate in a polling control mode (He et al. [001] Gabor et al, 2019) wherein MCTS returns only the first primitive action of the recommended macro-action. This approach allows for recovery from locally suboptimal decisions by performing planning at each time step.
We evaluate L-MAP extensively in the offline RL setting across diverse range of tasks. In stochas- tie MuJoCo environments (Rigter et al. 2023), L-MAP consistently outperforms both model-based baselines like Trajectory Transformer (TT) (Janner et al., 2021) and Trajectory Autoencoding Plan- ner (TAP) (Jing et al. 2023). as well as model-free methods such as Conservative Q-Learning (CQL) Kummrr et al. 2020) and Implicit Q-Learning (IQL.) (Kostrikov et al. 2022). This demon-
2

---

Published as a conference paper at ICLR 2025
strates L-MAP's robust capability in handling stochastic dynamics. For deterministic MuJoCo tasks, L-MAP shows comparable or superior performance to these baselines, highlighting that our planning approach effectively accounts for stochasticity in the behavior policy, leading to competitive perfor- mance even in deterministic environments. Notably, L-MAP scales effectively to high-dimensional tasks, as evidenced by its strong performance on the challenging Adroit hand manipulation tasks. Furthermore, L-MAP's use of temporal abstraction enables lower latency decision-making com- pared to methods like TT. These results underscore L-MAP's versatility and effectiveness across various types of control problems, from stochastic to deterministic environments, and from low to high-dimensional action spaces.
PRELIMINARIES
We consider a continuous state and action space Markov Decision Process (MDP) defined by {S, [S,A, P.r), where S _ IR" is the state space, A C R'' is the action space, : x A + ((S) is the transition function, and r: x A - TK the reward function. To manage the complexity of these continuous spaces, we introduce macro actions, which are fixed-length sequences of primitive actions. A macro action m E M is defined as m = (a,,...,,::+ --) where each a, E A and is the length of the macro action. Our goal is to compute optimal macro-level policy n² 8 + Pm that maximizes the expected discounted return E, [R(s, ग (s))]-
L (K E N which is composed of a sequence of states St E S. fixed-size macro actions my E M, and corresponding retur-to-go estimates Rs = 2"=.7" formally represented as = (R,, $1, m1, RL+1, $1+1, ML+1..... R 1)x-1) (x-
Trajectory Representation: Consider trajectory of length T
METHOD
Planning in continuous action space is hard and computationally challenging, and full enumera- tion of all possible actions is infeasible. Discretizing the action space is one way to address this challenge, but in practice, enumerating large set of discrete actions can also be challenging, partic- ularly for online approaches. bample-based methods offer an efficient approach for handling large and complex domains. These methods sample subset of actions rather than exhaustively enu- merating all possibilities. reducing computational costs while computing optimal policies or value functions (Hubert et al. [00II.. Building on these insights, we propose the Latent Macro Action Planner (L-MAP), which learns temporal abstractions in the form of macro-actions and plans using a latent transition model that serves as both a prior policy and a transition model.
3.1 DISCRETIZING STATE-MACRO ACTION SEQUENCES WITH VQ-VAE
A key insight from prior work is that a learned caaee-oonditioned discretization can be used to con- struct a discretization scheme with relatively few discrete actions while maintaining high granular- ity (Jiang et al. 2023 Luo et al., 2023b). As shown in Fig our approach leverages a learned state-conditioned discretization to enable planning in a lower-dimensional discrete space. Specifi- cally, our encoder processes sequences of state and macro-actions defined as It = (Re, St, m²) and its subsequent token as L = (Ry+L, S14L, M4+)). The encoder function is defined as:
input. For example, each token
Senc (x, = (Re, mc), Fe+ = (Re+L, S14 1.)) )) = (zt; Zt+1);
(I)
where the transition chunk size is two, resulting in two latent codes assigned per chunk. To elaborate, the encoder first concatenates the input return-to-go estimates, states, and macro actions into two transition vectors. It then applies sequence model, in our case, a causal Transformer, producing two latent feature vectors for each chunk of transitions.
In stochastic environments, executing the same macro-action m from state can yield different returns R, introducing variability that complicates the vector quantization VQ-VAE, i.e., note that using the full token 2, (Re, St, m.) directly (s,, m,) pairs solely due to differences in R, This challenge can cause the latent space to become fragmented and reflect return variability more than the underlying structure of available actions.
result in different latent codes for identical
3

---

Published as a conference paper at ICLR 2025
R
Kays tn Re
Ray Ri
Rs im Rai koy Ray
Enc
Ross t- Re43: Beads Staes
R
Kayu Rives Ma RaL-oa-1
Rup Ri
Codobooo
Figure 2: An overview of our VQ-VAE model that discretizes state-macro action sequences
Consequently, the agent might overestimate the returns during decision-making by emphasizing latent codes associated with higher observed returns, neglecting the true distribution of the primitive actions and their expected returns.
To address this issue, we aim focus the vector quantization process primarily on representations of the state 8 and macro-actions m, while still preserving the ability to reconstruct the return R. To tackle this challenge, our approach involves creating two versions of each token zy; the full input 24 = (Rt, St. m²) and masked version 2.Task = (mask, st, m2), where R₂ is masked out. The encoder processes both x, and xi a:mask to generate two embeddings, =<(3++) and ze (x:faakk) respectively. We use (x:"aask) for vector quantization to obtain the quantized latent code (mpkkk To ensure that the codebook embeddings incorporate information from the full input, including RG, we update of the quantized latent code towards ze (zz). We modify the loss function by
the embedding introducing an additional term that encourages the embedding from the masked input to be close to that from the full input. Specifically, used the embedding from the full input zx(xx) as the learning target for the embedding of the masked input že (x (a.aa)). The modified loss function is:
C - lo= loggoope p(æ 2n (x + |e
Ze (x)] - -ell2 + B1| (x: sg(e] 12 + (2 (asksk) - ze (2) (2)
where sg denotes the stopgradient operator and the additional term || z=(xx)1 acts as regularizer that aligns the embeddings of the masked and full inputs.
Incorporating macro-actions within each token is critical, as it enables the model to capture tempo- ral dependencies across multiple time steps without the need for downsampling. This approach is particularly important in stochastic settings, where downsampling techniques that aggregate states (as in Jiang et al.) (2023) ) can obscure the stochasticity imposed by the environment's 's dynamics. The decoder takes the initial state and latent codes as inputs, and outputs the reconstructed trajectories:
fdee (St, 214 = (ā: = (Re, St, me), LetL = (Re-L. SiLL, m.)L)..
(3)
The decoding process can be seen as the inverse of the encoding process, except that the initial state is merged into the embeddings of the codes with linear projection before decoding.
Latent Transition Model: Following the discretization process, the subsequent step involves mod- eling sequences of latent codes in autoregressive manner using causal Transformer. The Prior Transformer is conditioned on the initial state 8,, achieved by adding the state feature to all token embeddings (Jiang al. 2023). Primarily, functions as a transition model in the latent space, en- abling the sampling of the next latent code 2;+1 conditioned on the current code 2, and state s. This transition, represented as T Sx - Z, implicitly captures the full R S M - x S M transition in the original space, as each z encodes information about the return-to-go, state and macro-action. Additionally, p(z 8) acts as a prior policy for efficient action sampling, allow- ing rapid selection of probable macro-actions based on learned behaviors from the offline dataset. By operating in the learned latent space, the model potentially reduces computational complexity compared to modeling transitions in the original state-action space, especially for diseesnnnnii environments. The discrete nature of the latent space allows for efficient sampling, which can be beneficial for downstream tasks such as planning.
4

---

Published as a conference paper at ICLR 2025
3.2 PLANNING WITH A LATENT MACRO ACTION MODEL
Planning in high-dimensional environments using learned discrete representations introduces un- certainties from multiple sources. First, the representation learning process introduces uncer- lainty due to the non-injective mapping from the high-dimensional state-action space to lower- dimensional latent space. This can result in many-to-one correspondences, where multiple distinct high-dimensional inputs map to the same latent representation, creating apparent stochasticity even in deterministic environments. Second, the environment itself may be inherently stochastic. The detailed analysis of these uncertainty sources is provided in Appendix
We argue that taking expectations over latent transitions is beneficial in mitigating all these sources of uncertainty, regardless of whether the environment is deterministic or stochastic. By considering the expected outcomes over multiple latent transitions, we can average out the randomness intro- duced by the non-injective mapping and inherent stochasticity, leading to more reliable planning decisions. This insight applies broadly to planning methods that employ models with non-injective mapping characteristics. Building on this insight, we employ Monte Carlo Tree Search (MCTS) as our planning algorithm to mitigate the impact of stochasticity arising from non-injective mappings and potential environmental randomness. MCTS iteratively explores the latent space and takes ex- pectations over transitions, allowing for robust planning in the presence of uncertainty.
Pre-constructing the Latent Search Space. Our approach leverages a learned latent transition model to generate and evaluate macro actions for planning efficiently. Starting from an initial state sp. we sample M latent codes 2, each representing a potential macro action. For each sampled latent code 2, we sample N subsequent latent codes 2' to simulate possible future tra- jectories, capturing the outcomes of these macro actions. We obtain the corresponding state- action transitions and return estimates by decoding these latent pairs (2, conditioned on s. To construct the planning tree efficiently, we cache the initial state So along with the top-k latent codes (and their associated information) based on the de- coded returns, where k À x M and E (0, 1] controls the expansion ratio of the tree. The cached latent codes represent the most promising macro ac- tions to consider from the initial state. The latent codes a' are then decoded to obtain a set of recon- structed tokens, i.e., (R, s, m)). For each of these states s, we sample latent codes z", represent- ing potential macro actions from (note that B and M are exogenously defined syyee--arameters).. This process is recursively applied, allowing to pand the planning tree while controlling its growth through the parameter A. By focusing on the most promising macro actions at each state, we maintain a compact and informative planning structure that ef- ficiently explores the state-action space at macro level.
So
M p(z su)
pl $0.2)
Dec
top
Selection. Starting from the cached tree structure, MCTS iteratively expands and evaluates nodes, al- lowing for more comprehensive exploration of the state-action space. For each state 8 in the tree, MCTS selects one of the top-k cached latent codes 2 based on the Upper Confidence Bounds for Trees (UCT) (Kocsis & Szepesvári 2006); UC ning.
Figure 3: Pre-construction of the latent search space by sampling and evaluating la- tent macro-action codes, caching the top- candidates, and recursively expanding the planning tree for efficient macro-level plan-
Q(s,z)+ + +c/ [og(((1)) where Q(s,=) represents the [s.=)
value of executing macro action z in state (estimated through the decoded return-to-go), (s)) denotes the number times state has been visited, N (ss,) 2) denotes the number of times maero actionz has been chosen in state and c is exploration coefficient.
Progressively Widening the State Space for Search. Despite these powerful abstraction tech- niques, the search space remains challenging due to the underlying high-dimensional nature of the
5

---

Published as a conference paper at ICLR 2025
Repeat
Selection
Expansion
Backpropagation
if 12en, m| < N(t,))
Yes
No
R
&a
Figure 4: Illustration of our MCTS process for macro-level planning- The algorithm iteratively selects actions using the UCT policy, applies progressive widening to balance exploration and ex- ploitation, performs parallel expansion of multiple macro actions and their potential outcomes, and backpropagates estimated Q-values to efficiently explore and refine the planning tree.
original state space, residual stochastic characteristics of transitions in the abstracted space, and the complexity of long-horizon planning scenarios. we were to apply MCTS directly to this ab- stracted space, would encounter two main issues: inefficient utilization of our pre-built search space, with the search potentially diverging prematurely into unexplored regions, and difficulty in building sufficiently deep trees for high-quality long-term decision-making particularly in areas of high stochasticity or uncertainty (Couêtoux et al. (01T). Therefore, we use progressive widening extend MCTS to incrementally expand the search tree. It balances the exploration of new states with the exploitation of already visited states based on two hyperparameters: E [0, 1] and E R+ Let [C(s, 2) denote the number of children for the state-action pair (s, 2). The key idea is al- ternate between adding new child nodes and selecting among existing child nodes, depending on number of times state-action pair (s, has been visited. A new state added to the tree [C(s,=)] < ((s =)), where N(s, =) is the number of times the state-action pair has been visited. The hyperparameter a controls the propensity to select among existing children, with a = lead- ing to always selecting among existing child and 1 leading to vanilla MCTS behavior (always adding a new child). In this way, we could enhance our approach by efficiently utilizing the pre-built search space, prioritizing the exploration of promising macro actions while allowing for incremen- tal expansion of the search tree. This technique enables our method to make quick decisions in an anytime manner, leveraging the cached information, and further refine the planning tree if additional time is available.
Expansion. In approach, the expansion phase differs from standard MCTS by performing par- allel expansion of multiple nodes from a leaf node. From the Icaf node, a set of B latent codes [ sampled, each representing a distinct macro action, drawn from a latent transition model p(> s) to ensure diverse action space coverage. For each sampled macro action 2() N subsequent latent codes (
Jz=1 are sampled according to z/(dd) ~ p(2' 2(0), 2(0),)) ()),)) for = (1...., N). modeling potential outcomes and capturing the stochastic nature of macro actions. These latent transitions are then decoded to obtain the resulting next states {8/(.J1 for each macro action.
Finally, the search tree is expanded by adding all child nodes {()((,)), (()(6)))22(63)))) ( for each macro action 20) to the current leaf node 8. This breadth-wise expansion enables simultaneous exploration of multiple promising macro actions, enhancing the diversity and comprehensiveness of the search and facilitating efficient exploration in complex environments.
Backpropagation. Following the expansion phase, where multiple macro actions are expanded simultaneously, the backpropagation step updates the estimated Q-values based on the returnto-go as shown in Fig
EXPERIMENTS
The empirical evaluation of L-MAP consists of three sets of tasks from D4RL (Fu et al., 2020): gym locomotion control, AntMaze, and Adroit. We compare L-MAP to range of prior offline RI.

---

Published as a conference paper at ICLR 2025
algorithms, including both model-free actor-critic methods (Kumar al. 2020;
Kostrikov et al.
2022) and model-based approaches (Rigter et al.] 2003 Jiang et al. 2023: D Janner et al. 2021
Our work is conceptually most related to the Trajectory Transtormer (TT: Janner et al. (202 () and
the Trajectory Autoencoding Planner (TAP: Jiang et al. (202)), which are model-based planning methods that predict and plan in continuous state and action spaces. These two baselines serve as
our main points of comparison for deterministic environments.
To demonstrate L-MAP's ability to make performant decisions in stochastic environments, we com- pare it with One Risk to Rule Them All (IR2R; Rigter et al.] (2023)), a risk-averse model-based algorithm designed for stochastic domains, and model-free actor-critic methods Conservative Q- Learning (CQL; Kumar et al. (2020)) and Implicit Q-Learning (IQL; Kostrikov et al. (2022)). We evaluate L-MAP on Stochastic MuJoCo tasks (Rigter al. 2023), which serve as a proof of concept in the stochastic continuous control domain.
We then test L-MAP on Adroit, which presents a challenge with its high state and action dimen- sionality. Finally, we evaluate L-MAP on AntMaze, a sparse-reward continuous-control problem. In this task, L-MAP achieves similar performance to TT, surpassing model-free methods. Through these diverse evaluations, we aim to demonstrate L-MAP's versatility and effectiveness across dif- ferent types of control problems, including stochastic environments, nigh-dimensional spaces, and sparse-reward scenarios. Additionally, conduct an ablation study to analyze the impact of key components in L-MAP; detailed results of this study can be found in Appendix
Hyperparameters As for the L-MAP-specific hyperparameters, we set our macro action length to 3. The planning horizon in the raw action space is set to 9 for gym locomotion tasks and 15 for Adroit tasks. These horizons are either smaller or equal to those used in TT and TAP. Our choice of parameters is to ensure a control rate of approximately 10 Hz for locomotion tasks. For each task, we conduct experiments with 3 different training seeds, and each seed evaluated for 20 episodes.
Stochastic Mujoco On the Stochastic MuJoCo tasks, with results presented in Table I L-MAP consistently outperforms the model-based baselines, TAP and TT, across all datasets and en- vironments, demonstrating its superior capacity to handle stochasticity in continuous control tasks. Notably, L-MAP achieves the highest performance in multiple datasets for both the Hopper and Walker2D environments. When compared to 122R, a risk-averse model-based algorithm specifi- cally designed for stochastic domains, L-MAP shows competitive or superior results most cases. An exception is the Hepiuun-RellayHHigh Hopper dataset, where 1R2R attains higher score. This suggests that while L-MAP exhibits robustness across a variety of stochastic settings, there are spe- cific scenarios where risk-averse strategies like IR2R may hold an advantage. Additionally, L-MAP generally outperforms the model-free methods, CQL and IQL. However, CQL surpasses L-MAP in the Medium-Expor Expert-Moor Hopper dataset. It is worth noting that L-MAP is the only method among all baselines that achieves performance comparable to CQL in this specific setting.
Table 1: Results for Stochastic MuJoCo.
Model-Based
Model-Free
Dataset Type
Env
L-MAP
TAP
TT
1R2R
CQL
IQL
Medium-Esee Mod
Hopper
106.11 = 2.16
40.36
= 5.42
$6.10 : 3.33
$2.19
2 8.37
106.17 :: 2.16
60.61
= 3,46
Medium Espert Mod
Walker2D
1.41
91.40
1.42
80.93
56.48
86.66
1.84
Medium- Mod
Hopper
55.07 + 3.06
43.64
+ 2.25
44.49
+ 2.47
65.24
+ 3.31
49.92
+ 3.00
56.00
3,60
Medium- Mod
Walker2D
52.94 1.57
44.46 1.82
43.61
2.15
65.16 - 2.84
49.38 + 2.02
48.82
2.31
Medium-Roge Roplay-Mod
Hopper
52.30 2.65
3.21
37.85
1.19
22.82
2.08
40.53 + 1.52
49.12 : 3.38
Medium Replay-Mout Mod
Walker2D
51.44 1.65
43.49
2.27
27.43
1 3.33
52.23
2.22
40.24 1 1.67
40.77
1 2.72
Medium- Expert- High
Hopper
66.93 4- 3.46
37.31
+ 3.66
58.04
3.60
37.99
2.71
68.03
3.94
44,83 2.58
Medium- Ceatun-Expert-iie Expert-High
Walker2D
97.18 +: 2.08
91.09 - 2.78
$0.01
- 3.51
32.38
- 4.55
83.18 + 3,70
68.61
3.33
Medium- High
Hopper
3.55
2.69
5.5
33.99
45.21
2.97
2.47
Medium-HH High
Walker2D
68.87 + 2.21
$9.84
4 5.03
32.13
4.51 4.51
69
+ 3.24
47.53 + 3.05
Medium- Reghay-High
Hopper
58.05 3.366
48.69
2.97
39.24
2.16
68.25
3.78
51,70
43,27
2.78
Medium Replay-High
Walker2D
65.87 3.07
55.15 4 3.29
16.55
4 2.17
63.63
1 3.41
50.33 1: 3.88
45.13 2.38
Mean
68.63
$2.53
46.28
48.71
61.47
53.42
D4RL MuJoCo On the deterministic MuJoCo tasks, particularly when compared to established model-free approaches such as CQL and IQL, L-MAP demonstrates notable performance in en- vironments like Walker2D and Hopper, matching exceeding these baselines even in dense reward scenarios as shown in Table 2 This highlights L-MAP's effectiveness across various task structures. When compared to TT, I-MAP consistently delivers comparable results. However, L.-
7

---

Published as a conference paper at ICLR 2025
MAP offers significant practical advantage: its use of temporal abstraction enables lower la- tency decision-making for equivalent planning horizons, resulting in improved efficiency during deployment. Furthermore, L-MAP generally outperforms TAP, suggesting that even in determinis- tic environments, the expectation-bant based planning approach proves advantageous by accounting for stochasticity in the behavior policy. This leads more robust policies and, consequently, superior results.
Table 2: Normalised results for D4RL MuJoCo-v2 following the protocol of Fu al. 2020
Model-Based
Model-Free
Dataset Type
Env
L-MAP
ТAP
TT
1R2R
CQL
IQI.
Medium- Medium-Expert
HalfCheetah
92.14
86.40 2.22
95.0 + 0.2
93,99
1.40
91.6
86.7
Medium Expert
Hopper
105.74 1 2.24
85.55 = 3.83
110.0
2.7
57.40
105.4
91.5
Medium-Exie Expiuu-Exppett
Walker2D
109.35 + 0.08
105.32 = 2.03
101.9
6.8
73.18 6.29
108.8
109.6
Medium
HalfCheetah
45.50 0.10
44.73 0.39
46,9 0.4
73.45 0.15
44.4
47.4
Medium
Hopper
73.90 +
69.14 2.33
+.6
55.49 3.99
58.0
66.3
Medium
Walker2D
80.31 1.20
7.80
72.5
78.3
Medium-Replar
Replly
HalfChectah
38.45 0.80
40.83 0.72
4.5
63.85
0.19
45.5
44.2
Medium- Replay
Hopper
91.18
80,92 = 3.79
936
89.67
1.92
95.0
94.7
Mediumm-Relaa
Walker2D
81.04 2
72.32 3.26
82.6 6.9
90.67
1.98
77.2
77.2
Mean
79.73
70.77
78.88
72.60
77.60
77.32
Adroit Control In the Adroit robotic control tasks, which are characterized by their high- dimensional state and action spaces, our proposed method, L-MAP, demonstrates strong and competitive performance as shown in Table 3 Across the Human, Cloned, and Expert datasets, L-MAP exhibits notable effectiveness compared to both model-based approaches (TAP and TT) and model-free methods (CQL,, IQL, and Behavior Cloning (BC
Table 3: Adroit robotic hand control results.
Model-Based Approaches
Model-Free Approaches
Dataset Type
Env
L-MAP
TAP
TT
CQL
IQL
BC
Human
Pen
76.26 = 8.58
66.86 8.41
36.4
37.5
71.5
34.4
Human
Hammer
1.71 0.12
1.7 : 0.09
0.8
1.5
Human
Door
11.24 = 1.11
9.51 : 1.10
0.1
9,9
4,3
0.5
Human
0.02
Cloned
Pen
60.68 7.88
46.44 7.54
11.4
39.2
37.3
56.9
Cloned
Hammer
143-1:00 0.29
1.32 0.12
0.5
2.1
2.1
0.8
Cloned
Door
13.22
1.43
-0.1
Cloned
Relocate
0.15
0.13
-0.23 - 0.01
-0.1
-0.1
-0.2
-0.1
Expert
Pen
126.60 :1:60:::55 5.60
112.16 6.57
72.0
107.0
85.1
Expert
127.16
Expert
Door
105.24
105.86 1: 0.08
94.1
101.5
34.9
Expert
Relocate
107.57
0.76
106.21 + 1.61
10.3
95.0
101.3
Mean (All)
51.40
49,33
20.08
40.32
14.76
36.73
Mean (Non-Expert)
18.79
17.37
6.13
11.70
14.76
11.74
In the Human dataset, which includes suboptimal human demonstrations, L-MAP achieves the high- est score in the Door cavironment and performs well in other tasks. Although IQL leads in the Pen task and CQL leads in the Hammer and Relocate tasks, L-MAP maintains competitive results, par- ticularly surpassing TT and BC in most environments. This suggests that L-MAP effectively utilizes suboptimal data make robust decisions in complex settings. For the Cloned dataset, which con- tains a mix of optimal and suboptimal trajectories, L-MAP secures top performance in the Pen and Relocate tasks. In the Expert dataset, comprised of optimal demonstrations, L-MAP attains the highest scores in the Pen and Relocate environments while remaining competitive in the Hammer and Door tasks. Overall, L-MAP achieves the highest average score of 51.40 across all datasets and environments, and 18.79 across non-expert datasets, highlighting its effectiveness in handling vary- ing levels of data optimality. Furthermore, the experimental results indicate that L-MAP effectively
We included Behavior Cloning (BC) additional baseline since the original 1R2R method was not evaluated for Adroit tasks.
8

---

Published as a conference paper at ICLR 2025
manages the complexities of high--dmensional Adroit environments. Incorporating more action in- formation into the single token does not detract from performance; instead, it appears to enhance the model's ' ability to learn nuanced temporal dependencies required for successful task execution.
AntMaze In the AntMaze environments set of sparse-reward continuous-contro tasks where an agent must navigate robotic ant target location, L-MAP demonstrates strong and competitive performance as shown in T These tasks are particularly challenging due to the sparse rewards and the presence of suboptimal trajectories that lead various goals other than the target position used during testing.
Table 4: Performance comparison on AntMaze environments. This evaluation demonstrates that our approach can achieve comparable performance TT with separate network, while being more efficient during sampling and decision-maiee maaonn...
Dataset Environment
BC
CQL
IQL
TT (+Q)
TAP
L-MAP
Umaze AntMaze
54.6
74.0
87.5
100.0 0.0
78.33 = 5.32
93.33 + 3.22
Medium-Play AntMaze
0.0
61.2
71.2
93.3 - 6.4
43.33 6.40
75.00 + 6.85
Medium--iverss. AntMaze
0.0
53.7
70.0
100.0 0.0
30.00 = 5.92
88.33 + 4.14
Large-Play AntMaze
0.0
15.8
39.6
66.7 + 12.2
63.33 6.22
78.33 + 5.32
Large-Divers Diverse AntMaze
0.0
14.9
47.5
60.0 + 12.7
66.67 + 6.09
81.67 + 5.00
Mean
10.92
43.92
55.16
84.00
56.33
83.33
Similar to TAP, our approach integrates goal positions into the observation space, allowing it to condition trajectory generation on specific goals. This conditioning narrows the focus of sampled trajectories towards the target direction, simplifying the planning process. Instead of using the IQL critic for value estimation, L-MAP leverages Monte Carlo planning to provide refined value estimates. This alternative approach avoids the additional computational cost of sampling with separate Q-network, as required by TT (+Q).
Our method achieves an average success rate of 83.33% across all AntMaze environments, which comparable to the 84.00% average of TT (+Q). Notably, L-MAP outperforms TT (+Q) in the more complex Large-Play and Diverse environments, achieving success rates of 78.33% and 81.67% respectively, compared TT (+Q)'s 66.7% and 60.0%. This indicates that L-MAP is partic- ularly effective in larger mazes where navigation complexity is higher. While TT (+Q) attains per- feet success rates in smaller environments like Umaze and Medium-Diverse, L-MAP still performs exceptionally well with success rates of 93.33% and 88.33% in these settings. This consistency suggests that our method robust across different scales of environment complexity.
RELATED WORK
Recent advancements in reinforcement learning focus on learning temporally extended action prim- itives reduce decision-making horizons and improve learning efficiency. Both model-free and model-based methods leverage temporal abstraction to manage task complexity.
Model-free methods such as CompILE (Kipf et al. 2019), RPL (Gupta al. 2019). OPAL (Ajay et al. 2021), ACT (Zhao et al. 2023), and PRISE (Zheng et al. 2024) leverage temporal abstraction various ways. For instance, ComplLE learns latent codes representing variable-length behavior segments, enabling cross-task generalization. RPL employs a hierarchical policy architecture to simplify long- monghorizon tasks by decomposing them into sub-policies. OPAL introduces a continuous space of primitive actions to reduce distributional shift in offline RI, enhancing policy robustness. PRISE applies sequence compression to learn verithle-length action primitives, improving behavior cloning by capturing essential behavioral patterns. These approaches demonstrate the versatility of temporal abstraction addressing different challenges in reinforcement learning, particularly in managing the complexity inherent in sequential decision-making.
From a mosel-based perspective, recent work has treated reinforcement learning as a sequence mod- eling problem, utilizing Transformer architectures to model entire trajectories of states, actions, wards, and values. This approach is exemplified by methods like Trajectory Transformer (TT) (Zhou) et al. 2020), and TAP (Tiang et al. 2023). TAP, in particular, shares conceptual similarities with

---

Published as a conference paper at ICLR 2025
our proposed method, L-MAP, in its use of efficient planning solutions for complex action spaces. These sequence modeling approaches have shown promise in capturing long-term dependencies and handling the variability in trajectories, but they often face challenges in stochastic environments where the outcome is not solely determined by the agent's 's actions. As highlighted by |Paster et al. (2022), reinforcement learning via supervised learning methods may replicate suboptimal actions that accidentally led to good outcomes due to environmental randomness. To address this issue, they proposed ESPER, solution inspired by the decision transformer framework (Chen et al. 2021), ) ESPER mitigates the influence of stochasticity on policy learning in discrete action spaces by clus- tering trajectories and conditioning on average cluster returns.
From a theoretical perspective, several foundational works have studied continuous-space RL via maanilton-JacoiiBaae
equations. For example, Kim et al. (202T) grounded Q-learning and DQN in this theory. characterizing optimal control without explicit optimization, [Munos] (2000) estab- lished convergence results using viscosity solutions, and Hann et al. (2017) employed deep learning to solve high-dimensional PDEs via backward stochastic differential equations. While providing crucial theoretical foundations, these works focused on deterministic environments or required per- fect knowledge about the dynamics of the environment.
Our approach also has interesting connections to robust RL, though with key distinctions. While robust MDPs (lyengar 2005 Nilim & Ghaoui 2005) deal with varying transition kernels chosen adversarially from uncertainty sets, our work focuses on learning and planning with a fixed tran- sition kernel in offline setting where environmental stochasticity is captured through learned models, Early robust RL addressed planning with known dynamics in tabular settings (Xu & Man- nor 2010), and generalizing continuous, nigh-dimensional spaces challenging (Lim & Autef, 2019). Our temporal abstraction could complement robust RL by providing structured transition functions, potentially integrating classical robust RL planning into high- -ige-dimessional environments.
From a plamning perspective, our work relates to methods like MuZero (Schritwwieser et al. 2020). stochastic MuZero (Antonoglou et al., 2022). and Vector Quantized Models for Planning (Ozair et al., 2021). which primarily operate in discrete action spaces and online settings, limiting their applicability to continuous control tasks in offline RL. MuZero Unplugged (Schrittwieser et al.] 2021) extended MuZero to the offline setting and adapted to low-dimensional continuous action spaces using factorized policy representations (Tang & Agrawal 2020). However, scaling to high- dimensional action spaces is challenging due computational infeasibility and imprecise action selection (Luo et al., 2023b). Additionally, MuZero Unplugged focuses on deterministic environ- ments and may struggle in highly stochastic continuous settings.
Our method, L-MAP, extends these concepts to high-dimensional continuous action spaces by ef- fectively handling stochasticity and complexity. Using an encoder to group similar state-macro- action pairs and reconstructing return-to-go estimates via a decoder within the VQ- VAE framework, L-MAP captures essential dynamics while abstracting unnecessary details. This approach models future returns more accurately in stochastic settings. Combined with planning algorithms, L-MAP refines expected return estimates. bridging the gap between temporal abstraction techniques and ro- bust performance in stochastic environments. Our latent code representation and transition model reduce the need to learn separate policy, dynamies, and value components in the offline setting, increasing planning efficiency and accounting for environmental stochasticity, thereby enhancing generalization across complex tasks.
DISCUSSION AND LIMITATIONS
In conclusion, we introduced the Latent Macro Action Planner (L-MAP), which leverages temporal abstractions learned with a condiioond VQ-VAE to construct a discrete latent space of macro- actions. This approach enables efficient planning in high-dimensional continuous action spaces within stochastic environments. Future directions include exploring transfer learning to handle new tasks, and adapting L-MAP to online learning scenarios through strategies such as risk-averse ex- ploration (Luo al. 2024). These enhancements would enable continuous improvement and help tackle more complex challenges, ultimately improving generalization and efficiency in complex, real-world settings.
10

---

Published as a conference paper at ICLR 2025
7 ACKNOWLEDGEMENTS
This material is based upon work sponsored by the National Science Foundation (NSF) under Grant CNS-2238815 and by the Defense Advanced Research Projects Agency (DARPA) under the Assured Neuro Symbolic Learning and Reasoning program. Results presented in this paper were obtained using the Chameleon testbed supported by the National Science Foundation. Any opinions, findings, conclusions, recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the NSF or the DARPA.
REFERENCES
Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum. OPAL: offline prim- itive discovery for accelerating offline reinforcement learning. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
Robert Almgren and Neil Chriss. Optimal execution of portfolio transactions. Journal of Risk, 3: 5 40, 2001.
Ioannis Antonoglou, Julian Schrittwieser, Sherjil Ozair, Thomas K. Hubert, and David Silver. Plan- ning in stochastic environments with learned model. In The Tenth International Conference on Learming Representattions, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
URL http:///eaaaaaaaeaoeeee/ecccecna
net http:/eee//ttteeeeeeeeee /tte:/ooommm11xX000llalla
Wenhang Bao and Xiao-yang Liu. Multi-agent deep reinforcement learning for liquidation strategy analysis. arXiv preprint 19X1::990.1114/// 2019.
Andrew G Barto and Sridhar Mahadevan. Recent advances in hierarchical reinforcement learning. Discrete event dynamic systems, 13:341-379, 2003.
Ashwin Carvalho, Yiqi Gao, Stéphanie Lefèvre, and Francesco Borrelli. Stochastic predictive
control of autonomous vehicles in uncertain environments. 2014. URL https://api. semant www.entisschlllr artparoorrr//oopsr::1111111
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Rein- forcement learning via sequence modeling.
Marc'Aurela Ranzato, Alina Beygelz- imer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Anmual Conference on Neural Informa- tion Processing Systems 2021, NewrlPS 2021, December 6-14, 2021, virtual, pp. 15084 15097. 2021. URL 77499662200000002225553333
hew.:eeeeeddddn................. httpp//.....rrrpppppppppppee/eee wwtpppppe//22//aaa/eeee html
Adrien Couëtoux, Jean- Baptssee Hoock. Nataliya Sokolovska, Olivier Teytaud, and Nicolas Bon- nard. Continuous upper confidence trees. In Carlos A. Coello Coello (ed.), Learning and
Intelligent Optimization 5th International Conference, LION 5, Rome, Italy, January 17- 21, 2011, Selected Papers, volume 6683 of Lecture Notes in Computer Science, pp. 433 Springer, 2011. doi: 978 978-3-644 3-----42-22552 25566-3566632 32]
URLhttps:: /ectts:://ci.... w0t.di.ooooo....10 1007/
Thomas G Dietterich. Hierarchical reinforcement learning with the maxq value function decompo- sition. Journal of artificial intelligence research, 13:227-303, 2000.
Damien Ernst, Guy-Bart Stan, Jorge Goncalves, and Louis Wehenkel. Clinical data based optimal sti strategies for hiv: a reinforcement learning approach. Proceedings of the 45th IEFE Conference on Decision and Control, pp. 667-672, 2006.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement leamiin. CoRR, abs/2004.07219, 2020. URL https: /arxiv. org/ 2001//220040001199
11

---

Published as a conference paper at ICLR 2025
Thomas Gabor, Jan Peter, Thomy Phan, Christian Meyer, and Claudia Linnhoff- Poopie.. Subgoal- based temporal abstraction in monte-carlo tree search. In Sarit Kraus (ed.), Proceedings of the Twenty-Eighh International Joint Conference Artificial Intelligence. IJCAI 2019, Macao, China, August 10-16, 2019, pp. 5562-5568. org. 2019. doi: 10.24963/UCAA12219777 URI. . 249
https:
2019/772
Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. In Leslie Pack Kaelbling, Danica Kragic, and Komei Sugiura (eds.), 3rd Annual Conference on Robot Learn- ing. CoRL. 2019, Osaka, Japan, October 30 November 1, 2019, Proceedings, volume 100 of Proceedings of Machine Learning Research, pp. 1025-1037. PMLR, 2019. URL http: //proceedingg ngs. ml w.w......ss.............. wwwpppessslvl... v010//l00000// gupta2Ca. wtrrss//1000/gupa22aa..
html
Jiequn Han, Amulf Jentzen, and Weinan E. Overcoming the curse of dimensionality: Solving high-
dimensional partial differential equations using deep learning. CoRR, abs/1707.02568, 2017.
URL http:
org/ abs/ 1tt./17770000 1756802568)
Ruijie He, Emma Brunskill, and Nicholas Roy. Efficient planning under uncertainty with macro- actions. J. Artif. Intel.. Res., 40:523-570, 2011. doi: 10.1.13/JAIR..111. URLhttts:://do
org/ 1.10.6613/jai.. 3171
Thomas Hubert, Julian Schrittwieser, loannis Antonoglou, Mohammadamin Barekatain, Simon Schmitt, and David Silver. Learning and planning in complex action spaces. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML, 1021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Re- search, pp. 4476 4486. PMLR, 2021. URL http::httpp:toeeeeeeeoeee hubert2la.htm. html
ml wwl.r.rrr/ssv333
Garud Iyengar. Robust dynamic programming, Math. Oper. Res., 30:257-280, 2005. URL https: httpp//...aaatt.ccccao..rrrrrrrr
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 1273 1286, 2021. URL httpt: httppppps.//ppppeeengg.g..oo neurgids....nsnnn............... cc/ paper/
html
Zhengyao Jiang, Tianjun Zhang, Michael Janner, Yueying Li, Tim Rocktäschel, Edward Grefen- stette, and Yuandong Tian. Efficient planning in compact latent action space. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL httpps://openreveee CA77 7VVNNn)
net forum?id-
Jeongho Kim, Jaeuk Shin, and Insoon Yang. Hamilton-jacobi deep q-learning for deterministic continuous-ting time systems with lipschitz continuous controls. J. Mach. Learn. Res., 22:206:1-
206:34, 2021. URL https: http:ss://m...oorg///peeeeeeeeee www.oo/pppp././////././ 1235.
Thomas Kipf, Yujia Li, Hanjun Dai, Vinícius Flores Zambaldi, Alvaro Banchez-Gonzalez Edward Grefenstette, Pushmeet Kohli, and Peter W. Battaglia. Compile: Compositional imitation learning and execution. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th Machine Learning,
International Conference fornia, USA, volume 97 of Proceedings of Machine Learning Research, pp. 3418-3428. PMLR, 2019. URL URL http: httttt::::ppoceeiiinirn
9-15 June 2019, Long Beach, Cali-
mlr.. wtw.l/..rees////////// http99////kiflm 9a html
Levente Kocsis and Csaba Szepesvári. Bandit based monte-carlo planning. In Johannes Fürnkranz, Tobias Scheffer, and Myra Spiliopoulou (eds.), Machine Learning: ECML 2006, 17th Euro- pean Conference on Machine Learning, Berlin, Germany, September 18-22, 2006, Proceed- ings, volume 4212 of Lecture Notes in Computer Science, pp. 282-293. Springer, 2006. doi: 10.1007/1 187184) 29. URL htt.l:ttps://dett http:///.oooooog....11// httg///..010//111111111
12

---

Published as a conference paper at ICLR 2025
Iyy Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. In The Tenth International Conference on Learning Representations, ICLR 2022, Vir- tual Event, April 25-29, 2022. OpenReview.net, 2022. URL https://epenrevieee nnt//
forum? d 68 TTFB
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: An- nual Conference on Neural Informattio Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020. /ttp////tttttts//eeeeeeaa.aaass. neurips 0d2b2061 82 wwe.e22552f2222222 Grab2012222513221111655552222213 htt.://cc....aaa.oo.ooo....ooo/a
html
Zenan Li, Fan Nie, Qiao Sun, Fang Da, and Hang Zhao. Uncertainty-aware decision transformer for stochastic driving environments. CoRR, abs/2309. 2023. doi: (0.48550AARXIV23099
16397. httphtts: http:///oooo.r.11.11 48550/ artivv/iivvvv.a 2309.
Shiau Hong Lim and Arnaud Autef. Kernel-based reinforcement learning in robust markov decision processes. In International Conference on Machine Learning, 2019. URL https://api.
semanticscholes wttp://orpp11:111779905.33
Baiting Luo, Shreyas Ramakrishna, Ava Pettet, Christopher B. Kuhn, Gabor Karsai, and Ayan Mukhopadhyay. Dynamic simplex: Balancing safety and performance in autonomous cyber physical systems. In Sayan Mitra, Nalini Venkatasubramanian, Abhishek Dubey, Lu Feng, Mahsa Ghasemi, and Jonathan Sprinkle (eds.), Proceedings of the ACM/IEEE 14th International Conference on
Systems, ICCPS 2023, (with CPS-IoT Week 2023), San Antonio, TX, USA, May 9-12, 2023, pp. 177-186. ACM, 2023a. doi: 10.1145/357644413555944 URL https: hops://dos //doi. www.oogggg..oo...ooooooo.......o 10. 011.111111676111133333 718776638113333 3585934
Baiting Luo, Yunuo Zhang, Abhishek Dubey, and Ayan Mukhopadhyay. Act as you learn: Adaptive decision-making in non-stationary markov decision processes. In Mchdi Dastani, Jaime Simão Sichman, Natasha Alechina, and Virginia Dignum (eds.), Proceedings of the 23rd International Conference Autonomous Agents and Multiagent Systems, AAMAS 2024, Auckland, New Zealand, May 6-10, 2024, pp. 1301-1309. International Foundation for Au- tonomous Agents and Multiagent Systems ACM, 2024. doi: 00.5555/3635633.6629888 URL 3662988
https: acm. /tt.://oooo1...555/3533333
Jianlan Luo, Perry Dong, Jeffrey Wu, Aviral Kumar, Xinyang Geng, and Sergey Levine. Action- quantized offline reinforcement learning for robotic skill learning. In Jie Tan, Mare Toussaint, and Kourosh Darvish (eds.), Conference on Robot Learning, CoRL 2023, 6-9 November 2023, Atlanta, GA, USA, volume 229 of Proceedings of Machine Learning Research, pp. 1348 1361. PMLR, 2023b.
mlr. wttpppsess/s//aooeeeep
html
Rémi Munos. A study of reinforcemmet learning in the continuous case by the means of viscosity solutions. Mach. Learn., 40(3):265-29) 2000. doi: 10. 10.1023/A:1000768630999 //doi. 10. 1023/A: 100::100755500 0001166366666000
URL https:
Arnab Nilim and Laurent El Ghaoui. Robust control of markov decision processes with uncertain transition matrices. Oper. Res., 53:780-798, 2005. URL https://api.
semanticscholar. org/ CC:///sosSSSD:1545555555
Sherjil Ozair, Yazhe Li, Ali Razavi, Ioannis Antonoglou, Aäron van den Oord, and Oriol Vinyals. Vector quantized models for planning. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning. ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 8302 8313. PMLR, 2021. ml
URL http::/ppocceddd www.ingw..........e.ee. w.wp.rss/wsv... wttp:/sss/w13////rr2211.. html
Keiran Paster, Sheila A. Mcllraith, and Jimmy Ba. You can't count on luck: Why de- cision transformers and rvs fail in stochastic environments. In Sanmi Koyejo, S. Mo- hamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Process- ing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 December 9, wapper..nippe cttpcc/cppppaaa http:/pppppp_ieeeepeeeeeeeeeeeee obop html
2022, 2022. URL http: /ttpt//t/////pppprrrr...... ott.://405500.111000000000rarar-
13

---

Published as a conference paper at ICLR 2025
Mare Rigter, Bruno Lacerda, and Nick Hawes. One risk rule them all: A renk-sen perspective on model-based offline reinforcement learning. In Alice Oh, Tristan Nau- mann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Informa- fion Processing Systemy 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL E492873719168 1113319966710 (ww.99661159999aa1a... 158511 Abstract Destract-Conferencee Conferenceeeeeeee html
Julian Schritwieser. Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with learned model. Nat., 00000/44188600000055555
588(7839):604 609, 2020. doi:
URL https: :ttps:://o!. orrg/
10. 10_1038//41156-0200000000
Julian Schrittwieser, Thomas Hubert, Amol Mandhane, Mohammadamin Barekatain, loannis Antonoglou, and David Silver. Online and offline reinforcement learning by planning with learned model. Mare Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34; Anmual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 27580 27591, 2021. URL http:ttss: http://pppppppppppeeeeee Abb00 .html
wwwrrps............. ec//
David Silver, Thomas Hubert, Julian Schrittwieser, loannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy Lillicrap, Karen Si- monyan, and Demis Hassabis. Mastering chess and shogi by self-play with a general reinforce- ment learning algorithm, 2017. URL https:// htt///xxvvv.rr/////e/ abw.aais...1..///////. 01815
Richard Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: frame- work for temporal abstraction in reinforcement learning. Artificial intelligence, 112(1-2):181- 211. 1999.
Yunhao Tang and Shipra Agrawal. Discretizing continuous action space for on-policy optimization. In The Thirty- Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Phirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Sympo- sium Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, Febru- AAAI Press, 2020. doi: 10.1609/AAALV33104..... https: //op:.://do.... org. 1609/ aaai w34104..00000
ary 7-12, 2020, pp. 5581-5
URL
Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete repre- sentation learning. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Infor- mation Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 6300-6315, 2017. URL httpps://proceedings neurips www.nnepp...pppppeeer... tttccc//pppa pocc//ppe//201//aaas 7e99af/7e6aa00000e2ee900099999 Abstrss222bbAbttaattte html
Huan Xu and Shie Mannor. Distributionally robust markov decision processes. In John D. Lafferty, Christopher K. I. Williams, John Shawe-Taylor, Richard S. Zemel, and Aron Cu-
lotta (eds.), Advances in Neural Information Processing Systems 23: 24th Annual Confer-
ence on Neural Information Processing Systems 2010. Proceedings of a meeting held 6-
December 2010, Vancouver, British Columbia, Canada, pp. 2505-2513,
Curran Asso-
ciates, Inc., 2010. URL.
neeriisss. cc httpc///pppee paper/2010/har Abstract .tmll
Shuo Yang, George Pappas, Rahul Mangharam, and Lars Lindemann. Safe perception-based control under stochastic sensor uncertainty using conformal prediction. In 62nd IEEE Confer- ence on Decision and Control, CDC 2023, Singapore, December 13-15, 2023, pp. 6072-6078. IEEE, 2023. doi: 10. 001100/0/CC9999...
000.:77533..22130884033 CDC49753. 2 2023. 0000022..1118800 0384075
htllhhtps://det .rgg 1109/
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbecl, and Yang Gao. Mastering atari games with limited data. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang.
14

---

Published as a conference paper at ICLR 2025
and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Anmual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 25476 25488, 2021. / neurips cc/ paper / 2202 htta:////nncccaaaaoooaoooooeaoaa
-Abs html
Tony Z. Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fine-grained bimanual manipulation with low-cost hardware. In Kostas E. Bekris, Kris Hauser, Sylvia L. Herbert, and Jingjin Yu (eds.), Robotics: Science and Systems XIX, Daegu, Republic of Korea, July 10-14, 2023, 2023. doi: 20155660//S..2233. XIX.011.. htRphtts::///o// org/10..gg11.... 15607/RSS. 2023 XTX. 016
Ruijie Zheng, Ching-An Cheng, Hal Daumé III, Furong Huang, and Andrey Kolobov. PRISE: Ilm- style sequence compression for learning temporal action abstractions in control. In Forty-first International Conference Machine Learning, ICML. 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URLpttps::// owwnneeeieeeee / net. wttp:e//oooommmdded22220000
Wenxuan Zhou, Sujay Bajracharya, and David Held. PLAS: latent action space for offline rein- forcement learning. In Jens Kober, Fabio Ramos, and Claire Tomlin (eds.). 4th Conference on Robot Learning, CoRI. 2020, 16-18 November 2020, Virtual Event / Cambridge, MA, USA, volume 155 of Proceedings of Machine Learning Research, pp. 1719-1735. PMLR, 2020. URL
https: http::://poooeeddeee ngs mlr, press / v 55/ zhou2 1b. html
15

---

Published as a conference paper at ICLR 2025
75
66.4
Mean Normalized
63
64.57
62.75
62.51
80
39.39
57.51
57.7
50
1:4AP
RUCT
No Progressive Widhenng
No Parallet Expansion
Lo
Lag
horimm
hortzoneo
VO-VAE toss
Figure 5: Results of ablation studies, where the height of the bar is the mean normalized scores on high noise gym locomotion control tasks.
A ABLATION STUDY
We present analyses and ablations of key hyperparameters such as macro action length, planning horizon, the of pUCT (Silver et al. 2017) versus UCT, and the effect of customized VQ- VAE loss function. Figure summarizes the results from ablation studies conducted on high-noise stochastic MuJoCo tasks.
Macro Action Length
We tested macro action lengths L 1, = 3, and L = 5 evaluate their impact L-MAP's performance. The highest mean score of 68.7 was achieved with L = 3. Increasing to reduced the mean score 64.57, while decreasing it to further dropped it to 59.39. This indicates that a macro action length of optimally balances temporal abstraction and adaptability. A moderate length allows the model capture important action sequences while remaining responsive to en- vironmental changes. Shorter lengths may fail to model temporal dependencies effectively, while longer lengths may hinder quick adaptation in stochastic environments.
Planning Horizon
We assessed the effect of planning horizon by varying the number of planning steps in L-MAP. Reducing the planning horizon to 3 steps (expanding single latent variable) decreased the mean score 57.51, compared to 68.7 with the default longer planning horizon. This demonstrates that longer planning horizon significantly enhances performance by enabling the model to better antici- pate future events and handle uncertainty high-noise stochastic environments.
Tree Search Algorithm: UCT vs. pUCT
We compared standard UCT and pUCT tree search algorithms in I-MAP. UCT achieved a mean score of 68.7, slightly outperforming pUCT, which scored 66.4. While both methods effective, UCT performs marginally better in this context. A possible explanation is that pUCT leverages a learned prior policy to guide exploration, making it sensitive to the quality of the prior. If the prior policy is suboptimal, pUCT may be less effective due to this dependency.
VQ-VAE Loss Function We compared our loss function with the standard loss function without masking (mean scores: 68.7 vs 57.7). Our approach outperforms the standard loss by focusing primarily state and action during vector quantization. This results in less skewed reconstructed
16

---

Published as a conference paper at ICLR 2025
returns and more coherent latent space, accurately capturing action and state distributions. Conse- quently, the model generates more reliable latent representations for reconstruction.
Progressive Widening
We evaluated the impact of progressive widening on MAP's performance. Removing progressive widening led to significant drop in the mean score from 68.70 to 54.77. This substantial decrease demonstrates the importance of controlled state space expansion during planning for a large search space. Progressive widening enables MAP to balance between exploiting existing states in the pre- built search space and incrementally adding new states. Without progressive widening, the search suffers from excessive branching, making it difficult build sufficiently deep trees for meaningful planning in areas of high stochasticity.
Parallel Expansion
We assessed the contribution of parallel expansion by comparing L-MAP's performance with and without this feature. Removing parallel expansion reduced the mean score from 68.7 to 62.75, yielding performance similar to reducing the planning horizon to six steps. This comparison reveals that parallel expansion primarily affects the algorithm's 's ability efficiently explore the search space. Given the same number of MCTS iterations, removing parallel expansion results in less exploration of possible trajectories, reducing the algorithm's planning capability to that of a shorter horizon. This demonstrates that parallel expansion is crucial for maximizing the effectiveness of each MCTS iteration by enabling broader simultaneous exploration of potential outcomes.
B ADDITIONAL STOCHASTIC ENVIRONMENT EXPERIMENTS: HIV TREATMENT AND CURRENCY EXCHANGE
Table 5: Results for HIV Treatment and Currency Exchange.
Model-Based Approaches
Model-Free Approaches
Env
L-MAP
TAP
TT
1R2R
CQL
IQL
HIV
59.08 1.96
51.95 1.98
54.16 = 3.30
56, 15 2.17
59.74 1.11
31.1
Currency
106.78 = 5.00
89.72 = 3,90
79.28 = 2.61
78.52 + 2.08
93.96 + 1.69
89.41 4: 2.83
The HIV Treatment environment, originally introduced by Ernst et al. (206)), simulates treatment planning where an agent controls drug types (RTI and PI) in 6-dimensional state space rep- resenting cell and virus concentrations. The stochasticity arises from varying drug efficacy at each step. The Currency Exchange environment, based on the Optimal Liquidation problem (Almgren & Chriss 2001: [ao & Liu 2019), involves converting currency under stochastic exchange rates that process. Both environments were adapted by Rigter al. (2023) to the offline RL setting, with datasets collected using partially trained and random policies respec- tively.
follow an
For the HIV Treatment domain, L-MAP and CQL achieve comparable strong performance (59.08 1.96 and 59.74 + 1.11 respectively), outperforming other baselines. the Currency Exchange environment, L-MAP substantially outperforms all other approaches, achieving score of 106.78 + 5.00 compared to the next best performer CQL at 93.96 + 1.69. This superior performance demonstrates L-MAP's versatility across different types of stochastic environments.
C LATENT SPACE ANALYSIS
To empirically demonstrate the uncertainties introduced by non-injective mappings, behavior pol- icy, and environmental stochasticity, we generate heatmaps representing the transition probabil- ities between latent codes. We focus on the Hopper environment and consider three datasets: medium-expers medium, um, and medi um-reply. in both deterministic and stochastic settings. The heatmaps are constructed by encoding the state-macro-action pairs into latent codes using our learned representation and visualizing the transition probabilities between these codes.
17

---

Published as a conference paper at ICLR 2025
Current Latent Code 1a
0.30
0.
Current Laten Codee 10
Current Latent Code ID
0.04
0.02
0.00
Next Latent Code ID ()+
Next Latent Code ID Dt
Next Latent Code (2--1
(a) Medium expert
(b) Medium
(c) Medium replay
Figure 6: Heatmaps for Deterministic Hopper Environment (Top 50 Frequent Latent Codes). In each heatmap, the intensity of the color at position (i,j) represents the probability of transitioning from the current latent code display the frequency of each latent code occurring across the dataset with the learned encoder as the current (=,, right histogram) and next (Z++1, top histogram) codes. The observed spread in the heatmaps indicates that, despite the deterministic nature the environment, transitions from a single 2, lead to multiple 2++1-
to the next latent code 2++1 The accompanying histograms
Currant Latent coda to ch3
Current Latent Code 10 Go)
0.05
Current :ateent Code In
0.20
0.05
0.04
0.030
0.03
0.15
0.425
030
0.020
0.55
Next Lanent Code Cl
Next Latent Code ID 12+11)
Next Latent Code ID (2+
(a) Medium Expert
(b) Medium
(c) Medium Replay
Figure 7: Heatmaps for Stochastic Hopper Environment (Top 50 Frequent Latent Codes). The ob- served spread in the heatmaps indicates that inherent environmental stochasticity further contributes transitions from single z, leading multiple 2++1-
C. DETERMINISTIC EN ENVIRONEENT HEATMAPS
In analyzing the heatmaps for deterministic environments shown in Fig. 6 it becomes evident that transitions from a current latent code to multiple next latent codes 2 are not strictly deter- ministic. This observed spread in transitions originates from primary sources: the non-injective nature of the learned representation and the stochasticity of the behavior policy employed dur- ing data collection.
First, the non-injective mapping of the encoder function fene may result in multiple distinct high- dimensional state-macro-action pairs being mapped the same latent code as shown in the his- tograms of Fig 6 Specifically, for different state-macro-action pairs xt (0 .m)) and (2)
(2) (2)
It (5t m ), possible that:
Sene Senc (1: zt.
(1) even though xtt # (2) Consequently, their corresponding next state-mact cattemmaco-aatiee actionn pairs I++1 and (2)
J may differ, potentially leading to different next latent codes upon encoding:
2 = fenc (x (1)
2 = fene (x with 2 2
Second, because the behavior policy To used for data collection may be stochastic, it introduces variability in the selection of macro-accionn at both the current and subsequent time steps. Given state 8t, the behavior policy determines the macro-action m4 as follows:
my Th (m 82).
18

---

Published as a conference paper at ICLR 2025
(1) and m, This stochastic selection can result in different macro-actions m, same state 8t, which naturally introduces stochasticity. Note that even if the encoder maps both (1) (St, ) and 2t (2) (st, m² the same latent code
(2) being chosen from the
fenc (x) fene (xx]()) 24.
the next states (1) istic, i.e.,
841 and (2)
s41 might differ, though the environment dynamics Tew are determin-
i+1 Ton (s,, mi (1)
Si+1 [em m (2)
with 6111 * 9/41
These different next states lead different next state- iaaar-aaction pairs:
(
(s
T+11 = (s m m+1).
I ++
Upon encoding, they may yield different next latent codes:
2+1 = fenc (1,11,
2 = fene(
with 2/-1 * 2(+1
Therefore, even in deterministic environment, the combination of non-injective encoder and a stochastic behavior policy introduces variability in the latent transitions. The heatmaps for de- terministic environments empirically demonstrate this spread, showing that each does not map deterministically to a single 2++1 but rather to a distribution of possible next latent codes.
C.2 STOCHASTIC ENVIRONMENT HEATMAPS
The heatmaps for stochastic environments shown in Fig. exhibit more pronounced spread in transition probabilities. This inherent environmental stochasticity means that for given 8, and mis there are multiple possible next states 8t+1, leading to wider distribution of next latent codes 24+1 upon encoding. When combined with the non-injective mapping of the encoder and the stochasticity of the behavior policy, the uncertainties in the latent transitions are further amplified.
C.3 THE IMPACT OF L REGULARIZATION ON REPRESENTATION FIDELITY
Current Latent Code ID (zc)
Zz
Current Latent Code ID
0.6
0,4
0.10
0.05
0.00
Next Latent Code ID (2:+1)
Next Latent Code (2:+1)
(a) L1 norm
(b) L2
Figure 8: Transition Probability Heatmaps for Medium Environment (Top 50 Frequent Latent Codes). Left: Heatmap depicting transition probabilities when embeddings are regularized using the L1 norm. Right: Heatmap illustrating transition probabilities under 1.2 norm regularization.
Medium-Replay Replay Datasets from the Stochastic Hopper
The heatmaps shown in Frevell distinct patterns between transition probabilities for latent codes encoded by encoders trained with L, and L2 norm regularization in the latent space. The L2 norm demonstrates more distributed transition probabilities, with multiple moobbiityprobbbilityy transi- tions (shown as light blue dots) for each current state, indicating the encoder preserves more granu- lar information. contrast, the L1 norm exhibits highly deterministic transitions for certain latent codes, shown by the predominantly dark purple background with the bright yellow spot approach- ing probability 1.0. This suggests that the encoder trained with L, regularization tends to collapse dissimilar inputs into the same latent code, leading to less nuanced representations.
19

---

Published as a conference paper at ICLR 2025
D ANALYSIS OF PERFORMANCE TRENDS WITH INCREASING STOCHASTICITY
Table 6: Hopper Environment Results with Increasing Stochasticity
Model-Based
Model-Free
Dataset Type
Env
L-MAP
IR2R
Deterministic
Medium Medium-Exxpert
Hopper
105.74 1: 2.24
85.55 à: 3.83
110.0
57.40 6.05
105.4
91.5
Medium
Hopper
73.90
1,91
69.14 = 2.33
61.1 3.6
55.49 +: 3.99
58.0
66.3
Medium-Replay
Hopper
91.18 + 0,56
80.92 = 3.79
91.5 115=:30
89.67 + 1.92
95.0
94.7
(Deterministic)
90.27
67.52
84.17
Moderate Stochasticity
Medium-Expert Expert-Mood
Hopper
2,16
3.46
Medium-Mod
Hopper
55.07 + 3.06
43.64 = 2.25
44,49
2.47
65.24
3.31
49.92 + 3.00
56.00 + 3.60
Meddum--Relaay
71.16 Hopper
$2.30 2.68
38. 10 3.22
1.19 37,85
22.82
2.08
40.53 1.52
49.12 3,38
Mean (Moderate Stochasticity)
40.87
46.15
46.75
65.54
55.24
High Stochusticity
Medium- High
Hopper
66.93 3,46
37.31 :: 3.66
58.04
3.60
37.99 1:
68.03 1: 3.94
44.83 1: 2.58
Medium-Higg High
Hopper
55.32 = 3,56
43.93 = 2.66
41,26
5.53
33.99 0.92
45.21 + 2.97
49.69 + 2.47
Medium Replay-Hligh
Hopper
43.27
Mean (High Stochasticlty)
60.10
43.31
46,18
46.74
54.98
45.93
This section examines how L-MAP and baseline methods respond increasing levels of stochas- ticity in the Hopper environment. Table 6 presents the performance metrics across deterministic, moderate, and high stochasticity settings.
In the deterministic setting, L-MAP achieves a mean score of 90.27, indicating strong performance and outperforming all other model-based methods. Among the baselines, TT attains mean of 87.53, TAP achieves 78.54, and IR2R scores 67.52. The model-free methods CQL and IQL also perform well, with mean scores of 86. and 84.17, respectively. The high scores across all methods suggest that the deterministic environment poses minimal challenges, allowing both L-MAP and the baselines to excel.
As the environment introduces moderate stochasticity, L-MAP's mean performance decreases to 71.16, reflecting a reduction of approximately 21% from deterministic performance. The model- based baselines experience larger declines; TAP's mean drops to 40.87 (a 48% reduction), TT's to 46.15 (a 47% reduction), and IR2R's to 46.75 (a 31% reduction). The model-free methods also suffer performance losses; CQL's mean decreases to 65.54 (a 24% reduction), and IQL's to 55.24 (a 34% reduction). Despite the reductions, L-MAP maintains a higher mean score than all basclines in this setting, indicating better resilience moderate stochasticity among both model-based and model-free methods.
In the setting of high stochasticity, L-MAP's mean further decreases 60.10, representing total reduction of about 33% from the deterministic case. The model-based baselines continue to show declining trends; TAP's mean falls 43.31 45% reduction), TT's to 46.18 (a 47% reduction), and IR2R's ' to 46.74 74( (a 31% reduction). The model-free methods also see further decreases; CQL's mean drops 54.98 (a 36% reduction), and IQL's 45.93 (a 45% reduction). While all methods experience performance degradation, L-MAP consistently outperforms the model-based basclines TAP and TT, and maintains an edge over the model-free methods CQL and IQL. The performance of I.-MAP shows relatively better robustness among the baselines.
The overall trend indicates that increasing stochasticity adversely affects all methods, but L-MAP's performance diminishes a slower rate compared to the other model-based methods. These results suggest that L-MAP more robust to stochastic variations the environment than most of the baseline methods, particularly the model- ones.
E PLANNING HYPERPARAMETERS
For all environments, we utilize the following hyperparameters for sampling during the search pro- 0.1 and 1, which determine the exploration rate of progressive widening; and set
20

---

Published as a conference paper at ICLR 2025
the number of Monte Carlo Tree Search (MCTS) iterations to 100. Detailed parameters for each
environment are presented Table 7
Table 7: Planning Hyperparameters
Environment
M
?
Stochastic MuJoCo
16
0.5
0.99
D4RL MuJoCo
16
0.5
0.99
Adroit
10
0.5
0.99
AntMaze
16
0.5
0.998
Currency
32
4
0.5
0.99
HIV Treatment
1.0
0.99
21


Reviews:
[{"review_id": "xDf6QCWD8c", "reviewer": "Reviewer_q9DP", "summary": "The authors propose a model-based method for planning in stochastic environments called Latent Macro Action Planner (L-MAP). Their method first learns an encoder-decoder architecture to encode return-to-go estimates, states, and a sequence of actions, which they call macro actions. The learned discrete codes are then used to learn a state-conditioned latent prior $p(z|s)$ which is used to seed a modified Monte Carlo Tree Search (MCTS). The MCTS uses progressive widening to expand the search based on the prior initialization using a heuristic based on the number of children in the tree and the number of times a state-action node is visited and uses the prior to sample children for leaf nodes. \n\nThe authors evaluate their method in two stochastic offline reinforcement learning tasks, Hopper and Walker, and compare their method in the standard Hopper, Walker, HalfCheetah, AntMaze, and Adroit deterministic tasks. They also provide ablations over some of their hyperparameters used in their method.", "questions": "How does your method compare to versions of TT or TAP that learn macro actions?\nIn Table 4, you include a column named \"DT\". Is this decision transformer and if so, why do you include it in the results when it is not mentioned in the paper?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The authors propose pre-constructing the latent search space using their learned prior and show benefits over searching without this pre-constructed prior, which is interesting.", "weaknesses": "The proposed algorithm has many layers of complexity and it is unclear which components contribute to the method's success. The two transformer-based baselines that they compare against do not utilize any temporal action compression, begging the question of if they would benefit similarly from predicting macro actions. While the authors ablate over macro action length and planning horizon, they do not provide any experiments supporting their decisions to utilize UCT for selection, progressive widening, or parallel expansion, and do not provide the values for hyperparameters $M,N,B,\\alpha$ in their experiments. \n\nTheir results also leave a lot to be desired. The only environment where the method separates itself from existing methods is in the stochastic environments, but they only test in two such environments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "A2Cm2gnJ7W", "reviewer": "Reviewer_UcYT", "summary": "This paper proposes Latent Macro Action Planner (L-MAP), an offline RL algorithm designed for high-dimensional continuous action spaces in stochastic environments. L-MAP learns a discrete latent space of temporally extended macro-actions to reduce the action dimensionality and simplify the planning process. A causal transformer is learnt in tandem as a latent transition model, predicting future latent codes. Planning is performed using MCTS with progressive widening. Experiments show that L-MAP outperforms model based and model free baselines in stochastic environments while being on par in deterministic environments, highlighting its versatility.", "questions": "- Could you explain the source of stochasticity in the Mujoco tasks used in the paper? Is it primarily from the dataset generating policy or are the environments themselves modified in some way?\n- Why is L-MAP less performant in the deterministic environments? Is there an underlying theoretical reason for lower performance in deterministic settings given that it was more consistently the best performing algorithm in the stochastic setting?\n- Could the authors study how L-MAP compares to the baselines with increasing levels of stochasticity(in either the dataset or transition dynamics) ? This can be for a single environment/task.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- Comprehensive experiments: The paper provides a suite of experiments covering multiple tasks with good model based and model free baselines. \n- The inclusion of the Adroit hand task showcases the scalability of the proposed method to high dimensional problems.\n- The use of macro-actions and a pre-constructed search space allows for faster decision-making. \n- Additionally, L-MAP provides a way to trade off latency and longer planning depending on the user requirements.\n- The ablation study provides insights into the importance of different hyperparameters and design choices, strengthening the paper's conclusions.", "weaknesses": "See questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GqfxmjW8gy", "reviewer": "Reviewer_Wb4D", "summary": "Paper focusses on stochastic planning in offline RL settings (where the data was collected by a stochastic behavioral policy). To avoid the stochasticitiy related with continuous action spaces, paper focuses on a temporally extended (latent) macro-action policy that is obtained via a vector-quantized variational autoencoder (VQ-VAE).", "questions": "+ How is the length of the macro action chosen? Predetermined or randomly assigned in an on-the-go fashion?\n\n+ What are examples of real-world problems with high-dimensional continuous action spaces in the opinion of the author(s)?\n\n+ I think you missed a couple of connecting words between lines 76 and 77. Do you want to fix those?\n\n+ I am curious to see how this offline L-MAP policy behaves in off-policy RL settings. Can you provide an example or two?\n\n+ L127-128: \"particularly for online approaches\"; why was this qualification needed seeing we are dealing only with offline RL in this paper?\n\n+ How did you determine the instanteneous rewards, r_i, for the computation of R in equation (1)?\n\n+ Lines 167-171, what loss function did the authors use in the reconstruction process? There have been instances in reconstruction like this when the choice of loss makes a huge difference in the reconstruction. See https://arxiv.org/pdf/2311.03534\n\n+ What happens if instead of using the l2 norm on the embeddings regularizer, you try the l1 norm? Does it improve the repreesentation's fidelity?\n\n+ Can you supply a proof, empirical results, ot citation for the claim on lines 213-217?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- I love the integration of primitive actions in the raw state space which is then abstracted temporally via encoders (VQ-VAE) to generate a discrete latent code. Sticking a transformer on the latent codes sounds good for creating a distribution over state and primitive action representations.\n\n+ I think it's nice that the authors were thoughtful enough to add a conditionally widening MCTS tree to the search in the latent space to further aid the autoregressive and stochastic continuous L-MAP policy.\n\n- Perhaps supply the examples of stochastic environments that motivated the question you asked in the Introduction.", "weaknesses": "+ I love the motivation in the introduction that many environments appear stochastic but a couple of examples where offline RL has been tried and where significant gaps in efficiency has been noticed would further strengthen the authors' argument.\n\n+ While it's true that continuous action spaces and stochasticity in environments may make planning horizons longer, the authors would do well to mention the role of robust control or robust RL in mitigating some of these challenges whilst highlighting how their proposed method complements these existing methods;\n - In this sentiment, I am interested to see a motivation for why the authors decided to embark on the temporal abstractions proposal; \n\n + There is a significant body of work in continuous space RL in the realm of the viscosity solutions to Hamilton-Jacobi equations. The authors should discuss some of these in their intro/related work. Example texts:\n\n  - A Study of Reinforcement Learning in the Continuous Case by the Means of viscosity Solutions: https://www.ri.cmu.edu/pub_files/pub1/munos_remi_1999_3/munos_remi_1999_3.pdf\n\n  - Solving high-dimensional partial differential equations using deep learning: https://www.pnas.org/doi/pdf/10.1073/pnas.1718942115\n\n+ Please define what sg signifies in equation (2)!", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)