Title: ANQ: Approximate Nearest-Neighbor Q Learning
Venue: NeurIPS 2023
Accepted: no
Keywords: reinforcement learning

explainable reinforcement learning

semi-parametric learning

Memory-based learning
Abstract: In specific domains such as autonomous driving, quantitative trading, and healthcare, explainability is crucial for developing ethical, responsible, and trustworthy reinforcement learning (RL) models. Although many deep RL algorithms have attained remarkable performance, the resulting policies are often neural networks that lack explainability, rendering them unsuitable for real-world deployment. To tackle this challenge, we introduce a novel semi-parametric reinforcement learning framework, dubbed ANQ (Approximate Nearest Neighbor Q-Learning), which capitalizes on neural networks as encoders for high performance and memory-based structures for explainability. Furthermore, we propose the Sim-Encoder contrastive learning as a component of ANQ for state representation. Our evaluations on MuJoCo continuous control tasks validate the efficacy of ANQ in solving continuous tasks while offering an explainable decision-making process.

1

Abstract

In specific domains such as autonomous driving, quantitative trading, and healthcare, explainability is crucial for developing ethical, responsible, and trustworthy

reinforcement learning (RL) models. Although many deep RL algorithms have 3

attained remarkable performance, the resulting policies are often neural networks 4

that lack explainability, rendering them unsuitable for real-world deployment. To 5

tackle this challenge, we introduce a novel semi-parametric reinforcement learning 6

framework, dubbed ANQ ( A pproximate N earest Neighbor Q -Learning), which cap7

italizes on neural networks as encoders for high performance and memory-based 8

9

structures for explainability. Furthermore, we propose the Sim-Encoder contrastive

10

11

12

learning as a component of ANQ for state representation. Our evaluations on Mu-

JoCo continuous control tasks validate the efficacy of ANQ in solving continuous tasks while offering an explainable decision-making process.

1 Introduction 13

St

Black Box - Neural Network Policy

at

ANQ (ours)

ExplainableMemory

ng = argmaxn(qr..qm)

Key

Values

Approximate

St

G

→

et

Neighbor

Nearest

e1

S1

a1

r1

S2

Q1

at = ana +N(μ,o)

p←

Search

e2

82

@2

r2

82

q2

Explanations

***

ek

Sk

ak

Sk

qk

trainedwithOne-Step-AwayState

Encoding Contrastive Learning.

Conventional Deep Reinforcement Learning

In recent years, parametric reinforcement learning methods featuring end-to-end training, such as 14

Proximal Policy Optimization (PPO) [Schulman et al., 2017], Soft Actor-Critic (SAC) [Haarnoja 15

et al., 2018], and Deep Deterministic Policy Gradient (DDPG) [Lillicrap et al., 2015], have gar16

nered significant attention within the reinforcement learning community. These approaches have 17

demonstrated remarkable success in addressing decision-making challenges across diverse domains, 18

including robotics [Hwangbo et al., 2019], video games [Mnih et al., 2015], and board games [Schrit19 twieser et al., 2020]. Nevertheless, the incorporation of deep neural networks in these methods 20 presents a major obstacle to interpreting the underlying rationale of their decision-making processes. 21 This limitation hampers the application of such methods to numerous real-world scenarios, such 22 as autonomous driving [Kiran et al., 2021], quantitative trading [Zhang et al., 2020], and beyond. 23 Consequently, further investigation is necessary to enhance the interpretability and practical utility of 24 these reinforcement learning techniques in complex, real-world contexts. 25

26

27

28

29

30

31

This issue calls for the research of explainable reinforcement learning (XRL) which aims at obtaining RL models that are both explainable and of high performance. Fidelity is one of the major objectives in XRL [Milani et al., 2022] which measures to what extent the model makes decisions following its explanation. Among different XRL algorithms, white-box algorithms (i.e., making decisions directly using explainable models such as linear models or decision trees) enjoys high fidelity than the others. (We defer the introduction of other XRL algorithms to Section 5.3.)

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

Memory-based reinforcement learning, following the non-parametric paradigm, is a popular class of white-box algorithm and differs from widely researched parametric methods in deep reinforcement learning. The approximation function in memory-based reinforcement learning is determined directly by the training samples, rather than relying on a gradually-updated parameterized function. Prominent memory-based methods include EC [Blundell et al., 2016], NEC [Pritzel et al., 2017], and EMDQN [Lin et al., 2018] (see more in Ramani [2019]). Memory-based reinforcement learning has several benefits, including being able to approximate a universal class of functions, the ability to directly impact the policy with newly accumulated data without back-propagation updates Blundell et al. [2016], the mitigation of the curse of dimensionality in global estimation Sutton and Barto [1998], and higher data sampling efficiency and faster learning Lin et al. [2018]. Most importantly, memorybased reinforcement learning possesses the advantage of improved explainability due to its humanunderstandable decision making system (i.e., the memory consists of pre-collected samples).

Despite its potential for self-explainability through white-box decision-making, the utilization of memory-based reinforcement learning for enhancing explainability remains relatively unexplored. Existing studies investigating the use of episodic memory for explanations, such as Cruz et al. [2019], Pritzel et al. [2017], Blundell et al. [2016], have been limited to grid world environments or discrete tasks. In contrast, our work aims to expand this research scope to encompass continuous robotics tasks in Mujoco by proposing a comprehensive memory-based self-explainable framework.

Efficiently retrieving relevant data from extensive databases presents a significant challenge in developing an effective memory-based reinforcement learning algorithm, particularly in continuous control tasks as emphasized by Sutton and Barto [1998]. However, recent advancements in approximate nearest-neighbor searching algorithms, such as Hierarchical Navigable Small World (HNSW) Malkov and Yashunin [2018], have demonstrated their effectiveness in swiftly retrieving pertinent information from billions of records in natural language processing (NLP) tasks. Such methods have been successfully applied to question-answering Kassner and Schütze [2020] and text generation Borgeaud et al. [2022] tasks. In addition to NLP applications, retrieval-based systems have been integrated with deep reinforcement learning algorithms, resulting in enhanced sample efficiency Goyal et al. [2022], Humphreys et al. [2022].

The contributions of our paper are summarized as follows:

We introduce a novel framework, ANQ, which offers efficient control in continuous domains across a wide range of Mujoco experiments, while maintaining high explainability through its "data is policy" design principle.

We present the Sim-Encoder, a nearest neighbor contrastive learning approach for state representation, which demonstrates its effectiveness in memory retrieval learning tasks.

2 Preniminaries 66

We first introduce notations and summarize the conventional episodic control method. 67

68

69

70

71

72

73

74

75

76

77

78

79

80

2.1 Notation

In this work, we study policy learning in continuous action space A and observation space S . We consider a Markov decision process with transition s t +1 ∼ p ( s t +1 | s t , a t ) . After performing an action, the agent receives a reward, and the ultimate goal is to optimize the policy to maximize returns.

A key-value-based dataset D stores the key as the state embedding e . The database consists of rows of { k, e t , s t , a t , r t , s t +1 , q t } ∈ D and columns of { K , E , S t , A , R , S t +1 , Q } ∈ D . K represents the set of all record IDs. The maximum number of rows is M . The observation Sim-Encoder network is denoted as G θ parameterized by the network parameters θ .

For database operating, in total, six operations are defined in the memory module: APPEND, TRIM, GET, UPDATE, SEARCH, and INDEX. More corresponding explanations for these operations can be found in Sec.3.2.

2.2 Episodic Control

Episodic control methods enhance sampling efficiency and episodic returns by using an external 81 memory database for interactions such as writing, reading, and updating. The concept was first 82 introduced in Blundell et al. [2016], which resolved complex sequential decision tasks. 83

This method is defined for discrete spaces. It proposes the following Q table update mechanism: 84

After the update, it generates an effective Q Table. During the policy execution phase, if an 85 observation-action pair exists in memory, the Q value is retrieved directly from the table. However, if 86 the pair is not found in memory, an approximation matching and estimation process is required. The 87 agent queries the Q Table using the following approach to obtain the Q value. 88

The objective of episodic control is to accelerate learning speed and improve decision quality. An 89 external memory module can then compensate for drawbacks such as low sample efficiency and slow 90 gradient updates. 91

In previous discussions, the Episodic Control (EC) method has been investigated under both discrete 92 actions and continuous actions (Li et al. [2023], Kuznetsov and Filchenkov [2021]). However, the 93 explainability of EC in continuous action spaces suffers from low fidelity due to the utilization of a 94 policy network. In this paper, we set out to achieve two objectives concurrently. First, we explore 95 how Episodic Control can be effectively applied in continuous action spaces. Second, we strive to 96 leverage the memory of Episodic Control to attain explainability benefits. 97

98

99

100

101

102

103

104

105

3 Method

The complete algorithm is presented in Algorithm 1, and the illustration of the inference pipeline can be observed (cf. Fig.1). The proposed method involves generating an embedding vector e t from the observation using the Sim-Encoder. Subsequently, we employ the HNSW algorithm Malkov and Yashunin [2018] to search for the nearest neighbor set e n within the memory. Each neighbor is associated with an action and a Q value, and the action with the highest Q value is selected as the policy output. It is worth noting that this action is continuous, which distinguishes it from previous EC work Blundell et al. [2016].

First, the Sim-Encoder in embedding observations into a cosine space is augmented via One-Step106 Away State Encoding Contrastive Learning. This approach employs adjacent states as positive 107 samples for contrastive learning, with experimental outcomes demonstrating that the implementation 108 of the Sim-Encoder considerably enhances performance. 109

110

111

112

113

114

115

116

117

Algorithm 1 ANQ Algorithm

Input:

Subsequently, in order to acquire a comprehensive Q-table, we employ in-memory learning, which involves the batch computation of all Q-value estimations and Q-learning updates for each state stored in memory. The training process undergoes iterative cycles until the global Q-value converges.

3.1 Embedding Module

We introduce our novel approach, the "One-Step-Away State Encoding Contrastive Learning." The reason for using a one-step-away state as a positive sample is that the most informative actions and q-values for the current state are derived from a scenario that is most similar to it (Blundell et al. [2016]).

This method aims to effectively represent the state with contrastive learning. Specifically, we utilize 118 positive samples that consist of a state pair s t , s t +1 that are one step away. The resulting state 119 representation is designed such that the nearest neighbor of each state is reachable within one step. 120 We adopt a similar objective to SimCLR Chen et al. [2020], aiming to maximize the similarity 121 between two vectors as measured by cosine similarity sim ( u, v ) = u T v/ ( | u || v | ) . The Sim-Encoder is 122 a standalone component trained to maximize the similarity of embedding, without reward information 123 but only state transition tuples. 124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

3.2 Memory Module

The explainable memory module is in the form of a key-value database. And the keys in the database correspond to the observation embedding vectors obtained via the Sim-Encoder, and each key is associated with a corresponding value that includes information such as the current step's observation, action, reward, and all of other relevant data. To manage this database, we have defined 6 standard operations, namely APPEND, TRIM, GET, UPDATE, INDEX, and SEARCH, which are detailed in Algorithm 1. and Table.1.

The GET operation requires a target state's embedding as the key and returns the corresponding values. To prevent the database from becoming excessively large, we define a TRIM operation that automatically removes older data, retaining only the most recent M records. This design enables efficient storage and retrieval of data while ensuring that the database remains manageable and up-to-date.

In our approach for effective memory retrieval, Approximate K-Nearest Neighbors Search (AKNN) plays a crucial role. We introduce a SEARCH operation that takes a state embedding as input and returns the corresponding key(s) of the nearest neighbor(s) in the database. Additionally, we define an INDEX operation, activated when the database undergoes modifications, seeing Malkov and Yashunin [2018]. This operation reorganizes the HNSW index to align with the updated database, ensuring that subsequent KNN searches remain both fast and accurate.

3.3 Policy Evaluation

We introduce the Approximate Nearest Neighbor Search Q-Learning method. In contrast to conventional tabular Q-Learning, we employ a novel form of state value estimation, ˆ V ( s t ) , by aggregating the Q-values from the nearest neighbors of the state (cf. Fig.2). The Q-value of each state-action pair is updated following the Bellman equation, incorporating a decay factor, γ .

During the practical training process, we adopt a batch updating strategy wherein we simultaneously 148 compute the labels for all neighbors of each state and estimate the values of all states in memory. 149 Subsequently, we update all Q-values in the table accordingly. The learning iteration persists until the 150 maximum change in Q-values falls below a specified threshold. 151

3.4 Policy Improvement 152

For policy improvement, our proposed method directly selects the action with the maximum Q-value 153 from the neighbors ( e n ∈ E , a n ∈ A , q n ∈ Q ), as shown in Equation 10. Using the embedding e t 154

/n=

/n=

3

S1

/:

/n=

71

2

(0 + z + 10)

(=+0+)

3

3

/n:

/n=

/n=

SEARCH/GET

UPDATE

memoryoperation

memoryoperation

_d3-JM

Wallker2d-v3

60003

Anit-v3

HalfCheetah-v3

3500

ars-7M_

4000

S-IM

5000

_pa-LM.

-SAC.M.

10000

3000

wunqau abeuane

3000

-as:75M.

PRO-UL

wngau abeaae

4000

c-1M

ungau

8000

2500

3000

6000 =PRPUM

2000

2000

2000

4000

-sJ2M.

1500

1000

a2r-1M..

1000

2000

a2r-UM

1000

500

a2r-1M

alr-UM

2

million steps

6

10

2

4

million steps

6

10

4

million steps

6

10

Z

4

milion steps

6

10

generated by the contrastively learned Sim-Encoder, candidates are retrieved from the memory 155 module via the nearest neighbor search. To encourage exploration, action noise following a Gaussian 156 distribution N ( µ, σ ) is added. 157

In contrast to employing a black-box network as an actor, we have devised a data-driven, self158 explaining actor that seamlessly integrates the results of model search and generates decisions directly 159 using a rule-based approach. 160

161

162

163

164

165

4 Experiments

In these experiments, we aimed to evaluate the performance of our ANQ approach in solving continuous control tasks in Mujoco, provide action explainability, and investigate the significance of the Sim-Encoder module in the ANQ framework.

4.1 Solving Continuous Control Task in Mujoco

First of all, our approach is evaluated on several continuous control tasks in the MuJoCo physics 166 engine. Specifically, we compare our method with state-of-the-art reinforcement learning (RL) 167 algorithms, including SAC-1M, PPO-1M, and TRPO-1M, on the Walker2d-v3, Ant-v3, HalfCheetah168 v3, and Hopper-v3 environments. We use the benchmark performance reported by stable-baselines3 169 Raffin et al. [2021]. 170

The results (cf. Fig.3) show that our method slightly outperforms A2C-1M on the Walker2d-v3 171 task and PPO-1M on the Ant-v3 task while achieving comparable performance to TRPO-1M on the 172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

Question:

Why didyouexecute

Answer:

Based onmymemory,Ichooseto execute

due to its highest Q value of 320.9

-Upon

iflhadexecuted

1would haverecieved a return withexpectation320.9

-Upon

,iflhadexecuted

,Iwould haverecievedareturnwithexpectation318.8

-Upon

ifIhadexecuted

,Iwouldhaverecieved areturn with expectation318.6

-Upon

iflhadexecuted

,1would haverecievedareturn withexpectation318.5

HalfCheetah-v3 task. Furthermore, we analyze the performance of our method on the Hopper-v3

task by examining the game replay. We find that the agent fails to take the second step and falls after the first step. This indicates that our method may currently lack exploration capability. This will be

addressed in future research in order to surpass the performance of traditional RL methods. Overall, our approach successfully and stably converges on the MuJoCo continuous control tasks, but further

improvement is necessary to achieve better performance, seeing the discussion in Sec.6.

For the hyperparameters, we utilized a 4-layer MLP network with layer normalization as the encoder. The learning rate was set to 0.0003, the batch size was 512, and the Adam optimizer was used. The size of the explainable memory was limited to 500,000, and old data were discarded once this limit was exceeded. We set the parameters of HNSW to M=16 and ef=10. The total number of training steps was 10 million, and the agent performed ANQ learning every 40,000 environment interactions. We set the number of neighboring actions sampled during each action selection to 10.

4.2 Action explainability

In this explainability experiment, we designed a question-and-answer (QA) case (cf. Fig.4) to simulate a scenario where humans need to double-check the correctness of the robot's decision during human-robot collaboration. Specifically, humans ask "why" questions to query the basis of the robot's action, and the robot responds with the policy that it has chosen, as well as the evidence supporting its decision.

To provide a convincing explanation, the robot searches its memory for similar states and explains to the human the actions that it had taken in the past in similar scenarios, as well as the corresponding returns. By providing such detailed explanations, the robot is able to offer valuable insights to humans and effectively bridge the gap in understanding between human and machine decision-making processes, for ensuring safe and reliable human-robot collaboration.

4.3 Sim-Encoder

We conducted experiments to investigate the significance of the Sim-Encoder module within the ANQ framework. We have illustrated the retrieved samples (cf. Fig.5). Without the Sim-Encoder, semantically similar states do not share relevant information in cosine space, as discussed in Su et al. [2021]. Our ablation study (cf. Fig.6) demonstrated that the Sim-Encoder led to substantial performance improvements across all four tested tasks, as it effectively retrieves and embeds temporally

QueryState

RetrievedStates

1400

AQL(with sim-encoder)

1350

AQL(no sim-encoder)

400

1300 -

averageretum

1250 -

averagereturn

300

1200

1150 -

200

1100

1050 -

100

AQL(with sim-encoder)

1000

AQL(no sim-encoder)

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

million steps

million steps

HalfCheetah-v3

260

1500

AQL(with sim-encoder)

240 -

AQL(no sim-encoder)

1250

averagereturn

220 -

averageretum

1000

200

750

180

500

160

250 -

140

AQL(with sim-encoder)

120

AQL(no sim-encoder)

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

million steps

million steps

proximate states into a space with adjacent cosine distances. Overall, the Sim-Encoder is an essential 202 component of the ANQ framework and holds potential for use in other RL algorithms. 203

204

205

5 Related Work

5.1 Episodic Control

The idea of episodic control (EC) was bio-inspired by the mechanism of the hippocampus Lengyel 206 and Dayan [2007]. EC, as a non-parametric approach, possesses virtues including rapid assimilation 207 of past experiences and a solution for sparse-reward situations. Notable works like MFEC Blundell 208 et al. [2016] and NEC Pritzel et al. [2017] employed kNN search to acquire the value for the current 209 state derived from similar states. The value function is in tabular form and updated using the classical 210 Q-learning method. While MFEC adopts random projection Johnson [1984] and VAE Kingma and 211 Welling [2013] as state embedding methods, NEC employs a differentiable CNN encoder instead. 212

Beyond that, Lin et al. proposed EMDQN Lin et al. [2018], which is a synergy of EC and DQN. 213

Their approach combined the merits of both algorithms, i.e., fast learning at an early stage and good 214

final performance. ERLAM then further promoted the efficacy by introducing an associative memory 215

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

5.2 Retrieval-based Learning

The retrieval-based learning and inference architecture provides a viable solution for managing an explainable and extensible knowledge base. One prominent instantiation of this architecture is the retriever-reader model Zhu et al. [2021], which has gained traction in the open domain question answering (openQA) research community. The retriever component returns a set of relevant articles, while the reader extracts the answer from the retrieved documents. Numerous natural language processing (NLP) algorithms, including kNN-LM Khandelwal et al. [2019], RAG Lewis et al. [2020] and RETRO Borgeaud et al. [2022], leverage a retrieval-based approach to enhance their performance and efficiency. These techniques have proven to be effective in the domain of NLP and continue to be an active area of future NLP research Liu et al. [2023].

5.3 Explainable Reinforcement Learning

The methods for explainability in reinforcement learning can be broadly categorized into three groups, as discussed in Milani et al. [2022]: Feature Importance (FI), Learning Processing and Markov Decision Process (LPM), and Policy-Level (PL). FI methods involve utilizing decision tree models for explainability, learning an explainable surrogate network through expert and learner frameworks, or directly generating explanations through natural language or saliency maps. LPM addresses explainable transition models to answer "what-if" questions, interpretation of Q values, and identification of key training points. PL provides an understanding of long-term behavior and summarizes the policy. However, many existing explainable reinforcement learning methods require additional network training Guo et al. [2021] or the use of decision trees Silva et al. [2019]. These methods can also impose a cognitive burden on users to understand the model's behavior Dodge et al. [2021]. In contrast, the memory-based reinforcement learning algorithm, ANQ, presented in this paper provides self-explainability without additional explanation specifically training.

6 Limitation

In this study, we present an innovative and explainable architecture, termed ANQ, which, despite its novelty, does not significantly outperform state-of-the-art benchmarks. Our primary aim is to demonstrate the efficacy of ANQ with its highly interpretable policy. We acknowledge this performance gap and recognize that our method has not yet incorporated the latest techniques, such as maximum entropy learning from SAC Haarnoja et al. [2018], etc. These refinements will be addressed in future work, rather than here. Moreover, we have not compared our approach with other contrastive learning methods for representation learning. Since we proposed the Sim-Encoder, a thorough comparison with alternative methods and further study will also be included in future research.

7 Conclusion

Explainability is crucial in specific domains of reinforcement learning, such as autonomous driving, quantitative trading, and healthcare. To address this challenge, we propose ANQ, a novel semiparametric reinforcement learning framework that combines the high performance of neural networks with the explainability of a memory-based structure. Additionally, we validate the effectiveness of Sim-Encoder, a key module of ANQ, in state representation and learning efficiency enhancement. Empirical evaluations demonstrate ANQ's effectiveness in solving continuous tasks and providing explainable decision-making. Our contributions include proposing a framework that achieves both efficient control and robust explainability. While further improvements are necessary for superior performance, our results indicate that ANQ is a promising approach for developing explainable and trustworthy RL models in critical applications.

8 Reference

References

Charles Blundell, Benigno Uria, Alexander Pritzel, Yazhe Li, Avraham Ruderman, Joel Z Leibo, Jack Rae, Daan Wierstra, and Demis Hassabis. Model-free episodic control. arXiv preprint arXiv:1606.04460 , 2016.

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. In International conference on machine learning , pages 2206-2240. PMLR, 2022.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning , pages 1597-1607. PMLR, 2020.

Francisco Cruz, Richard Dazeley, and Peter Vamplew. Memory-based explainable reinforcement learning. In AI 2019: Advances in Artificial Intelligence: 32nd Australasian Joint Conference, Adelaide, SA, Australia, December 2-5, 2019, Proceedings 32 , pages 66-77. Springer, 2019.

Jonathan Dodge, Andrew Anderson, Roli Khanna, Jed Irvine, Rupika Dikkala, Kin-Ho Lam, Delyar Tabatabai, Anita Ruangrotsakun, Zeyad Shureih, Minsuk Kahng, et al. From 'no clear winner' to an effective explainable artificial intelligence process: An empirical journey. Applied AI Letters , 2(4):e36, 2021.

Anirudh Goyal, Abram Friesen, Andrea Banino, Theophane Weber, Nan Rosemary Ke, Adria Puigdomenech Badia, Arthur Guez, Mehdi Mirza, Peter C Humphreys, Ksenia Konyushova, et al. Retrieval-augmented reinforcement learning. In International Conference on Machine Learning , pages 7740-7765. PMLR, 2022.

Wenbo Guo, Xian Wu, Usmann Khan, and Xinyu Xing. Edge: Explaining deep reinforcement learning policies. Advances in Neural Information Processing Systems , 34:12222-12236, 2021.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 1861-1870. PMLR, 10-15 Jul 2018.

Peter C Humphreys, Arthur Guez, Olivier Tieleman, Laurent Sifre, Théophane Weber, and Timothy Lillicrap. Large-scale retrieval for reinforcement learning. arXiv preprint arXiv:2206.05314 , 2022.

Jemin Hwangbo, Joonho Lee, Alexey Dosovitskiy, Dario Bellicoso, Vassilios Tsounis, Vladlen Koltun, and Marco Hutter. Learning agile and dynamic motor skills for legged robots. Science Robotics , 4(26):eaau5872, 2019.

William B Johnson. Extensions of lipschitz mappings into a hilbert space. Contemp. Math. , 26:189-206, 1984.

Nora Kassner and Hinrich Schütze. Bert-knn: Adding a knn search component to pretrained language models for better qa. arXiv preprint arXiv:2005.00766 , 2020.

Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. arXiv preprint arXiv:1911.00172 , 2019.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

B Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A Al Sallab, Senthil Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey. IEEE Transactions on Intelligent Transportation Systems , 23(6):4909-4926, 2021.

Igor Kuznetsov and Andrey Filchenkov. Solving continuous control with episodic memory. arXiv preprint arXiv:2106.08832 , 2021.

Máté Lengyel and Peter Dayan. Hippocampal contributions to control: the third way. Advances in neural information processing systems , 20, 2007.

Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-augmented generation for knowledgeintensive nlp tasks. Advances in Neural Information Processing Systems , 33:9459-9474, 2020.

Zhuo Li, Derui Zhu, Yujing Hu, Xiaofei Xie, Lei Ma, Yan Zheng, Yan Song, Yingfeng Chen, and Jianjun Zhao. Neural episodic control with state abstraction. arXiv preprint arXiv:2301.11490 , 2023.

Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.02971 , 2015.

Zichuan Lin, Tianqi Zhao, Guangwen Yang, and Lintao Zhang. Episodic memory deep q-networks. arXiv preprint arXiv:1805.07603 , 2018.

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

Yiheng Liu, Tianle Han, Siyuan Ma, Jiayue Zhang, Yuanyuan Yang, Jiaming Tian, Hao He, Antong Li, Mengshen He, Zhengliang Liu, Zihao Wu, Dajiang Zhu, Xiang Li, Ning Qiang, Dingang Shen, Tianming Liu, and Bao Ge. Summary of chatgpt/gpt-4 research and perspective towards the future of large language models, 2023.

Yu A Malkov and Dmitry A Yashunin. Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs. IEEE transactions on pattern analysis and machine intelligence , 42(4):824-836, 2018.

Stephanie Milani, Nicholay Topin, Manuela Veloso, and Fei Fang. A survey of explainable reinforcement learning, 2022. URL https://arxiv.org/abs/2202.08434 .

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature , 518(7540):529-533, 2015.

Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adria Puigdomenech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural episodic control. In International Conference on Machine Learning , pages 2827-2836. PMLR, 2017.

Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stablebaselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research , 22 (268):1-8, 2021. URL http://jmlr.org/papers/v22/20-1364.html .

Dhruv Ramani. A short survey on memory based reinforcement learning. arXiv preprint arXiv:1904.06736 , 2019.

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):604-609, 2020.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

Andrew Silva, Taylor Killian, Ivan Dario Jimenez Rodriguez, Sung-Hyun Son, and Matthew Gombolay. Optimization methods for interpretable differentiable decision trees in reinforcement learning. arXiv preprint arXiv:1903.09338 , 2019.

Jianlin Su, Jiarun Cao, Weijie Liu, and Yangyiwen Ou. Whitening sentence representations for better semantics and faster retrieval. arXiv preprint arXiv:2103.15316 , 2021.

Sutton and A. Barto. Reinforcement Learning:An Introduction . Reinforcement Learning:An Introduction, 1998.

Zihao Zhang, Stefan Zohren, and Stephen Roberts. Deep reinforcement learning for trading. The Journal of Financial Data Science , 2(2):25-40, 2020.

Fengbin Zhu, Wenqiang Lei, Chao Wang, Jianming Zheng, Soujanya Poria, and Tat-Seng Chua. Retrieving and reading: A comprehensive survey on open-domain question answering. arXiv preprint arXiv:2101.00774 , 2021.

Guangxiang Zhu, Zichuan Lin, Guangwen Yang, and Chongjie Zhang. Episodic reinforcement learning with associative memory. 2020.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces the ANQ (Approximate Nearest Neighbor Q-Learning) framework, which aims to provide explainability in reinforcement learning models. ANQ combines neural networks for high performance and memory-based structures for explainability, offering a promising solution for domains like autonomous driving, quantitative trading, and healthcare. The paper discusses the challenges of explainability in reinforcement learning and how ANQ addresses them. It also presents the Sim-Encoder contrastive learning used in ANQ's state representation and provides insights into the evaluations of ANQ on MuJoCo continuous control tasks and its effectiveness in solving continuous tasks. Overall, the paper presents a novel approach to reinforcement learning that balances performance and transparency, making it suitable for real-world applications.\n\nSTRENGTHS:\nSome strengths of this paper include:  \n\n1. Novelty: The ANQ framework is a novel approach to reinforcement learning that combines neural networks and memory-based structures to provide explainability in decision-making processes. This approach is different from traditional reinforcement learning methods that focus solely on performance.  \n\n2. Explainability: The paper addresses the challenge of explainability in reinforcement learning models, which is crucial for real-world applications where transparency and trustworthiness are essential. ANQ's \"data is policy\" design principle ensures that the model's decisions are explainable and interpretable.  \n\n3. Sim-Encoder contrastive learning: The paper introduces the Sim-Encoder contrastive learning approach for state representation, which demonstrates its effectiveness in memory retrieval learning tasks. This approach enhances ANQ's performance and explainability.  \n\n4. Evaluations: The paper provides insights into the evaluations of ANQ on MuJoCo continuous control tasks and its effectiveness in solving continuous tasks. The results show that ANQ outperforms traditional reinforcement learning methods while maintaining explainability.  \n\n5. Real-world applications: The paper highlights the potential of ANQ in domains like autonomous driving, quantitative trading, and healthcare, making it suitable for real-world applications.\n\nWEAKNESSES:\nSome potential weaknesses of this paper include:  \n\n1. Despite the advantage of interpretability, the performance of this framework is still far from that of SOTA RL algorithm. I think it would be better if this framework could have similar performance to SOTA RL method. \n\n2. Lack of real-world case studies: Although the paper highlights the potential of ANQ in various domains, it does not provide specific real-world case studies or examples to demonstrate the practical application and effectiveness of the framework. Including such case studies would strengthen the paper's claims and provide more concrete evidence of ANQ's utility in real-world scenarios. \n\n3. No comparison was made with other interpretable reinforcement learning algorithms. For example, some Neuro-Symbolic Search methods. \n\n3. The presentation of this paper could be better, for example, Figure 3 could be larger and clearer.\n\nQUESTIONS:\n1. I think it is still important to have similar performance with recent SOTA algorithms. Is it possible to make your framework applicable to recent RL algorithms? \n\n2. How does your framework perform in different types of environments? For example, sparse reward environments, and discrete action space environments?\n\nLIMITATIONS:\nIncorporation of latest techniques: The paper mentions that ANQ has not yet incorporated the latest techniques, such as maximum entropy learning from SAC and other contrastive learning methods for representation learning. While the paper acknowledges that these refinements will be addressed in future work, the absence of these techniques in the current implementation may limit the overall performance and effectiveness of ANQ.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-18"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe submission creates a framework called Approximate Nearest Neighbor Q-Learning (ANQ). ANQ uses a sim-encoder contrastive learning and approximate nearest neighbor search to find which states are similar. Utilizing this approach, they can use it to find similar states in aiding for the decision the framework has made. They showcase their performance by performing experiments on MuJoCo with continuous control tasks. The exact environments experimented were Walker2d, Ant, HalfCheetah, and Hopper. There is also an ablation study conducted to show the benefit of the sim-encoder and contrastive learning that aids this framework. They provide a component called Explainable Action to show why it executed a particular action based on the state.\n\nSTRENGTHS:\nSignificance:\nGood ablation experiments to showcase the benefit of the sim-encoder. The results make it convincing that for AQL, the sim-encoder is quite beneficial. Plus you did this ablation among 4 environments.\n\nOriginality:\nThe explainable action is an interesting piece that can provide a good impact. By searching for similar states and can provide the explanations as to why the decision was made.\n\nWEAKNESSES:\nClarity:\nConfusion, in Section 4.1, you mention that the algorithms included were SAC, PPO, and TRPO, then in the next paragraph you mention A2C. Consider in the first paragraph to mention the exact algorithms you compare because in the next paragraph, you mention an algorithm that was not discussed exactly in the previous paragraph. In Figure 3, you show TD3-1M, ARS-75M which were not discussed so please in Section 4.1 to denote exactly.\n\nWith the figures, please provide more with the caption like a summarization or a sentence to showcase why it is important. It can help the reader if they have not read the parts within the main text.\n\nSignificance:\nThe approach can be on par with one or two deep RL approaches. Consider to improve the performance to have bigger impact. Usually for methods that are explainable there is a drop in performance so others may be hesitant to use it due to the performance drop.\n\nQUESTIONS:\nA suggestion: For experiments, it could benefit to utilize this framework for discrete actions spaces. Plus for impact, consider other MuJoCo examples that aren't locomotion, consider the fetch to show that it can be utilized with a plethora of applications. In addition, to extend this to image observation applications like atari.\n\nAlso what is the limitation approach with respect to the explainable memory, since in section 4.1, you mention the size is 500,000. Do you need that much for the performance?\n\nAlso please check in the weaknesses section for other suggestions.\n\nLIMITATIONS:\nThey do address the performance drop compared to the deep RL models. No negative societal impacts.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes a memory-based Q-learning algorithm aimed at better explainability. The method extends the prior work \"Episodic Control\" and enables learning with continuous action space.\n\nThis work provides 2 main contributions:\n1. a one-step-away contrastive-learning objective to learn embeddings from states.\n2. modified policy evaluation and improvement rules to account for the continuous action space.\n\nThe authors show empirical results to support their design choices:\n1. The proposed algorithm is able to achieve some meaningful learning in 4 continuous control tasks.\n2. The question-answering example shows that the method is able to find the nearest-neighbor states and their Q values to explain a chosen action.\n3. The ablation study on the embedding module shows that it is necessary to learn dynamics-aware embedding.\n\nSTRENGTHS:\n1. The proposed method can learn in environments with continuous state and action spaces, making it ubiquitously applicable to real-world applications.\n2. The algorithm and experiment settings are clear.\n3. The ablation study on the embedding module justifies the contrastive learning objective.\n\nWEAKNESSES:\n1. The main weakness of this work is in the experiment results. In the Mujoco tasks, the performance hovers around the weakest baseline among all compared methods. For Ant, HalfCheetah, and Hopper, the policies also appear to have converged to suboptimal ones. This could have been caused by insufficient exploration.\n\n2. The algorithm is only demonstrated in state-based environments, whereas the prior work \"Episodic Control\" can work with image observations. Contrastive learning has been shown to be useful for learning good feature extractors for images. The results would have been much more convincing if they were from vision-based tasks.\n\n3. The notations are not fully explained. For example, $k$ and $e_t$ both exist in the dataset but the authors say that they use embeddings as keys. Also, $R$ in Equation 1 is not introduced.\n\n4. Finding out the nearest neighbors and printing out their Q values is not a convincing way to explain the chosen actions because the Q values are computed in expectation. One can arguably explain a neural policy in a similar way, by sampling a few different actions and printing out their Q values.\n\nQUESTIONS:\nI'm not fully convinced by the way the authors compute the value target in Equation 6. The $(s, a)$ pairs in the memory come from old policies and thus, could be sub-optimal. Computed from sub-optimal $q^n$, $\\hat{v}$ is an underestimation if older interactions are not ejected quickly enough. This might explain why the task performance is bad.\n\nLIMITATIONS:\nThe authors are upfront about the limitation in task performance and provide viable options for improvement. I would also encourage the authors to try their method on image-based tasks.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}

comments: SUMMARY:
Instead of using deep neural networks to approximate Q functions as it's done in deep RL methods, the paper investigates the potential of using nearest neighbor methods to approximate Q functions. They used contrast learning with a Sim encoder. They argue that such a method is more explainable than the ones with deep neural nets.

STRENGTHS:
The illustration is the proposed method is clear.

WEAKNESSES:
* The proposed nearest neighbor might be useful and efficient for low-dimensional domains like Mujoco. I doubt its effectiveness when it goes to high-dimensional domains like video games;

* I don't see why the nearest neighbor method is more explainable;

* The proposed method lacks proper baselines;

* The overall presentation needs improvements.

QUESTIONS:
Please address the weaknesses points in the rebuttal.

LIMITATIONS:
NA
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-11