Value Memory Graph: A Graph-Structured World Model for Offline Reinforcement Learning
Authors: Deyao Zhu
Li Erran Li
Mohamed Elhoseiny
Venue: ICLR 2023 poster (2023)
Keywords: Offline Reinforcement Learning
Model Based Reinforcement learning

Abstract: Reinforcement Learning (RL) methods are typically applied directly in environments to learn policies. In some complex environments with continuous state-action spaces, sparse rewards, and/or long temporal horizons, learning a good policy in the original environments can be difficult. Focusing on the offline RL setting, we aim to build a simple and discrete world model that abstracts the original environment. RL methods are applied to our world model instead of the environment data for simplified policy learning. Our world model, dubbed Value Memory Graph (VMG), is designed as a directed-graph-based Markov decision process (MDP) of which vertices and directed edges represent graph states and graph actions, separately. As state-action spaces of VMG are finite and relatively small compared to the original environment, we can directly apply the value iteration algorithm on VMG to estimate graph state values and figure out the best graph actions. VMG is trained from and built on the offline RL dataset. Together with an action translator that converts the abstract graph actions in VMG to real actions in the original environment, VMG controls agents to maximize episode returns. Our experiments on the D4RL benchmark show that VMG can outperform state-of-the-art offline RL methods in several tasks, especially when environments have sparse rewards and long temporal horizons. Code is available at https://github.com/TsuTikgiau/ValueMemoryGraph

Published as a conference paper at ICLR 2023
VALUE MEMORY GRAPH: A GRAPH-STRUCTURED WORLD MODEL FOR OFFLINE REINFORCEMENT LEARNING
Deyao Zhu'
Erran Li
Mohamed Elhoseiny!
King Abdullah University of Science and Technology
2 AWS AI, Amazon and Columbia University deyao. zhu, mohamed. elhoseiny / (aaaattte...... edu. sa, erranlli@gmail com
ABSTRACT
Reinforcement Learning (RI.) methods are typically applied directly in environ- ments learn policies. In some complex environments with continuous state-action spaces, sparse rewards, and/or long temporal horizons, learning good policy in the original environments can be difficult. Focusing on the offline RL setting, we aim to build simple and discrete world model that abstracts the original environment. RL methods are applied our world model instead of the envi- ronment data for simplified policy learning. Our world model, dubbed Value Memory Graph (VMG), is designed as a directed-gret sereedd-gaa--bbasse Markov decision process (MDP) of which vertices and directed edges represent graph states and graph actions, separately. As state-action spaces VMG are finite and relatively small compared to the original environment, we can directly apply the value it- eration algorithm on VMG to estimate graph state values and figure out the best graph actions. VMG is trained from and built the offline RL dataset. Together with an action translator that converts the abstract graph actions in VMG to real actions the original environment, VMG controls agents maximize episode returns. Our experiments on the D4RL benchmark show that VMG can outperform state-of-the-as offline RL methods in several goal-oriented tasks, especially when environments have sparse rewards and long temporal horizons. Code is available at https::/qittuue www./swwoo////ccooott//// kg au Va ueMemo
INTRODUCTION
Humans usually good at simplifying difficult problems into easier by ignoring trivial details and focusing on important information for decision making. Typically, reinforcement learning (RL) methods are directly applied in the original environment to learn a policy. When have difficult environment like robotics or video games with long temporal horizons, sparse reward signals, or large and continuous state-action space, becomes more challenging for RL methods to reason the value of states or actions in the original environment to get a well-performing policy. Learning world model that simplifies the original complex environment into an easy version might lower the difficulty to learn policy and lead to better performance.
In offline reinforcement learning, algorithms can access a dataset consisting of pre-collected episodes to learn a policy without interacting with the environment. Usually, the offline dataset is used as replay buffer to train a policy an off-policy way with additional constraints to avoid distribution shift problems (Wu et al. 2019: Fujimoto et al. 2019; Kumar et al., 2019 Nair et al. 2020; Wang et al., 2020; Peng et al. 2019). As the episodes also contain the dynamics information of the original environment, it possible e to utilize such dataset directly learn abstraction of the environment in the offline RI. setting. To this end, we introduce Value Memory Graph (VMG). a graph-structured world model for offline reinforcement learning tasks. VMG is a Markov decision process (MDP) defined on a graph as an abstract of the original environment. Instead of directly applying RL methods the offline dataset collected in the original environment, learn and build VMG first and use
Work done outside of Amazon

---

Published as a conference paper at ICLR 2023
open microwave
move kettle
switch light
open cabinet
t=1-78
79-1128
t= 129-202
-03-2446
Start & End Subtask Path Past Path
Task List
Training
Go High-Value
3 Open Microwave Move Kettle
Switch Light
Offline Dataset
Evaluation
V Open Cabinet
Value Memory Graph
Successful Episode
Figure I: Demonstration of a successful episode where a robot trained in the dataset "kitchen patttalatta"""" accomplishes 4 subtasks in sequence guided by VMG. Vertex values shown via color shade. By searching graph actions that lead to the high-value future region (darker blue) calculated by value iteration the graph, VMG controls the robot maximize episode rewards and finish the task.
simplified substitute of the enviromment apply RL methods. VMG built by mapping offline episodes to directed chains in metric space trained via contrastive learning. Then, these chains are connected to graph via state merging. Vertices and directed edges of VMG are viewed as graph states and graph actions. Each vertex transition VMG has rewards defined from the original rewards the environment.
To control agents in environments, we first run the classical value iteration algorithm Puterrma 2014 once on VMG to calculate graph state values. This can done less than one second without training value neural network thanks to the discrete and relatively smaller state and action spaces in VMG. At each timestep. VMG is used to search for graph actions that can lead to high-value value future states. Graph actions are directed edges and cannot be directly executed in the original environment. With the help of an action translator trained supervised learning (e.g., Emmons et al. (202T) using the same offline dataset, the searched graph actions converted environment actions control the agent. An overview of our method shown in Fig
Our contribution can be summarized follows:
We present Value Memory Graph (VMG), a graph-structured world model offline rein- forcement learning setting. VMG represents the original environments graph-based MDP with relatively small and discrete action and state spaces.
We design a method to learn and build VMG on offline dataset via contrastive learning and state merging.
We introduce a VMG-based method to control agents by reasoning graph actions that lead high-value future states via value iteration and convert them environment actions via an action translator.
Experiments the D4RL benchmark show that VMG can outperform several state-of- -fie--ff-heeee---ee art offline RL methods on several goal-orienter tasks with sparse rewards and long temporal horizons.
2

---

Published as a conference paper at ICLR 2023
2
RELATED WORK
Offline Reinforcement Learning One crucial problem in offline RL is how to avoid out- the----he--tttt
training-distribution (OOD) actions and states that decrease the performance in test environments
(Fujimoto et al., 2019; Kumar et al., 2019 Levine et al. 2020)), Recent works like Wu et al. (2019);
Fujimoto et al. (2019); Kumar et al. 2019 Nair et al. (2020); Wang et al. (020); Peng et al. (2019)
directly penalize the mismatch between the trained policy and the behavior policy via explicit
density model or via divergence. Another methods like Kumar et al. (2000); Kostrikov et al. (202 (02laby
constrains the training via penalizing the Q function. Model-baset reinforcement learning methods like Yu et al. (2000) 2021): KKdambi et al. (2020) constrains the policy to the region of the world model that is close to the training data. Compared to previous methods, graph actions in VMG
always control agents to move to graph states, and all the graph states come from the training dataset. Therefore, agents stay close to states from the training distribution naturally.
Hierarchical Reinforcement Learning Hierarchical RL methods (e.g.. Savinov al. (2018); Nachum et al. 2 Eysenbach et al. (2019); Huang et al. (2019); Liu et al. (2020); Mandlekar et al. (2020) Yang al. (2020); Emmons et al. (2020); Zhang al. 2)21)) use hierarchical policies to control agents with high-level policy that generate commands like abstract actions or skills, and a low-level policy that converts them to concrete environment actions. Our method can be viewed hierarchical RL approach with VMG-based high-level policy and a low-level action translator. Compared to previous methods which learn high-level policies in environments, our high-level policy instead trained in VMG via value iteration without additional neural network learning.
Model-based Reinforcement Learning Recent research in model- reinforcement learning (MBRL) has shown a significant advantage (Ha & Schmidhuber 2018 Janner et al.) 2019 Hafner
al. 2019 (chrittwieser al., 2020 Ye et al. (2021) in sample efficiency over model-free rein- forcement learning. In most of the previous methods, world models are designed to approximate the original environment transition. In contrast, VMG abstracts the environment as a simple graph-based MDP. Therefore, we can apply RL methods directly to VMG for simple and fast policy learning. As we demonstrate later in our experiments, this facilitates reasoning and leads good performance in tasks with long temporal horizons and sparse rewards.
Graph from Experience Similar to VMG, methods like Hong et al. (022); Jiang et al. (002))
Shrestha et al. (2020); Marklund et al. (2020); Char et al. (2022) study the credit assignment problem on graph created from the experience. Hong et al. [2022) ; Jiang et al. (2022) are designed for discrete environments. [Shrestha et al. 2200); Marklund et al. (2020) considers the environments with finite actions and continuous state space by discretizing states state features via kNN, Char ctt al. (2022) introduces a stitch operator to create a graph directly by adding new transitions. It work with environments with low-dimensional continuous action spaces like Mountain Car Continuous (1 dimension) and Maze2D (2 dimensions). However, the stitch operator is hard to scale to ligh-dimmnsional action spaces. In contrast, VMG discretizes both state and action spaces and thus can work with continuous high-dimensional action spaces.
Representation Learning Contrastive learning methods learns a good representation by maximiz- ing the similarity between related data and minimizing the similarity of unrelated data (Oord et al. 2018 Chen et al. 2020 Radford et al. 2021) in the learned representation space. Bisimulation- based
methods like Zhang al. (2020) learn a representation with the help of bisimulation metrics (Ferns &
Precup] 2014 Ferns et al. 2011 Bertsekas & Tsitsiklis 1995) measuring the 'behavior similarity of
states W.r.t. future reward sequences given any input action sequences. In VMG, we use a contrastive
learning loss to learn a metric space encoding the similarity between states as L2 distance.
VALUE MEMORY GRAPH (VMG)
Our world model, Value Memory Graph (VMG), is a graph-structured Markov decision process constructed simplified version of the original environment with discrete and relatively smaller state-action spaces. RL methods can be applied on the VMG instead of the original environment to lower the difficulty of policy learning. To build VMG, we first learn metric space that measures the
3

---

Published as a conference paper at ICLR 2023
Afaa
Bnc,
f.
Erc,
f
Contrastive
Avayy
fe S
Afns
Dec,
Ene,
EnG.
(a) Contrastive learning
(b) Action reconstruction
Figure 2: The training pipeline of the state encoder Enc, and the action encoder Enc, to build the memory map. Enc, converts original states into points in the memory map. Enc, maps actions as transitions in the memory map.
reachability among the environment states. Then, graph is built in the metric space from the dataset as the backbone of our VMG. In the end. a Markov decision process is defined on the graph as an abstract representation of the environment.
3.1 VMG METRIC SPACE LEARNING
VMG is built in a metric space where the L2 distance represents whether one state can be reached from another state in few timesteps. The embedding in the metric space is based contrastive-learnis mechanism demonstrated in Fig 2a We have two neural networks: state encoder Enc, : that maps the original state to state feature in the metric space, and an action encoder Enc, fu, a + SS., that maps the original action a to a transition A fs,a in the metric space conditioned on the current state feature fs. Given transition triple (s, a, s'), ), we add the transition Afsa the state feature fa as the prediction of the next state feature for fo Afs,a- The prediction is encouraged to be close to the ground truth fs and away from other unrelated state features. Therefore, we use the following learning objective train Ene, and Enca:
Le = D'
N Σ max(m D2 Smnesss.
(1)
Here, :) denotes the L2 distance. 8neg.n denotes the n-th negative state. Given batch of transition triples (si, ai 8{) randomly sampled from the training set and a fixed margin distance m, all the other next states "'15zi the negative states for 8; and encourage f&; to be least m away from negative states the metric space. In addition, Dec, fe. Afs.a + à to reconstruct the action from the transition Afs.a conditioned on the state feature fa as shown in 2ig This conditioned auto-encoder structure encourages the transition ASs, to be meaningful representation of the action. Besides, we penalize the length of the transition when larger than the margin m to encourage adjacent states close in the metric space. Therefore, have the additional action loss L. shown below.
use action decoder
La = D² a) max( [Afaalll fff.llllmmm. m,
Linetrce the total training loss for metric learning, is the sum of the contrastive and action losses.
Linerric Le Lo
3.2 CONSTRUCT THE GRAPH VMG
To construct the graph in VMG, we first map all the episodes in the training data to the metric space as directed chains. Then, these episode chains are combined into graph with reduced number of state features. This is done by merging similar state features into one vertex based on the distance in the metric space. The overall algorithm are visualized in Fig 3a and can be found in Appx B. Given distance threshold 7m. vertex set V. and checking state 8;, check whether the minimal distance in the metric space from the existing vertices to the checking state 8; is smaller than 7m- not or if the vertex set is empty, process is repeated over the whole dataset. After the vertex set 1/ is constructed, each state s; can be classified into vertex by of which the distance in the metric space is smaller than 7mm In the training
set the checking state s;
vertex vy and add it to V. This
4

---

Published as a conference paper at ICLR 2023
%
(o(%) 2 augtry ) angtr a
(a) Merging vertices to construct the graph
(b) Define the reward of graph actions
Figure 3: Create a graph and define rewards in VMG. In Fig 3a three episodes are mapped as three chains in the metric space colored differently. merge nodes that are close to each other together and combine these chains into directed graph. In 3b the graph reward Ra Ra(un Vja) of the action from the green vertex V/, to the blue vertex 2/2 is defined as the average over rewards in the original episodes.
set, each state transition (st, az, dg,55) represents directed connection from s; to s/.. Therefore, we create the graph directed edges from the original transitions, For any two different vertices U/a Uja in V. if there exist transition (s,, 0z. s/) where s; and sy can be classified into v,, and U/ respectively, -ajo from Uj to V/ja
we add a directed edge j--
3.3 DEFINE A GRAPH-BARE BASED MDP
defined on the graph. Sa, Pa, Rg denotes the state set, the action set, the state transition probability, and the re- ward of this new graph MDP, respectively. Based on the graph, each vertex is viewed as a graph state. Besides, view each directed connection Gjiztz starting from vertex Ej an available graph action in Unj- Therefore, the graph state set Sg equals the graph vertex set and the graph action set is the graph edge set &. For the graph state transition probability Pg from fn U/g we define it as 1 if the corresponding edge exists in & otherwise O. Therefore,
VMG a Markov decision process (MDP)
Pa Ra(jaalun Way <n- - {
if en E
otherwise
We define the graph reward of each possible state transition ej-+jz the average over the original rewards from en to 2jz in the training set D, plus "internal rewards". The internal reward comes from the original transitions that are inside vy, djz after state merging. An example of graph reward definition is visualized in 3b Concretely,
Rji+jz = avg(r) ra Vss classified to U/ s/ classified to U/ja (8i, as, D)
(5)
{ Rソ Rii+ja - if ef E
Ru 1/,)
(6)
Not defined
otherwise
Note that the rewards of graph transitions outside of & are not defined, as these transitions will not happen according to Eq For internal rewards where both the source 8; the target &f of the original transition (si, s are classified to the same vertex, we split the reward into two and allocate them to both incoming and outgoing edges, respectively. This is shown as 1 R Rjitj and R -+j in 6
Now we have well-defined MDP on the graph. This MDP serves as our world model VMG.
3.4 How TO USE VMG
VMG, together with action translator, generate environment actions that control agents to maximize episode returns. We first run the classical RL method value iteration (Puterman, 2014) VMG to compute the value ((0z) of each graph state Uy. This be done in second without learning an additional neural-network- state-action spaces.
value function due to VMG's finite and discrete
To guide the agent, VMG provides a graph action that leads to higu-value graph states in the future at each time step. Due the distribution shift between the offline dataset and the environment, there
5

---

Published as a conference paper at ICLR 2023
can be gaps between VMG and the environment. Therefore, the optimal graph action calculated directly by value iteration on VMG might not be optimal in the environment. We notice that instead of greedily selecting the graph actions with the highest next state values, searching for a good future state after multiple steps first and planning a path to it can give us a more reliable performance. Given the current environment state Sc, first find the closest graph state V/ on VMG. Starting from Ve, we search for N, future steps to find the future graph state v² with the best value. Then, we plan a shortest path = [Vey 1o+1,... v²] from U. U via Dijkstra (Dijkstra al. 1959) on the graph. We select the Neg-th graph state NctN,, and make an edge Cc-c+NN, as the searched graph action. The graph action Eg, is converted the environment action de via an action translator: a, = Tran(s (Se, Vo..N.g). The pseudo algorithm can be found in Appx B.2
The action translator (ran(s, s') reasons the executed environment action given the current state 8 and a state s' in the near future. Tran(s, s') is trained purely in the offline dataset via supervised learning and separately from the training of VMG. detail, given an episode from the training set and time step t, we first randomly sample step + & from the future K steps. k ~ Uni form K). Then, (r, ,, s') is trained to regress the action a, at step given the state s and the future state 884K using L2 regression loss LTran D2 (t, St+4), a4) Note that when
1, p(a (st, S(+*) ) is determined purely by the environment dynamics and Tram (s. s') becomes an inverse dynamics model. As k increase, the influence of the behavior policy that collects the offline dataset on p(a, St, 8t+&) will increase. Therefore, the sample range K should be small to reflect the environment dynamics and reduce the influence of the behavior policy. In all of our experiments, K is to 10.
4 EXPERIMENTS
4.1 PERFORMANCE ON OFFLINE RL BENCHMARKS
Test Benchmark We evaluate VMG on the widely used offline reinforcement learning benchmark D4RL (Fu al. 2020 In detail, we test VMG on three domains: Kitchen, AntMaze, and Adorit. In Kitchen, a robot arm in virtual kitchen needs to finish four subtasks in an episode. The robot receives a sparse reward after finishing each subtask. D4RL provides three different datasets in Kitchen: kitchen-complete, kitchen-pee fitchen-partial, and kitchen-mixed. In AntMaze, a robot ant needs to go through a maze and reaches a target location. The robot only receives sparse reward when reaches the target. D4RL provides three mazes of different sizes. Each of them contains two datasets. In Adroit, policies control a robot hand to finish tasks like rotating a pen or opening a door with dense rewards. For evaluation, D4RL normalizes all the performance of different tasks to range of 0- 0-100, where 100 represents the performance of an "expert" policy. More benchmark details can found in D4RL (Fu et al. 2020) and Appx
Table 1: Experimental results on domains Kitchen, AntMaze, and Adroit from D4RL benchmark. VMG outperforms baselines in Kitchen and AntMaze where only sparse rewards are provided and achieves comparable performance Adroit. Results and the standard deviation are calculated over three trained models.
Dataset
BC
BRAC-p
BEAR
DT
AWAC
CQL
IQL
VMG
kitchen-comee compleee
65.0
0.0
43.8
62.5
73.0
6.7
kitchen-partial
38.0
0.0
0.0
49.8
46.3
68.8
11.9
kitchen-mixed
51.5
51.0
51.0
50.6
kitchen-total
154.3
0.0
0.0
144.6
159.8
192.4
antmaze-umaze
34.6
59.2
36.7
74.0
87.5
93,7
antmaze-umaze diverse
45.6
53.0
49.3
84.0
62.2
94.0
2.0
antmaze medium play
0.0
0.0
61.2
71.2
82.7
3.1
coammne-medium-dfie diverse
0.0
0.7
53.7
70.0
84.3
2.1
antmazo-largo-pae
0.0
0.0
0.0
15.8
39.6
67.3 4: 3.2
uutmaze-large-diverr
1.0
74.3
animaae-tooial
100.2
112.2
107.7
303.6
378.0
496.3
pea-human
63.9
8.1
-L.0
37.5
71.5
70.7
5.2
pen-cloned
37
1.6
26.5
39.2
37.3
58.2
1.6
hammer hurnan
1.2
0.3
0.3
4.4
1.4
4.1
1.2
hammer-eloned
0.6
2.1
2.1
2.2
door-human
2
-0.3
-0.3
9.9
4.3
1.5
0.5
door-cloned
-0.1
2.2
adroit-toial
104.7
9.9
25.7
93.5
118.2
138.9
kiichen-anttazz--al
359.4
541.7
656.0
827.6

---

Published as a conference paper at ICLR 2023
Baselines We mainly compare method with two state-of-the-as methods CQL (Kumar et al. 2020) and IQL (Kostrikov et al., 2021b) in all the the Q function to alleviate the OOD action issue in the
datasets. Both CQL and IQL
based on Q-learning with constraints offline setting. In addition, we also report the performance of BRAC-p (Wu et al. 2019), BEAR (Kumar et al. 2019), DT (Chen al. (22TT, and AWAC (Nair al. 2020) in the datasets they used. Performance of behavior cloning (BC) is from (Kostrikov et al. 202
Hyperparameters In all the experiments, the dimension of metric space is set to 10. The margin m in and is 1. The distance threshold "/ is set to 0.5, 0.8, and 0.3 in Kitchen, AntMaze, and Adorit, separately. Hyperparameters are selected from 12 configurations. We use Adam optimizer (Kingma & Ba 2014) with a learning rate 10-3, train the model for 800 epochs with batch size 100, and select the best-performing checkpoint. More details about hyperparameters and experiment
settings are in Appx D
Performance Experimental results are shown in T VMG's scores averaged over three individually trained models and over 100 individually evaluated episodes in the environment. In general, VMG outperforms baseline methods in Kitchen and AntMaze and shows competitive performance in Adroit. Note that a good reasoning ability in Kitchen and AntMaze domains is crucial as the rewards in both domains are sparse, and the agent needs to plan over a long time before getting reward signals. In AntMaze, baseline methods perform relatively well in the smallest maze 'umaze', which requires less than 200 steps solve. In the maze 'large' where episodes can be longer than 600 steps, the performance of baseline methods drops dramatically. VMG keeps a reasonable score all three mazes, which suggests that simplifying environments to graph-structured MDP helps RL methods better reason over long horizon in the original environment. Adroit is the most challenging domain for all the methods in D4RL with a high high-dimensional action space. VMG still shows competitive performance in Adroit compared baselines. Experiments show that learning policy directly in VMG helps agents perform well, especially in environments with sparse rewards and long temporal horizons.
4.2 UNDERSTANDING VALUE MEMORY GRAPH
To analyze whether VMG can understand and represent the structure of the task space correctly, we visualize an environment, the corresponding VMG, and their relationship Fig We study the task "antmaze-large-dverrs
shown in Fig 4a as the state space of navigation tasks is easier to visualize and understand. The target location where the agent receives a positive reward is denoted by a red circle. A successful trajectory is plotted as the green path. To visualize VMG, all the state features f. are reduced to two dimensions via UMAP Mcclnnes et al. 2018) and used as the coordinate to plot corresponding vertices shown in Fig 4b The graph state values are denoted by color shades. Vertices with darker blue have higher values. As shown in Fig 4b VMG allocates high values to vertices that close the target location and low values to far away vertices. Besides, the topology of VMG is similar to the maze. This is further visualized in Fig 4c where graph vertices are mapped the corresponding maze locations show their relationship. Our analysis suggests that VMG can
a G
GH
(a) AntMaze Large
(b) VMG
(c) VMG mapped to the maze
Figure 4: An example of VMG learned from the dataset 'antmaze-large-divens Fig 4a shows the environment with the target location highlighted by red circle. VMG visualized via UMAP in Fig[4b Graph state values are represented by color shades with higher values in darker blue. Graph states that are close to the target have high values calculated by value iteration. In Fi graph states mapped to the corresponding maze locations to show the relationship.
7

---

Published as a conference paper at ICLR 2023
T=15
T=30
T=60
85.0
65k
65.0
40%c
45.0
25.0
0.3
0.4
0.5
0.0
Figure 6: Influence of Tim in "kitchen-partial" in performance and VMG size.
R9: FIST
Figure 5: VMG and a successful trial in the task "pen- "pen-human" The blue pen is rotated the same orientation the green one.
learn meaningful representation of the task. Another VMG visualization in the more complicated task "pen-human" shown in Fig and and more visualizations can be found in Appx
4.3 REUSABILITY OF VMG WITH NEW REWARD FUNCTIONS
In offline RL, policies are trained mas- ter skills that can maximize accumulated returns via an offline dataset. When the dataset contains other skills that don't 't lead to high rewards, these skills will be simply ignored. We name them ignored skills. We can retrain a new policy to master ignored skills by redefining new reward functions correspondingly. However, rerunning RL methods with new reward functions is cumbersome the original complex environment, scratch. In contrast, rerunning value iteration VMG with new reward functions takes less than one second without retraining any neural networks. Note that the learning of VMG and the action translator is reward-free. Therefore, we don't need to retrain VMG but recalculate graph rewards using Eq 6 with new reward functions.
Table 2: VMG success rate of ignored skills. Agents can perform these skills by rerunning value iteration with the new reward function in trained VMG.
Value Iteration
Bottom Burner
Top Burner
Hinge Cabinet
Orig. Reward
0.7
3.7
0.0
New Roward
69.7
7,3
need to retrain the policy network and Q/value networks from
We design an experiment in the dataset "kitchen Sitchenn-aatial" to verify the reusability of VMG with new reward functions. In this dataset, the robot only receives rewards in the following four subtasks: open microwave, move a kettle, turn a light, and open a slide cabinet. Besides, there are training episodes containing ignored skills like turning on burner or opening a hinged cabinet. We first train model in the original dataset. Then, we define a new reward function, where only ignored skills have positive rewards and relabel training episodes correspondingly. After that, we recalculate graph rewards using 6 rerun value iteration on VMG, and test our agent. Experimental results in Tab show that agents can perform ignored skills after rerunning value iteration in the original VMG with recalculated graph rewards without retraining any neural networks.
4.4 ABLATION STUDY
Distance Threshold The distance threshold 7m directly controls the "radius" of vertices and affects the size the graph. We demonstrate how 7m affects the performance in the task kkitchenppatial in Fig The dataset size of "kitchen-partia" is 137k. A larger Ym can reduce the number of vertices but hurts the performance due to information loss. More results can be found Appx
Graph Reward Here study how will different designs of the graph reward Ro Uy.) affect the final performance. addition to the original version defined in Eq 5 and Eq that averages the environment rewards, try maximization and summation and denote them RG RG.maz and RG.sum, separately. Besides, we also study the effectiveness of the the internal reward through the following three variants of 9 RG.rm R2, RG.rm., R + R RG.mm. = Ry,.. + RR if ejj E E. Experimental results shown in suggest that the original design of the graph rewards represent the environment well and leads the best performance.
8

---

Published as a conference paper at ICLR 2023
Table 3: Ablation study of graph reward design in VMG. The original design gives the best performance.
Importance of Contrastive Loss The contrastive loss is the key training objective to learning meaningful metric space. The contrastive loss pushes states that can be reached in few steps be close each other and pushes away other states. To ver- ify this, we train variant of VMG without the contrastive loss and show the results in Tabl The variant without the contrastive loss does not work at all (0 scores) in the "kitchen-partial" and 'aatmaze-medium-plee tasks, and the performance 'pen-human" significantly decreases from 70.7 to 41.2. Results indicate the importance of contrastive loss in learning robust metric space.
Variants
kitchen-pariall
antmaze-mediun-ppay
pcu-human
Re,max
31,5
48.3
50.8
RG.cm
50.3
56.3
66.4
55.0
55.0
65.2
54.6
72.2
RG.
67.0
78.5
65.9
Ro (Orig.)
68.8
82.7
70.7
Effectiveness of Action Decoder The ac- tion decoder is trained to reconstruct the original action from the transition in the metric space conditioned on the state fea- ture shown in Fig 2b In this way, the train- ing of the action decoder encourages transi-
Table 4: Ablation study of contrastive loss, action de- coder, and Dijkstra search.
Model
kitchen-partial
antmaze-medium play pen-human
VMG
68.8
82.7
70.7
cuntrastive
0.0
action decoder
15.4
66.3
68.5
tions in the metric space to better represent
step search
39.8
actions and leads to a better metric space. To show the effectiveness of the action decoder, train VMG variant without the action decoder and show the results in Tab The performance without the action decoder drops in all three tested tasks, especially in "kitchen-pariaa" (from 68.8 to 15.4). The results verify our design choice.
Multi-Step Search In Tab list the performance of our method without Multiple-Ste Search. Compared to the original version, we observe performance drop in Sitchen-partial and "antmaze- medium-play' and similar performance in 'pen-human', which suggests that instead of greedily searching one step in the value interaction results, searching multiple steps first to find high value state in the long future and then plan a path via Dijkstra can help agents perform better. We think the advantage might caused by the gap between VMG and the environment. An optimal path on VMG searched directly by value iteration may not be still optimal in the environment. At the same time, shorter path from Dijkstra helps reduce cumulative errors and uncertainty, and thus increases the reliability of the policy.
Limitations As an attempt to apply graph-structured world models offline reinforcement learning, VMG still has some limitations. For example, VMG doesn't 't learn to generate new edges in the graph but only creates edges from existing transitions in the dataset. This might limitation when there is not enough data provided. In addition, VMG is designed in offline setting. Moving to the online setting requires further designs for environment exploring and dynamic graph expansion, which can be interesting future work. Besides, the action translator trained via conditioned behavior cloning. This may lead to suboptimal results in tasks with important low-level dynamics like gym locomotion (See Appx H. Training the action translator by offline RL methods may alleviate this issue.
CONCLUSION
We present Value Memory Graph (VMG), a graph-structured world model in offline reinforcement learning. VMG is Markov decision process defined on directed graph trained from the offline dataset as an abstract version of the environment. As VMG is smaller and discrete substitute for the original environment, RL methods like value iteration can be applied on VMG instead of the original environment to lower the difficulty of policy learning. Experiments show that VMG can outperform baselines in many goal-orienten long temporal horizons in the widely used offline RL benchmark D4RL. We believe VMG shows promising direction to improve RL performance via abstracting the original environment and hope can encourage more future works.
tasks, especially when the environments have sparse rewards and

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
We would like thank Ahmed Hefny and Vaneet Aggarwal for their helpful feedback and discussions on this work.
REFERENCES
Dimitri P Bertsekas and John N Tsitsiklis. Neuro-dynamic programming: an overview. In Proceedings of 1995 34th IEEE conference on decision and control, volume 1, pp. 560-564. IEEE, 1995.
lan Char, Viraj Mehta, Adam Villaflor, John M Dolan, and Jeff Schneider. Bats: Best action trajectory stitching, arXiv preprint (rXi::2204.12026, 2022.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34, 2021.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Edsger W Dijkstra et al. A note on two problems in connexion with graphs. Numerische mathematik, 1(1):269-271, 1959.
Scott Emmons, Ajay Jain, Misha Laskin, Thanard Kurutach, Pieter Abbeel, and Deepak Pathak. Sparse graphical memory for robust planning. Advances in Neural Information Processing Systems, 33:5251-5262, 2020.
Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning? arXiv preprint arXiv:2112.10755 2021.
Ben Eysenbach, Russ R Salakhutdinov, and Sergey Levine. Search on the replay buffer: Bridging planning and reinforcement learning. Advances Neural Information Processing Systems, 32, 2019.
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. SIAM Journal on Computing, 1666:1161-11714, 2011.
Norman Ferns and Doina Precup. Bisimulation metrics are optimal value functions. In UAI, pp. 210-219, 2014.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. International Conference on Machine Learning, pp. 2052-2062. PMLR, 2019.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803. 2018.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 1911:199.. 2019.
Zhang-Wei Hong, Tao Chen, Yen-Chen Lin, Joni Pajarinen, and Pulkit Agrawal. Topological experience replay. arXiv preprint arXiv:2203. 15845, 2022.
Zhiao Huang, Fangchen Liu, and Hao Su. Mapping state space using landmarks for universal goal reaching. Advances in Neural Information Processing Systems, 32, 2019.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. Advances in Neural Information Processing Systems, 32, 2019.
Zhengyao Jiang, Tianjun Zhang, Robert Kirk, Tim Rocktäschel, and Edward Grefenstette. Graph backup: Data efficient backup exploiting markovian transitions. arXiv preprint a5Xi: 2022.
10

---

Published as a conference paper at ICLR 2023
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data, 7(3):535- 02755-57,222 2019.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. Advances in neural information processing systems, 33: 21810-21823, 2020.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Hya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. In International Conference on Machine Learning, pp. 5774-5783. PMLR, 2021a.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06169, 2021b.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances Neural Information Processing Systems, 33: 179-1199,, 2020.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.0066 2020.
Kara Liu, Thanard Kurutach, Christine Tung, Pieter Abbeel, and Aviv Tamar. Hallucinative topologi- cal memory for zero-shot visual planning. In International Conference on Machine Learning, pp. 6259-6270. PMLR, 2020.
Stuart Lloyd. Least squares quantization in pem. IEEE transactions on information theory, 28(2): 129-137, 1982.
Ajay Mandlekar, Danfei Xu, Roberto Martín- Silvio Savarese, and Li Fei-Fei. Learning to generalize across long- hong-horizon tasks from human demonstrations. arXiv preprint arXiv:2003.06085, 2020.
Henrik Marklund, Suraj Nair, and Chelsea Finn. Exact (then approximate) dynamic programming for deep reinforcement learning. In Bian and Invariances Workshop, ICML, 2020.
Leland McInnes. John Healy, and James Melville. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint
2018.
Ofrr Nachum, Shixiang Shane Gu, Honglak Lee, and Sergey Levine. Data-efficient hierarchical reinforcement learning. Advances in neural information processing systems, 31, 2018.
Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09359, 2020.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 1807.03748 2018.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weighter regression: Simple and scalable off- off-policy reinforcement learning. arXiv preprint 1911v:1910.00177,, 2019.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, pp. PMLR, 2021.
11

---

Published as a conference paper at ICLR 2023
Nikolay Savinov, Alexey Dosovitskiy, and Vladlen Koltun. Semi- Semi-parametric topological memory for navigation. arXiv preprint arXiv:1803.00653. 2018.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):004-69, 2020.
Aayam Kumar Shrestha, Stefan Lee, Prasad Tadepalli, and Alan Fern. Deepaveragers: Offline reinforcement learning by solving derived non- paraparramtiie mdps. International Conference on Learning Representations, 2020.
Michita Imai Takuma Seno. d3rlpy: An offline deep reinforcement library. In NeurlPS 2021 Offline Reinforcement Learning Workshop, December 2021.
Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh S Merel, Jost Tobias Springenberg, Scott E Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, al. Critic regularized regression. Advances in Neural Information Processing Systems, 33:7768-7778, 2020.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint srXiv:/911. 1361, 2019.
Ge Yang, Amy Zhang, Ari Morcos, Joelle Pineau, Pieter Abbeel, and Roberto Calandra. Plan2vec: Unsupervised representation learning by latent plans. In Learning for Dynamics and Control, pp. 935-946. PMLR, 2020.
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in Neural Information Processing Systems, 34, 2021.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 33: 14:19-1144422 2020.
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in Neural Information Processing Systems, 34, 2021.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006. 10742, 2020.
Lunjun Zhang, Ge Yang, and Bradly C Stadic. World model graph: Learning latent landmarks for planning. International Conference on Machine Learning, pp. 1261 1-12620. PMLR, 2021.
Tian Zhang, Raghu Ramakrishnan, and Miron Livny. Birch: an efficient data clustering method for very large databases. ACM sigmod record, 25(2):103-11 114, 1996.
12

---

Published as a conference paper at ICLR 2023
CONTENTS
Introduction
2
Related Work
3
Value Memory Graph (VMG)
3.1
VMG Metric Space Learning
3.2
Construct the Graph in VMG
3.3
Define Graph-Based MDP
How to Use VMG
Experiments
4.1
Performance on Offline RI. Benchmarks
6
4.2
Understanding Value Memory Graph
4.3
Reusability of VMG With New Reward Functions
8
Ablation Study
Conclusion
Environment details
15
Algorithms
15
B.i Graph Construction
15
B.2 Policy Execution
15
Architecture of Neural Networks
15
Experiment Settings and Hyperparameters
16
Ablation Studies
16
Distance Threshold
16
State Merging Method
17
Influence of m
17
E,4
Influence of Discount Factor
18
Influence of the Dimension of the Metric Space
18
Influence of K
18
Training Curve
18
Environment/Graph Transition Ratio
19
Experiments in Gym Locomotion Tasks
19
13

---

Published as a conference paper at ICLR 2023
Future Works
20
Visualization of VMG
20
14

---

Published as a conference paper at ICLR 2023
A ENVIRONMENT DETAILS
The datasets in D4RL (Fu al. 2020 is under CC BY license and the related code is under Apache 2.0 License. We use the latest version of the datasets (v1/v0/vl for AntMaze, Kithcen, Adroit, separately). Different versions of datasets contain exactly the same training transitions. The newer version fixes some bugs in the meta data information like the wrong termination steps. Performance is measured by returns normalized to the range between 0 and 100 defined by the D4RL benchmark [9]. score- rondom score A score of 100 corresponds to the raadoo
In detail, normalized score 100 x ezpert
average returns of domain- expert. For AntMaze, and Kitchen, an estimate of the maximum score possible is used as the expert score. For Adroit, this estimated from a policy trained with behavioral cloning human-demonstration and online fine-tuned with RL in the environment. For more details about the datasets please refer to D4RL (Fu et al. 2020).
B ALGORITHMS
B. GRAPH CONSTRUCTION
The detailed algorithm of graph construction is shown in Alg
Algorithm 1:
Graph Construction
Input: : Training Set D = [(*;, az ri, s() i = 1, 1.2,..,,N) Empty vertices set = (}.
Current vertex index J 1, Distance threshold 7/. Empty edges set & {}
1 for (8i, 0i, st) in do
Js. =Enncsss))
Compute the distance dı between fs. and fa, for every fe, in V min in V} > Ym or J = then
4
VJ & Sis fos + fs.
V.append( f,,)) J J J +1
end
end
10 for
(8i, Og Ti, x in D do
Find Uj Uja that s; and si are classified to in V, respectively
12
if Uj # Djg and the connection ej -á2 & & then
E.append(e)_-ja)
end
is end
B.2 POLICY EXECUTION
The detailed algorithm of policy execution is shown in Alg
Details of Dijkstra When we use Dijkstra in Sec 3.4 to plan path P from u. to v", we define weights each edge make sure both short and high-rewarded. The weights used plan the path P are based on rewards. For each edge ej-+ we define the edge weight W,,+ja -j as the gap between the maximal graph reward and the edge reward and denote the weight set W. Wi + max Ro Uj. (Ve E E Ro( )..).
C ARCHITECTURE OF NEURAL NETWORKS
For all the networks including the state encoder Enc,, the action encoder Enca, the action decoder Dec,, and the action translator Tran(s, s'), we use 3-layer MLP with hidden size 256 and ReLU activation functions.
15

---

Published as a conference paper at ICLR 2023
Algorithm 2: Policy Execution
Input: Current state Se, State encoder Enc,, Action translator Tran, Vertex and edge sets VMG (V,E), Vertices value V, Edge weight W
Ss. (se)
Ve arg (VV f3)
3 Search future horizon of N, steps starting from u, and select the best value vertex o"
4 Compute the weighted shortest path P from de to o" via Dijkstra. P = [Ve, Pe+1s v*]
5 Og Tran(se: Ve+N.g)
D EXPERIMENT SETTINGS AND HYPERPARAMETERS
Our model is trained in a single RIX Titan GPU in about 1.5 hours. For inference, building the graph from clustering takes about 0.5-2 minutes before the evaluation. After that, it takes about 0.5-10 minutes to evaluate 100 episodes. We implement VMG on the top of the offline RL python package d3rlpy (Takuma Seno 2021) with MIT license.
the experiments, We use Adam optimizer (Kingma & Ba 2014) with learning rate 10 Batch size is 100. Each model is trained for 800 epochs. We save models per 50 epochs and report the performance of the best one evaluated in the environment from the checkpoints saved from the 500th to the 800th epochs. The remaining hyperparameter settings can be found in Tab N₂ = oo means we search the future steps till the end of the graph. For the domain Kitchen, the hyperparameters For Adroit, hyperparameters are tuned individually. We
tuned in "kitchen-partial" For
AntMaze is mantaaze-umaze-diersse use the environment tune the hyperparameters. We searched four hyperparameters in our main experiments: Tm in [0.3, 0.5, 0.8, 1.0, 1.2], reward discount in [0.8, 0.95], N in [I, 2, 31, N, in [12, searched one by one, in total 12 configurations. For hyperparameters like batch size or learning rate, we follow the default one in the RL library d3rlpy. The dimension of the metric space is set 10 in all the experiments. Tuning the hyperparameters offline is an ongoing and important research topic in offline RI, and we left it for future work.
o0]. Hyperparameters
Table 5: Detailed Hyperparameter Setting
Dataset
m
K
"/m
discount
N.,
N,
kitchen-complete
10
0.5
0.95
kitchen-parte
10
0.5
0.95
kitchen-mixet
10
0.5
0.95
antmaze-umaze
10
0.8
0.8
antmaze-umaz-ddiven
10
0.8
0.8
antmaze-medium-
10
0.8
0.8
untmuze-medium--dvenns
10
0.8
0.8
antmaze-large-play
10
0.8
0.8
antmaze-large-diverse
10
0.8
0.8
pen-human
10
pen-cloned
10
0.8
hammer-human
10
0.8
hammer-cloned
10
0.8
door-human -oor-huann
10
0.8
door-cloned
10
0.3
0.8
ABLATION STUDIES
E.I DISTANCE THRESHOLD
More experimental results of the distance threshold 7m in the tasks "antmaze Luntmaze-medium-aae medium-play" and "pen-cloned" can be found in Fig Results suggest that the model is not so sensitive to 7 if it not too large.
16

---

Published as a conference paper at ICLR 2023
80.0
8k
70.0
80k
80.0
6k
Number of Vertices
65.0
of Vertices
Parformance
Performance
40k
70.0
60.0
20k
60.0
55.0
50.0
50.0
0K
0.6
0.7
0.8
0.9
0.3 0.4
0.5
0.6 0.7 0.8 0.9
Parfermane
Namber Wrveiiee
Dateraace df
aa antmaze-medium--plg
(b) pen-cloned
Figure 7: More results of the influence of 7m in performance and VMG size
E.2 STATE MERGING METHOD
Vertices in VMG are merged from the original states based on a distance threshold 7m as described in Sec 3.2 It is also possible to use other clustering methods to merge states. However, the dataset sizes in some tasks can be up to I million. Many advanced clustering methods (like BIRCH (Zhang) et al. 1996)) are slow in this case (up to hours for BIRCH on machines with Intel Xeon Gold 6242). Therefore, we compare with the classical K-means (Lloyd 1982) implemented
Faiss (Johnson ct al. 2019) library with the GPU support in both the AntMaze domain and the Kitchen domains. Faiss-based K-means can be finished up to 20 seconds in our setting. Our merging method takes up to minute. Experimental results are shown in Tabl VMG created by our merging method performs better than the one created by K-means. The vertices of our method can be viewed hyperspheres in the metric space with the same radius Tmm In contrast, K-means cannot directly specify the size of each cluster, which can result in vertices with different "volumes" in the metric space. This might lead to undesired distortion in the graph and reduce the performance. As K-means doesn't have parameter to control the size the clusters directly, we have to search for the best number of clusters for every dataset. The number of clusters used in K-means is shown in Tab
Table 6: Ablation study of different state merging methods. Our original design gives us better performance.
AntMare
Kitchen
Model
umnaze
umuzze umuze-diverse diverse
medium-ppaa
mnedium-diverse
lage--play
large-diverse
complete
partial
mixed
VMG with K-meuns
88.7 7
79.7
81.2
77.0
72.3
76.3
61.1
18.3
43.6
Table 7: Number of clusters used in K- K-means
AntMaze
Kitchen partial 25000
umnaze 6000
umazo-diverse 2000
medium-play 1000
diversse- 1000
largo-play 10000
largo-diverse 10000
complete 3000
mixed 25000
E.3 INFLUENCE OF m
The value of the margin m Eq and Eq implicitly defines the minimal distance of negative state pairs in the learned metric space. To study the influence of m the performance, here we set m to 0.5, 1, and in the datasets autmaze-medium the results in Tab In the antmaze experiment, performance becomes better with larger m. But in the task pen-human, a smaller m gives us better results. In kitchen-partial, m=1 shows the best performance. Experimental results suggest that m=l is reasonable value for the tasks we evaluate on. And if we tune m separately, it is possible to improve the performance.
sitchen-partial and pen-human and show
17

---

Published as a conference paper at ICLR 2023
kitchen purtial
17.0
54.5
63.8
pes- pei-hnman
75.1
70.7
63.7
Table Influence of the margin m.
E.4 INFLUENCE OF DISCOUNT FACTOR
Here we study how different discount factor values will affect the performance of VMG. We set the discount factor the values 0.8, 0.95, and 0.99. Experiments in Tho that 0.99 leads to better performance in pen-human and comparable performance in kitchen-partial In antmaze- sedmaaze-edduumplle 0.99 performs worse than 0.8 and 0.95, which suggest that small discount factor in antmaze might help reduce cumulative errors.
discount factor
0.8
0.95
0.99
medium
62.7
76.3
75.0
kitchen
peo:
Table 9: Influence of the discount factor.
E.5 INFLUENCE OF THE DIMENSION OF THE METRIC SPACE
Here we study how different numbers of the metric space dimension will affect the performance of VMG. We set the metric space dimensions to 5, 10, and 20. Experiments in Tab show that models with latent space dimensions and 20 perform better than those with 5, which suggests that reasonable performance requires big enough dimensions of the latent space represent the states and actions better. Besides, space with 10 dimensions works better than 20 in kithcen-partial but worse than 20 in pen-human, this suggests the performance has space improve if we tune the dimensions individually in each task.
metric space dim
10
20
antruzz--sedium-lla
71.0
82.7
$2.0
kitchen punial
5.75
68.8
46.0
pen-human
Table 10: Influence of the metric space dimension.
E.6 INFLUENCE OF K
The hyperparameter K used in training the action translator in Sec 3.4 defines the range of the future states the action translator conditions on during training. To study the influence of K, here show experiments with K=5, 10, and 20 Tab TT We notice that K = doesn't work antmaze-medium and kitchen-partial, which suggests that K=5 is not big enough cover 2 steps in the graph transition. In addition, the experiments with = 20 show better results than K = 10 kitchen-par partial and pen-human. In maeinazz-meduum-peee K 10 performs the best. Experimental results suggest that big enough K helps the model perform better.
F TRAINING CURVE
Fig 8 shows the training curves of the contrastive loss L,, the action loss L,, and the anction translator loss LTran in tasks kitchen-partial, antmaze-mediun-play and pen-human.
18

---

Published as a conference paper at ICLR 2023
K
7.0
K2.7
74.0
kitchen purtial
0.3
53.8
76,8
pes- pei-hnman
80.3
70.7
83.5
Table 11: Influence of K.
6.8
533
330
6.4
645
0.0
Vnchenpatiall
inchenparlll
0.40
3448
0.03
0.00
3 0030
0.33
0:30
0415
030
5400
0.31
0.00
610.0
0.05
6.4
0.00
6.0
0.73
0.00
6:0
12008
Figure 8: Training curve of the contrastive loss Tec, action loss T. and action translator loss L-yaa-
G ENVIRONMENT/GRAPH TRANSITION RATIO
VMG abstracts the original continuous environment into finite and relatively small graph. An one-step transition in the graph corresponds to multiple steps in the environment. Here we compute the average numbers of environment transitions per graph transition in our main experiments and list the results in Tab 12
Table 12: Average number of environment transitions per graph transition.
AmMase
Kitchen
Adroit
medium
large
play
diverse 15.1
play
complete 1.1
partial
human 1.8
cloned
human
human
1.0
H EXPERIMENTS IN GYM LOCOMOTION TASKS
VMG introduced to help agents reason the long future better so as to improve their performance in complex environments with sparse rewards and large search space due to long temporal horizons and continuous state/action spaces. VMG may not help in gym locomotion tasks, as these tasks don't 't require agents to reason the long future and thus out of scope. Gym locomotion tasks provide rich and dense reward signals, and the motion patterns to learn in these tasks are periodic and short. Therefore, the problems VMG designed to solve are not issue here. Our performance in these tasks
19

---

Published as a conference paper at ICLR 2023
is expected to be close to behavior cloning, since the low-level component, the action translator, is trained via (conditioned) behavior cloning. The action translator is used to handle local dynamics that not modeled in VMG. Here Experimental results verify our assumption. Results and analysis suggest that an improved design and/or learning strategy of the action translator might help improve the performance. For example, training the action translator using conditioned offline RL methods instead of conditioned behavior cloning, However, this is orthogonal to our VMG framework contribution future reasoning, and we leave it for future work.
experiments in these tasks and show the results below.
Dataset
VMG
BC
CQL
IQL.
haffcheetah-medium
42.2
42.6
44.0
47.4
hopper-medium
49.4
52.9
58.5
66.3
walker2d-medium
70.4
75.3
72.5
78.3
halfchheeal--meiiuun-rpaaa
38.2
36.6
45.5
44.2
hopper-meddium-relly
15.2
18.1
95.0
94.7
-ellkerdd-meiuum--- repllay
28.6
26.0
77.2
73.9
halfcheetaah-medig medium-exxpert
80.5
55.2
91.6
86.7
hopper-meding medium--eedi expertt
49.5
52.5
105.4
91.5
saaiker2d-mmdiiu--xxpeee
70.4
107.5
108.8
109.6
Table 13: Peformance of VMG in gym locomotion tasks. The performance of VMG expected to be
closed to behavior cloning.
I FUTURE WORKS
There several directions to improve VMG. Building hierarchical graphs model different levels of environment structures might help represent the environment better. For example, a robot needs to cook meal, might have high-level graph to represent abstract tasks like washing vegetables, cutting vegetables, etc. A low-level graph can be used to guide a goal-conditioned policy. This might improve the high-level planning of the tasks. Extending VMG into the online setting is also an important future step. In online reinforcement learning, data with new information is collected throughout the training stage. Therefore, the graph needs to have mechanism continually expand and include the new information. Besides, exploration is a crucial component in online reinforcement learning. If we model the uncertainty of the graph, VMG can be used guide the agent explore regions with high uncertainty to explore more effectively. Combined with Monte Carlo tree search on VMG might also help policy explore and exploit better.
VISUALIZATION OF VMG
More visualization VMG in different tasks are demonstrated in Fig denoted as green path on the graph with "+" sign at the end.
episode
20

---

Published as a conference paper at ICLR 2023
(a) kitchen-complete
(b) VMG of kitchen-complete
(c) fitchen-particg
(d) VMG of sitchen-partial
(c) kitchen-mixed
(I) VMG of kitchen-mixed
(h) VMG of antmaze-umaze
Figure 9: Visualization of VMG in different tasks
21

---

Published as a conference paper at ICLR 2023
(a) antmazz-umarz-divvrss
(b) VMG antmaze-uaaze-diverse
+
(c) antmaze-medium play
(d) VMG of antmaze-
+
(e) intmare-meddum-diiveee
(f) VMG of antmare-medium--dierre
2P
antmaze-large-play
(h) VMG of antmaze-large-play
Figure 10: Visualization of VMG in different tasks
22

---

Published as a conference paper at ICLR 2023
(a) antmazz-large-diveree
(h) VMG of antmaze-large-iivrse
Vi-pu
(c) pen-human
(d) VMG of pen-human
(e) pen-cloned
(f) VMG of pen-cloned
gg)
hammer-human
(h) VMG of hammer-human
Figure 11: Visualization of VMG in different tasks
23

---

Published as a conference paper at ICLR 2023
(a) hammer-cloned
(b) VMG of hammer-cloned
(c) hoor-human
(d) VMG of door-human
(c) door-cloned
(f) VMG of door-cloned
Figure 12: Visualization of VMG in different tasks
24


Reviews:
[{"review_id": "doXuXrMliX", "reviewer": "Reviewer_b4KN", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes to group states in continuous space into clusters, according to a learned reachability metric. Such an operation allows turning the independent transition trajectories into a graph, allowing running value iteration for robust credit assignment. With shortest-path planning and an action translator, the agent can act to maximize value in a test environment while preventing OOD actions.\nThe empirical evaluations in the d4rl setting show the proposed VMG method surpasses some offline RL methods that applies conventional DQN-style bootstrapping.", "summary_of_the_review": "I think this is a good paper with significant contributions on the methodology level and the empirical performance of the method is also nice. \nWhereas it also has some problems like a false claim of SOTA and missed related work (thus overclaiming contributions). \nAlso, I have some questions mentioned in Strength And Weaknesses, which might be turned into weaknesses after the author's clarification.\n\nIn general, the current version of the paper is not perfect but I'm leaning to accept the paper because the problems should be easily fixed and other potential weaknesses won't surpass the main contributions.", "strength_and_weaknesses": "Strength\n* VMG provides a rather new approach to address OOD action problems in offline RL, which may have further impacts.\n* The proposed clustering method allows graph-based credit assignment to continuous control.\n\nWeakness\n* False claim of SOTA in the abstract. There are lots of offline RL methods that perform better than CQL in d4rl. For model-free methods, there are ATAC[1] and SAC-N[2], etc. And for model-based, there is a Trajectory Transformer[3] and COMBO[4].\n* Credit assignment in the transition data graph has been investigated in discrete state/action space in the online RL setting [5,6]. While being able to work in a continuous control is impressive, this work is not the first one that tried to leverage the graph-structured replay data.\n\nMinor problems and questions:\n* Why gym-locomotion control results are not presented? I agree the tasks themselves are less interesting but they are most of the offline RL methods are developed and tested on. The empirical results of this work are done on goal-reaching tasks (although some of them, like adroit, have shaped rewards), which might be helpful for the VMG to gain advantages. I would encourage authors to show gym-locomotion results, say, in the appendix if they do have the results. If the VMG indeed struggle with gym-locomotion, it will be better to state the scope of the current state of method clearly.\n* I'm not sure if it's appropriate to call VMG a world model. Although the term world model itself is very vague, it usually means a dynamics model plus a reward model. A dynamics model should be a parameterised model that can predict the future state given the current state and action. On the other hand, VMG only memorises and organises existing data and uses it for value estimation and building a policy at test time. To me, this is more like a model-free method with data set organised as a graph.\n* How's the training and inference cost of VMG?\n\nLinks:\n- [1] ATAC ([https://arxiv.org/abs/2202.02446](https://arxiv.org/abs/2202.02446))\n- [2] SAC-N ([https://arxiv.org/pdf/2110.01548.pdf](https://arxiv.org/pdf/2110.01548.pdf))\n- [3] Trajectory Transformer (https://arxiv.org/abs/2106.02039)\n- [4] COMBO [https://arxiv.org/pdf/2102.08363.pdf](https://arxiv.org/pdf/2102.08363.pdf)\n- [5] Topological Replay buffer [https://openreview.net/forum?id=OXRZeMmOI7a]\n- [6] Graph Backup [https://openreview.net/forum?id=0UQqmPGuL4n]", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "3xkHAqWsgM1", "reviewer": "Reviewer_ZzyC", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes learning a graph structured abstract world model for offline reinforcement learning. Here the learned abstract world model is built on top of a learned representation and used in conjunction with Value Iteration solver, Dijkstra’s search, and temporal inverse dynamics model. In general learned abstract models are powerful when combined with optimal planning and look-ahead policies. Here the authors use a clustering based approach to build the world model and show empirical evaluation of such an approach on a subset of D4RL tasks. Here the empirical evaluation is outperforms the current SOTA methods for the selective domains shown in the paper.", "summary_of_the_review": "In the current state of the paper, I will favor rejection. however I look forward to the authors rebuttal on my questions / concerns and am open to updating the score afterwards. ", "strength_and_weaknesses": "While the empirical results in isolation is very convincing, the paper lacks rigor in different design choices of the algorithm. The paper will definitely provide more value to the reader with better explanation on the design choices along with more discussion on the implementation details as outlined below.  would be a good contribution  VI and leveraging the graph for planning in continuous domain.\n\nStrengths:\n\n- The paper definitely boasts some empirical novelty as optimal planning on abstract graphs have not been explored well for d4rl tasks. The approach discussed performs well for the  shown subset of these d4rl tasks.\n    \n- The ablation studies presented does hint towards the contribution of different loss functions used towards representation learning / metric learning which is employed by the graph building algorithm.\n    \n- The notion of building a abstract graph for exact planning and combining it with lookahead policy (employing goal conditioned inverse dynamics) is powerful and broad in its scope\n    \n\nWeakness:\n\n- The design choices made throughout the approach lacks some rigorous justification as no theoretical evidence has been provided for these choices.\n    \n    - [C1]  (eq 1)Why was m set to 1 ? this value will definitely have bunch implications depending on the dimensionality of the latent space and the complexity of the environment dynamics.\n        \n    - [C2] (eq 2) why would we take max of the l2 norm and M. could we not have replaced this with max(l2(norm)-m, 0). In the current formulation this component of the loss will never go to zero if I understand it correctly.\n        \n    - [C3] (eq 4) We are basically ignoring any “self loops” that arise with this kind of clustering and only considering dynamics that take us from one state to another. Why is this? How often do these self loops occur. Moreover, the reward design from vj1 to vj2 seems arbitrary and not backed by any theory. Here we are basically assuming uniform weighting over all states in a cluster. I understand that adding half of the rewards from both self loops in source and target state, allows the trajectory reward to sum up to the true trajectory reward, but the reward definition here does seem arbitrary, which is hard to justify with empirical evaluations alone.\n        \n    - [C4] I did not quite get the notion of using Dijkstra’s search to find the best path across the world model, using one more weighting mechanism that is not theoretically justified. Would the Value iteration already not provide a relatively optimal path from any state to the best state reachable ?\n        \n    - [C5] The choice of k = 10 (for inverse dynamics) seems arbitrary as we are only using 1 or 2 step lookahead goals for the models. It seems that we are baking more into the prior as 1 or 2 step lookahead in the abstract space can technically be much more than 1 or 2 as we are aggregating states to a single cluster. i.e. a transition in abstract space may amount to more than one step in the real world, how much of this is actually happening in empirical evaluations is unclear.\n        \n    - [C6] The choices of penalty beta (from table 6 in appendix) seems low and arbitrary. generally we set this to 0.99. It does make sense that we would like to have smaller discount factors in abstract models as well as constrain it further to reduce cumulative errors, however the choices made here even without a hyper-parameter search is not rigorous by any means.\n        \n- Lacks some key implementation details.\n    \n    - What was the dimension of the latent space used for the latent dynamics model.\n        \n    - No training curves has been presented.\n        \n    - When we mention ““best one from the checkpoints saved from the 500th to the 800th epochs.” ([pdf](zotero://open-pdf/library/items/NHHMIX9V?page=15)) does it mean best in terms of the training loss / validation losses?\n        \n    - The values used for hyper-parameter search are not listed down, only the best performing ones. Moreover the number of runs used for hyper-parameter search must be mentioned in the main paper for clarity, as it is one of the crucial yet generally overseen.\n        \n    - Again the choices made for these hyper-parameter sweeps does not seem to be well justified.\n        \n- The paper is missing some key connections in the related work section.\n    \n    - Bisimulation metrics has been well explored for learning abstract models for a long time now. [4, 5, 6]. They provide a rigorous frame  and other abstraction literature is not cited well.\n        \n    - Secondly optimal planning on graph structured world models has also been explored in works like [1,2,3]. [1] Here is especially relevant as the world model is constructed using metric learned in abstract space. Careful comparison / discussion of these related works would greatly bolster the context of the paper.\n        \n\n[1] Shrestha, A., Lee, S., Tadepalli, P., & Fern, A. (2021). DeepAveragers: Offline Reinforcement Learning by Solving Derived Non-Parametric MDPs. *ArXiv, abs/2010.08891*.\n\n[2] Char, I., Mehta, V., Villaflor, A., Dolan, J.M., & Schneider, J.G. (2022). BATS: Best Action Trajectory Stitching. *ArXiv, abs/2204.12026*.\n\n[3] Marklund, H., Nair, S., & Finn, C. (2020). Exact (Then Approximate) Dynamic Programming for Deep Reinforcement Learning.\n\n[4] Ferns, N., & Precup, D. (2014). Bisimulation Metrics are Optimal Value Functions. *UAI*.\n\n[5] Ferns, N., Panangaden, P., & Precup, D. (2011). Bisimulation Metrics for Continuous Markov Decision Processes. *SIAM J. Comput., 40*, 1662-1714.\n\n[6] Bertsekas, D.P. (2009). Neuro-Dynamic Programming. *Encyclopedia of Optimization*.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "3j24NEI6FzG", "reviewer": "Reviewer_Wbjx", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The main focus of the paper is to tackle some critical challenges in continuous control RL tasks including sparse rewards and long-horizon planning. The idea is to build a simple and discrete world model leveraging an offline dataset (and the paper focuses on the offline RL setting) to learn an abstraction of the original environment. The authors propose Value Memory Graph that represents the MDP as a directed graph with \"abstract\" state and actions. Then, value iteration can be applied given the small state and action spaces. The authors also learn an action translator to convert abstract actions to real actions later. The authors conduct experiments on some of the tasks from D4RL to evaluate its effectiveness.", "summary_of_the_review": "The authors propose a novel method to tackle some of the important challenges in DRL. I think the experimental section can be improved by evaluating VMG on commonly-used MuJoCo tasks from D4RL. It can also be improved to extend VMG to online RL problems.", "strength_and_weaknesses": "### Strengths\n- Originality: The authors propose an interesting method to tackle sparse rewards and long-horizon planning, which is different from existing methods (such as improving exploration in sparse reward tasks). Specifically, the paper follows the \"abstraction\" perspective, which builds an abstract graph of the underlying MDP, which leads to a small state and action spaces that allow for tractable value iteration. To achieve this goal, the authors propose a contrastive learning method for state merging, and also propose an action translator to further translate the abstract action to a real action.\n\n- Clarity: The paper is clearly written and easy to follow. I think the paper motivates the problem well, and describes the procedure clearly including how to build a metric space, how to construct the graph in value memory graph, and how to use such value memory graph.\n\n### Weaknesses\n- Quality and significance: My main concern for the paper is the applicability of the method and the significance of the experiments. The paper can be improved by studying the method in the online RL setting (I understand that it is more difficult and challenging to apply the method in online RL settings, but I think that could greatly improve the paper.) In addition, I wonder how the method compare against baselines in commonly-used MuJoCo tasks from D4RL like hopper, walker2d, and halfcheetah. These are tasks with dense rewards, but they also have continuous state and action spaces, and it is worth evaluating VMG in these tasks (since this is also one of the challenges highlighted in the paper). I also have a question in Eq. (2). Why does it need to consider all the other next states? What if sj' and si' are very similar? I will be happy to increase the score if additional experiments on MuJoCo are included and analyzed.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "1aKFpKR_VDR", "reviewer": "Reviewer_UqHS", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper introduces Value Memory Graphs (VMG), a learned latent MDP structured as a graph. Since the VMG is relatively small, we can apply value iteration in it. The obtained value function is then used in a planner (Dijkstra’s), and the resulting plan in abstract space is converted back to the original space using a learned action decoder. Experiments in several tasks show improved performance over baselines, and investigate the importance of each component through ablation studies. ", "summary_of_the_review": "I recommend this paper for acceptance. It is well-written, covers a relevant topic, has good results, and extensive insight/ablation studies. There are some things to improve, as listed above, but overall I think this would be a valuable addition to the ICLR conference programme. ", "strength_and_weaknesses": "Strength:\n* The paper studies a relevant topic (structured abstraction & offline RL). \n* The paper is very well written and clear. It does not make things unnecessarily complex, while all necessary details are provided. \n* Performance of the methods clearly outperforms the baselines in the tested tasks. \n* Good ablation experiments. \n* Related work is well covered. \n* Good visualizations, such as Figure 4, which really show what the method does. \n\nWeaknesses: \n* Sec 3.2: Since each vertex is essentially a D-dimensional ball with radius gamma. These balls can not cover the full space without overlap. What happens if a certain point can be assigned to multiple vertices?\n* Sec 3.3: I understand why you construct the combined reward function from internal and transition rewards (with weight ½ on each internal half), but there is a strong assumption here, being that there are as many true steps internal of a vertex as there are between vertices. This will probably not be true in practice: especially with larger values of gamma, the number of steps within a vertex will be much bigger than the number of steps between them. I guess one could argue that you need extra weight on the internal components in the reward function.  \n* Sec. 3.4: I was highly surprised you used Dijkstra’s, after first having solved the problem with VI. I had written a comment about this, only to find out why you make this choice in the “Multi-step Search” section of the ablation (p 9). Try to explain why you make certain choices directly when you introduce them. \n* You only test on three tasks of the D4RL. This is fine, but some variation in environment type would have made the results a bit stronger. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster