Title: Diffusion Model-Augmented Behavioral Cloning
Venue: NeurIPS 2023
Accepted: no
Keywords: Imitation Learning

Learning from Demonstration

Diffusion Models

Behavioral Cloning
Abstract: Imitation learning addresses the challenge of learning by observing an expert’s demonstrations without access to reward signals from the environment. Most existing imitation learning methods that do not require interacting with the environment either model the expert distribution as the conditional probability p(a|s) (e.g., behavioral cloning, BC) or the joint probability p(s, a) (e.g., implicit behavioral cloning). Despite its simplicity, modeling the conditional probability with BC usually struggles with generalization. While modeling the joint probability can lead to improved generalization performance, the inference procedure can be time-consuming and it often suffers from manifold overfitting. This work proposes an imitation learning framework that benefits from modeling both the conditional and joint probability of the expert distribution. Our proposed diffusion model-augmented behavioral cloning (DBC) employs a diffusion model trained to model expert behaviors and learns a policy to optimize both the BC loss (conditional) and our proposed diffusion model loss (joint). DBC outperforms baselines in various continuous control tasks in navigation, robot arm manipulation, dexterous manipulation, and locomotion. We design additional experiments to verify the limitations of modeling either the conditional probability or the joint probability of the expert distribution as well as compare different generative models.

Abstract

Imitation learning addresses the challenge of learning by observing an expert's demonstrations without access to reward signals from environments. Most existing imitation learning methods that do not require interacting with environments either model the expert distribution as the conditional probability p ( a | s ) ( e . g ., behavioral cloning, BC) or the joint probability p ( s, a ) ( e . g ., implicit behavioral cloning). Despite its simplicity, modeling the conditional probability with BC usually struggles with generalization. While modeling the joint probability can lead to improved generalization performance, the inference procedure can be time-consuming and it often suffers from manifold overfitting. This work proposes an imitation learning framework that benefits from modeling both the conditional and joint probability of the expert distribution. Our proposed diffusion model-augmented behavioral cloning (DBC) employs a diffusion model trained to model expert behaviors and learns a policy to optimize both the BC loss (conditional) and our proposed diffusion model loss (joint). DBC outperforms baselines in various continuous control tasks in navigation, robot arm manipulation, dexterous manipulation, and locomotion. We design additional experiments to verify the limitations of modeling either the conditional probability or the joint probability of the expert distribution as well as compare different generative models.

1 Introduction 19

Recently, the success of deep reinforcement learning (DRL) [Mnih et al., 2015, Lillicrap et al., 2016, 20 Arulkumaran et al., 2017] has inspired the research community to develop DRL frameworks to 21 control robots, aiming to automate the process of designing sensing, planning, and control algorithms 22 by letting the robot learn in an end-to-end fashion. Yet, acquiring complex skills through trial and 23 error can still lead to undesired behaviors even with sophisticated reward design [Christiano et al., 24 2017, Leike et al., 2018, Lee et al., 2019]. Moreover, the exploring process could damage expensive 25 robotic platforms or even be dangerous to humans [Garcıa and Fernández, 2015, Levine et al., 2020]. 26

27

28

29

30

31

32

33

To overcome this issue, imitation learning ( i . e ., learning from demonstration) [Schaal, 1997, Osa et al., 2018] has received growing attention, whose aim is to learn a policy from expert demonstrations, which are often more accessible than appropriate reward functions for reinforcement learning. Among various imitation learning directions, adversarial imitation learning [Ho and Ermon, 2016, Zolna et al., 2021, Kostrikov et al., 2019] and inverse reinforcement learning [Ng and Russell, 2000, Abbeel and Ng, 2004] have achieved encouraging results in a variety of domains. Yet, these methods require interacting with environments, which can still be expensive or unsafe.

On the other hand, behavioral cloning (BC) [Pomerleau, 1989, Bain and Sammut, 1995] does not 34

require interacting with environments. BC formulates imitation learning as a supervised learning 35

problem - given an expert demonstration dataset, an agent policy takes states sampled from the 36

dataset as input and learns to replicate the corresponding expert actions. One can view a BC policy as 37

a discriminative model p ( a | s ) that models the conditional probability of an action a given a state s . 38 Due to its simplicity and training stability, BC has been widely adopted for various applications. 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

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

However, BC struggles at generalizing to states unobserved during training [Nguyen et al., 2023]. To address this issue, implicit behavioral cloning (IBC) [Florence et al., 2022] aims to model the joint probability of the expert state-action pairs p ( s, a ) with energy-based models. IBC demonstrates superior performance when generalization is required. Yet, imitation learning methods in a similar vein [Ganapathi et al., 2022] that model the joint probability of state-action pairs p ( s, a ) instead of directly predicting actions p ( a | s ) require time-consuming actions sampling and optimization to retrieve a desired action arg max p ( s, a ) during inference despite the choice of models.

∈A

a

This work proposes an imitation learning framework that combines both the efficiency of modeling the conditional probability and the generalization ability of modeling the joint probability . Specifically, we propose to model the expert state-action pairs using a state-of-the-art generative model, a diffusion model, which learns to estimate how likely a state-action pair is sampled from the expert dataset. Then, we train a policy to optimize both the BC objective and the estimate produced by the learned diffusion model. Therefore, our proposed framework not only can efficiently predict actions given states via capturing the conditional probability p ( a | s ) but also enjoys the generalization ability induced by modeling the joint probability p ( s, a ) and utilizing it to guide policy learning.

We evaluate our proposed framework and baselines in various continuous control domains, including navigation, robot arm manipulation, and locomotion. The experimental results show that the proposed framework outperforms all the baselines or achieves competitive performance on all tasks. Extensive ablation studies compare our proposed method to its variants, justifying our design choices, such as different generative models, and investigating the effect of hyperparameters.

2 Related Work

Imitation learning addresses the challenge of learning by observing expert demonstrations without access to reward signals from environments. It has various applications such as robotics [Schaal, 1997], autonomous driving [Ly and Akhloufi, 2020], and game AI [Harmer et al., 2018].

Behavioral Cloning (BC). BC [Pomerleau, 1989, Torabi et al., 2018] formulate imitating an expert as a supervised learning problem. Due to its simplicity and effectiveness, it has been widely adopted in various domains. Yet, it often struggles at generalizing to states unobserved from the expert demonstrations [Ross et al., 2011, Florence et al., 2022]. In this work, we augment BC by employing a diffusion model that learns to capture the joint probability of expert state-action pairs.

Adversarial Imitation Learning (AIL). AIL methods aim to match the state-action distributions of an agent and an expert via adversarial training. Generative adversarial imitation learning (GAIL) [Ho and Ermon, 2016] and its extensions [Torabi et al., 2019, Kostrikov et al., 2019, Zolna et al., 2021] resemble the idea of generative adversarial networks [Goodfellow et al., 2014], which trains a generator policy to imitate expert behaviors and a discriminator to distinguish between the expert and the learner's state-action pair distributions. While modeling state-action distributions often leads to satisfactory performance, adversarial learning can be unstable and inefficient [Chen et al., 2020]. Moreover, AIL methods require online interaction with environments, which can be costly or even dangerous. In contrast, our work does not require interacting with environments.

Inverse Reinforcement Learning (IRL). IRL methods [Ng and Russell, 2000, Abbeel and Ng, 2004, Fu et al., 2018, Lee et al., 2021] are designed to infer the reward function that underlies the expert demonstrations and then learn a policy using the inferred reward function. This allows for learning tasks whose reward functions are difficult to specify manually. However, due to its doubleloop learning procedure, IRL methods are typically computationally expensive and time-consuming. Additionally, obtaining accurate estimates of the expert's reward function can be difficult, especially when the expert's behavior is non-deterministic or when the expert's demonstrations are sub-optimal.

Diffusion Policies. Recently, Pearce et al. [2023], Chi et al. [2023], Reuss et al. [2023] propose to represent and learn an imitation learning policy using a conditional diffusion model, which produces a predicted action conditioning on a state and a sampled noise vector. These methods achieve encouraging results in modeling stochastic and multimodal behaviors from human experts or play

data. In contrast, instead of representing a policy using a diffusion model, our work employs a 89 diffusion model trained on expert demonstrations to guide a policy as a learning objective. 90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

3 Preliminaries

3.1 Imitation Learning

Without loss of generality, the reinforcement learning problem can be formulated as a Markov decision process (MDP), which can be represented by a tuple M = ( S, A, R, P, ρ, γ ) with states S , actions

A , reward function R ( S, A ) ∈ (0 , 1) , transition distribution P ( s ′ | s, a ) : S × A × S → [0 , 1] , initial state distribution ρ , and discounted factor γ . Based on the rewards received while interacting with the environment, the goal is to learn a policy π ( ·| s ) to maximize the expectation of the cumulative discounted return ( i . e ., value function): V ( π ) = E [ T ∑ t =0 γ t R ( s t , a t ) | s 0 ∼ ρ ( · ) , a t ∼ π ( ·| s t ) , s t +1 ∼ P ( s t +1 | s t , a t )] , where T denotes the episode length. Instead of interacting with the environment and receiving rewards, imitation learning aims to learn an agent policy from an expert demonstration dataset, containing M trajectories, D = { τ 1 , ..., τ M } , where τ i represents a sequence of n i stateaction pairs { s i 1 , a i 1 , ..., s i n i , a i n i } .

3.2 Behavioral Cloning: Modeling Conditional Probability p ( a | s )

To learn a policy π , behavioral cloning (BC) directly estimates the expert policy π E with maximum likelihood estimation (MLE). Given a state-action pair ( s, a ) sampled from the dataset D , BC optimizes max θ ∑ ( s,a ) ∈ D log( π θ ( a | s )) , where θ denotes the parameters of the policy π . One can view a

BC policy as a discriminative model p ( a | s ) , capturing the conditional probability of an action a given a state s . Despite its success in various applications, BC tends to overfit and struggle at generalizing to states unseen during training [Ross et al., 2011, Codevilla et al., 2019, Wang et al., 2022].

3.3 Modeling Joint Probability p ( s, a )

Aiming for improved generalization ability, implicit behavioral cloning [Florence et al., 2022] and methods in a similar vein [Ganapathi et al., 2022] model the joint probability p ( s, a ) of expert stateaction pairs. These methods demonstrate superior generalization performance in diverse domains. Yet, without directly modeling the conditional probability p ( a | s ) , the action sampling and optimization procedure to retrieve a desired action arg max a ∈A p ( s, a ) during inference is often time-consuming.

Moreover, explicit generative models such as energy-based models [Du and Mordatch, 2019, Song and Kingma, 2021], variational autoencoder [Kingma and Welling, 2014], and flow-based models Rezende and Mohamed [2015], Dinh et al. [2017] are known to struggle with modeling observed highdimensional data that lies on a low-dimensional manifold ( i . e ., manifold overfitting) [Wu et al., 2021, Loaiza-Ganem et al., 2022]. As a result, these methods often perform poorly when learning from demonstrations produced by script policies or PID controllers, as discussed in Section 5.4.

We aim to develop an imitation learning framework that enjoys the advantages of modeling the conditional probability p ( a | s ) and the joint probability p ( s, a ) . Specifically, we propose to model the joint probability of expert state-action pairs using an explicit generative model ϕ , which learns to produce an estimate indicating how likely a state-action pair is sampled from the expert dataset. Then, we train a policy to model the conditional probability p ( a | s ) by optimizing the BC objective and the estimate produced by the learned generative model ϕ . Hence, our method can efficiently predict actions given states, generalize better to unseen states, and suffer less from manifold overfitting.

3.4 Diffusion Models

As described in the previous sections, this work aims to combine the advantages of modeling both the conditional probability p ( a | s ) and the joint probability p ( s, a ) . To this end, we leverage diffusion models to model the joint probability of expert state-action pairs. The diffusion model is a recently developed class of generative models and has achieved state-of-the-art performance on various tasks Sohl-Dickstein et al. [2015], Nichol and Dhariwal [2021], Dhariwal and Nichol [2021].

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

In this work, we utilize Denoising Diffusion Probabilistic Models (DDPMs) J Ho [2020] to model expert state-action pairs. Specifically, DDPM models gradually add noise to data samples ( i . e ., concatenated state-action pairs) until they become isotropic Gaussian ( forward diffusion process ), and then learn to denoise each step and restore the original data samples ( reverse diffusion process ), as illustrated in Figure 1. In other words, DDPM learns to recognize a data distribution by learning to denoise noisy sampled data. More discussion on diffusion models can be found in the Section G.

4 Approach

Our goal is to design an imitation learning framework that enjoys both the advantages of modeling the conditional probability and the joint probability of expert behaviors. To this end, we first adopt behavioral cloning (BC) for modeling the conditional probability from expert state-action pairs, as described in Section 4.1. To capture the joint probability of expert state-action pairs, we employ a diffusion model which learns to produce an estimate indicating how likely a state-action pair is sampled from the expert state-action pair distribution, as presented in Section 4.2.1. Then, we propose to guide the policy learning by optimizing this estimate provided by a learned diffusion model, encouraging the policy to produce actions similar to expert actions, as discussed in Section 4.2.2. Finally, in Section 4.3, we introduce the framework that combines the BC loss and our proposed diffusion model loss, allowing for learning a policy that benefits from modeling both the conditional probability and the joint probability of expert behaviors. An overview of our proposed framework is illustrated in Figure 2, and the algorithm is detailed in Section B.

4.1 Behavioral Cloning Loss

The behavioral cloning (BC) model aims to imitate expert behaviors with supervision learning. BC learns to capture the conditional probability p ( a | s ) of expert state-action pairs. Given a sampled expert state-action pair ( s, a ) , a policy π learns to predict an action ˆ a ∼ π ( s ) by optimizing

where d ( · , · ) denotes a distance measure between a pair of actions. For example, we can adapt the mean-square error (MSE) loss || a -ˆ a || 2 for most continuous control tasks.

4.2 Learning a Diffusion Model and Guiding Policy Learning

Instead of directly learning the conditional probability p ( a | s ) , this section discusses how to model the joint probability p ( s, a ) of expert behaviors with a diffusion model in Section 4.2.1 and presents how to leverage the learned diffusion model to guide policy learning in Section 4.2.2.

4.2.1 Learning a Diffusion Model

We propose to model the joint probability of expert state-action pairs with a diffusion model ϕ . Specifically, we create a joint distribution by simply concatenating a state vector s and an action vector a from a state-action pair ( s, a ) . To model such distribution by learning a denoising diffusion probabilistic model (DDPM) J Ho [2020], we inject noise ϵ ( n ) into sampled state-action pairs, where n indicates the number of steps of the Markov procedure, which can be viewed as a variable of the level of noise. Then, we train the diffusion model ϕ to predict the injected noises by optimizing

where ˆ ϵ is the noise predicted by the diffusion model ϕ . Once optimized, the diffusion model can recognize the expert distribution by perfectly predicting the noise injected into state-action pairs sampled from the expert distribution. On the other hand, predicting the noise injected into stateaction pairs sampled from any other distribution should yield a higher loss value. Therefore, we

q

(

x

n

|

x

n

-1

)

ϕ

(

x

n

-1

|

x

n

)

x

n

-1

x

n

x

N

x

0

Forward di

ff

usion process

Reverse di ff

usion process

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

̂

(a) Learning a Diffusion Model

(b) Learning a Policy with the Learned Diffusion Model

s

a

/uni03F5

Di

ff

usion

Model

ϕ

/uni03F5

Expert's Demonstrations

D

τ

1

s

1

1

a

1

1

s

1

2

a

1

2

s

1

n

1

a

1

n

1

τ

2

s

2

1

a

2

1

s

2

2

a

2

2

s

2

n

2

a

2

n

2

τM

s M

1

a M

1

s M

2

a

3

2

s M

nM

a M

nM

Sample

ℒ

diff

Learning objective

Learnable mapping

Frozen mapping

s

a

/uni03F5

Di

ff

usion

Model

ϕ

s

a

/uni03F5

ℒ

agent

diff

ℒ

expert

diff

ℒ

DM

Policy

π

a

s

a

ℒ

BC

Agent's

action

Expert's

action

State

̂

propose to view L diff ( s, a, ϕ ) as an estimate of how well the state-action pair ( s, a ) fits the state-action distribution that ϕ learns from.

4.2.2 Learning a Policy with Diffusion Model Loss

Adiffusion model ϕ trained on the expert distribution can produce an estimate L diff ( s, a, ϕ ) indicating how well a state-action pair ( s, a ) fits the expert distribution. We propose to leverage this signal to guide a policy to imitate the expert. Specifically, given a state-action ( s, a ) sampled from D , the π predicts an action given the state ˆ a ∼ π ( s ) by optimizing

Intuitively, the policy learns to predict actions that are indistinguishable from the expert actions for the diffusion model conditioning on the same set of states.

We hypothesize that learning a policy to optimize Eq. 3 can be unstable, especially for state-action pairs that are not well-modeled by the diffusion model, which yield a high value of L diff even with expert state-action pairs. Therefore, we propose to normalize the agent diffusion loss L agent diff with an expert diffusion loss L expert diff , which can be computed with expert state-action pairs ( s, a ) as follows:

We propose to optimize the diffusion model loss L DM based on calculating the difference between the above agent and expert diffusion losses:

4.3 Combining the Two Objectives

Our goal is to learn a policy that benefits from both modeling the conditional probability and the joint probability of expert behaviors. To this end, we propose to augment a BC policy that optimizes the BC loss L BC in Eq. 1 by jointing optimizing the proposed diffusion model loss L DM in Eq. 5, which encourages the policy to predict actions that fit the expert joint probability captured by a diffusion model. To learn from both the BC loss and the diffusion model loss, we train the policy to optimize

where λ is a coefficient that determines the importance of the diffusion model loss relative to the BC loss. We analyze the effect of the coefficient in Section 5.6.1.

̂

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

(a) MAZE

(b) FETCHPICK

(c) FETCHPUSH

(d) HANDROTATE

(e) WALKER

5 Experiments

We design experiments in various continuous control domains, including navigation, robot arm manipulation, dexterous manipulation, and locomotion, to compare our proposed framework (DBC) to its variants and baselines.

5.1 Experimental Setup

This section describes the environments, tasks, and expert demonstrations used for learning and evaluation. More details can be found in Section A.

Navigation. To evaluate our method on a navigation task, we choose MAZE, a maze environment proposed in Fu et al. [2020] (maze2d-medium-v2), as illustrated in Figure 3a. This task features a point-mass agent in a 2D maze learning to navigate from its start location to a goal location by iteratively predicting its x and y acceleration. The agent's beginning and final locations are chosen randomly. We collect 100 demonstrations with 18,525 transitions using a controller.

Robot Arm Manipulation. We evaluate our method in a robot arm manipulation domain with two 7-DoF Fetch tasks: FETCHPICK and FETCHPUSH, as illustrated in Figure 3c and Figure 3b. FETCHPICK requires picking up an object from the table and lifting it to a target location; FETCHPUSH requires the arm to push an object to a target location. We use the demonstrations provided in Lee et al. [2021] for these tasks. Each dataset contains 10k transitions (303 trajectories for FETCHPICK and 185 trajectories for FETCHPUSH).

Dexterous Manipulation. In HANDROTATE, we further evaluate our method on a challenging environment proposed in Plappert et al. [2018], where a 24-DoF Shadow Dexterous Hand learns to in-hand rotate a block to a target orientation, as illustrated in Figure 3d. This environment has a high-dimensional state space (68D) and action space (20D). We collected 10k transitions (515 trajectories) from a SAC [Haarnoja et al., 2018] expert policy trained for 10M environment steps.

Locomotion. For locomotion, we leverage the WALKER environment Brockman et al. [2016], which requires a bipedal agent to walk as fast as possible while maintaining its balance, as illustrated in Figure 3e. We use the demonstrations provided by Kostrikov [2018], which contains 5 trajectories with 5k state-action pairs.

5.2 Baselines

We compare our method DBC with the following baselines.

BC learns to imitate an expert by modeling the conditional probability p ( a | s ) of the expert behaviors via optimizing the BC loss L BC in Eq. 1.

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

265

266

267

268

269

270

271

272

273

274

275

276

277

278

Implicit BC (IBC) [Florence et al., 2022] models expert state-action pairs with an energy-based model. For inference, we implement the derivative-free optimization algorithm proposed in IBC, which samples actions iteratively to select the desired action with the minimum predicted energy. This baseline serves a representative of the methods that solely model the joint probability p ( s, a ) of the expert behaviors.

Diffusion policy refers to the methods that learn a conditional diffusion model as a policy [Chi et al., 2023, Reuss et al., 2023]. Specifically, we implement this baseline based on Pearce et al. [2023]. We include this baseline to analyze the effectiveness of using diffusion models as a policy or as a learning objective (ours).

5.3 Experimental Results

We report the experimental results in terms of success rate (MAZE, FETCHPICK, FETCHPUSH, HANDROTATE), and return (WALKER) in Table 1. The details of model architecture can be found in Section C. Training and evaluation details can be found in Section D. Additional analysis and experimental results can be found in Section E and Section F.

Overall Task Performance. Our proposed method DBC achieves the highest success rates, outperforming our baselines in all the goal-directed tasks (MAZE, FETCHPICK, FETCHPUSH, and HANDROTATE) and perform competitively in WALKER compared to the best-performing baseline (BC). We hypothesize the improvement in the goal-directed tasks can be mostly attributed to the better generalization ability since starting positions and the goals are randomized during evaluation and therefore requires the policy to deal with unseen situation. To verify this hypothesis, we further evaluate the baselines and our method in FETCHPICK and FETCHPUSH with different levels of randomization in Section E.

Locomotion. Unlike the goal-directed tasks, we do not observe significant improvement but competitive results from DBC compared to the best-performing baseline (BC). We hypothesize that this is because locomotion tasks such as WALKER, with sufficient expert demonstrations and little randomness, do not require generalization during inference. The agent can simply follow the closed-loop progress of the expert demonstrations, resulting in both BC (7066.61) and DBC (7057.42) performing similarly to the expert with an average return of 7063.72. On the other hand, we hypothesize that Diffusion Policy performs slightly worse due to its design for modeling multimodal behaviors, which is contradictory to learning from this single-mode simulated locomotion task.

Action Space Dimension. While Implicit BC models the joint distribution and generalizes better, it requires time-consuming actions sampling and optimization during inference. Moreover, such procedure may not scale well to high-dimensional action spaces. Our Implicit BC baseline with a derivative-free optimizer struggles in HANDROTATE and WALKER environments, whose action dimensions are 20 and 6, respectively. This is consistent with Florence et al. [2022], which reports that the optimizer failed to solve tasks with an action dimension larger than 5. In contrast, our proposed DBC can handle high-dimensional action spaces.

Inference Efficiency. To evaluate the inference efficiency, we measure and report the number of evaluation episodes per second ( ↑ ) for IBC (9.92), Diffusion Policy (1.38), and DBC ( 30.79 ) on an NVIDIA RTX 3080 Ti GPU in MAZE. This can be attributed to the fact that DBC and BC model the conditional probability p ( a | s ) and can directly map states to actions during inference. In contrast, Implicit BC requires action sampling and optimization, while Diffusion Policy is required to iteratively denoise sampled noises. This verifies the efficiency of modeling the conditional probability.

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

(a)

(b)

(c)

Train Goal Region

Test Goal Region

Test Goal

BC

πDM

×

Expert Action

Predicted Action

Texpert

πBC

TDM

5.4 Comparing Modeling Conditional Probability and Joint Probability

This section aims to empirically identify the limitations of modeling either the conditional or the joint probability in an open maze environment implemented with [Fu et al., 2020].

Generalization. We aim to investigate if learning from the BC loss alone struggles at generalization ( conditional ) and examine if guiding the policy using the diffusion model loss yields improved generalization ability ( joint ). We collect trajectories of a PPO policy learning to navigate from (5 , 3) to goals sampled around (1 , 2) and (1 , 4) (green), as shown in Figure 4a. Given these expert trajectories, we learn a policy π BC to optimize Eq. 1 and another policy π DM to optimize Eq. 5. Then, we evaluate the two policies by sampling goals around (1 , 1) , (1 , 3) , and (1 , 5) (red), which requires the ability to generalize. Visualized trajectories of the two policies in Figure 4a show that π BC (orange) fails to generalize to unseen goals, whereas π DM (blue) can generalize ( i . e ., extrapolate) to some extent. This verifies our motivation to augment BC with the diffusion model loss.

Manifold overfitting. We aim to examine if modeling the joint probability is difficult when observed high-dimensional data lies on a low-dimensional manifold ( i . e ., manifold overfitting). We collect trajectories from a script policy that executes actions (0 . 5 , 0) , (0 , 0 . 5) , ( -0 . 7 , 0) , and (0 , -0 . 7) (red crosses in Figure 4b), each for 40 consecutive time steps, resulting the green spiral trajectories visualized in Figure 4c.

Given these expert demonstrations, we learn a policy π BC to optimize Eq. 1, and another policy π DM to optimize Eq. 5 with a diffusion model trained on the expert distribution. Figure 4b shows that the diffusion model struggles at modeling such expert action distribution with a lower intrinsic dimension. As a result, Figure 4c show that the trajectories of π DM (blue) drastically deviates from the expert trajectories (green) as the diffusion model cannot provide effective loss. On the other hand, the trajectories of π BC (orange) is able to closely follow expert's. This verifies our motivation to complement modeling the joint probability with modeling the conditional probability ( i . e ., BC).

5.5 Comparing Different Generative Models

Our proposed framework employs a diffusion model (DM) to model the joint probability of expert state-action pairs and utilizes it to guide policy learning. To justify our choice, we explore using other popular generative models to replace the diffusion model in MAZE. We consider energy-based models (EBMs) [Du and Mordatch, 2019, Song and Kingma, 2021], variational autoencoder (VAEs) [Kingma and Welling, 2014], and generative adversarial networks (GANs) Goodfellow et al. [2014]. Each

309

310

311

312

313

314

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

generative model learns to model expert state-action pairs. To guide policy learning, given a predicted state-action pair ( s, ˆ a ) we use the estimated energy of an EBM, the reconstruction error of a VAE, and the discriminator output of a GAN to optimize a policy with or without the BC loss. Training details can be found in Section D.3.

Table 2 compares using different generative models to model the expert distribution and guide policy learning. All the generative model-guide policies can be improved by adding the BC loss, justifying our motivation to complement modeling the joint probability with modeling the conditional probability. With or without the BC loss, the diffusion model-guided policy achieves the best performance compared to other generative models, verifying our choice of the generative model.

5.6 Ablation Study

In this section, we investigate the effect of the diffusion model loss coefficient λ (Section 5.6.1) and examine the effect of the normalization term L expert diff in the diffusion model loss L DM (Section 5.6.2).

5.6.1 Effect of the Diffusion Model Loss Coefficient λ

We examine the impact of varying the coefficient of the diffusion model loss λ in Eq. 6 in MAZE. The result presented in Table 3 shows that λ = 5 yields the best performance. A higher or lower λ leads to worse performance, demonstrating how modeling the conditional probability ( L BC) and the joint probability ( L DM) can complement each other.

5.6.2 Effect of the Normalization Term L expert diff

We aim to investigate whether normalizing the diffusion model loss L DM with the expert diffusion model loss L expert diff yields improved performance in MAZE. We train a variant of DBC where only L agent diff in Eq. 3 instead of L DM in Eq. 5 is used to augment BC. This variant learning from an unnormalized diffusion model loss achieves an average success rate of 80 . 20% , worse than the full DBC ( 86 . 99% ). This justifies the effectiveness of the proposed normalization term L expert diff in L DM.

6 Conclusion

We propose an imitation learning framework that benefits from modeling both the conditional probability p ( a | s ) and the joint probability p ( s, a ) of the expert distribution. Our proposed diffusion model-augmented behavioral cloning (DBC) employs a diffusion model trained to model expert behaviors and learns a policy to optimize both the BC loss and our proposed diffusion model loss. Specifically, the BC loss captures the conditional probability p ( a | s ) from expert state-action pairs, which directly guides the policy to replicate the expert's action. On the other hand, the diffusion model loss models the joint distribution of expert's state-action pairs p ( s, a ) , which provides an evaluation of how well the predicted action aligned with the expert distribution. DBC outperforms baselines or achieves competitive performance in various continuous control tasks in navigation, robot arm manipulation, dexterous manipulation, and locomotion. We design additional experiments to verify the limitations of modeling either the conditional probability or the joint probability of the expert distribution as well as compare different generative models. Ablation studies investigate the effect of hyperparameters and justify the effectiveness of our design choices.

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

References

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 , 2015.

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. In International Conference on Learning Representations , 2016.

Kai Arulkumaran, Marc Peter Deisenroth, Miles Brundage, and Anil Anthony Bharath. Deep reinforcement learning: A brief survey. IEEE Signal Processing Magazine , 2017.

Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. In Advances in Neural Information Processing Systems , 2017.

Jan Leike, David Krueger, Tom Everitt, Miljan Martic, Vishal Maini, and Shane Legg. Scalable agent alignment via reward modeling: a research direction. arXiv preprint arXiv:1811.07871 , 2018.

Youngwoon Lee, Shao-Hua Sun, Sriram Somasundaram, Edward S. Hu, and Joseph J. Lim. Composing complex skills by learning transition policies. In Proceedings of International Conference on Learning Representations , 2019.

Javier Garcıa and Fernando Fernández. A comprehensive survey on safe reinforcement learning. Journal of Machine Learning Research , 2015.

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

Stefan Schaal. Learning from demonstration. In Advances in Neural Information Processing Systems , 1997.

Takayuki Osa, Joni Pajarinen, Gerhard Neumann, J Andrew Bagnell, Pieter Abbeel, Jan Peters, et al. An algorithmic perspective on imitation learning. Foundations and Trends® in Robotics , 2018.

Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances in Neural Information Processing Systems , 2016.

Konrad Zolna, Scott Reed, Alexander Novikov, Sergio Gomez Colmenarejo, David Budden, Serkan Cabi, Misha Denil, Nando de Freitas, and Ziyu Wang. Task-relevant adversarial imitation learning. In Conference on Robot Learning , 2021.

Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson. Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning. In International Conference on Learning Representations , 2019.

Andrew Y. Ng and Stuart J. Russell. Algorithms for inverse reinforcement learning. In International Conference on Machine Learning , 2000.

Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In International Conference on Machine Learning , 2004.

Dean A Pomerleau. Alvinn: An autonomous land vehicle in a neural network. In Advances in Neural Information Processing Systems , 1989.

Michael Bain and Claude Sammut. A framework for behavioural cloning. In Machine Intelligence 15 , 1995.

Tung Nguyen, Qinqing Zheng, and Aditya Grover. Reliable conditioning of behavioral cloning for offline reinforcement learning. arXiv preprint arXiv:2210.05158 , 2023.

Pete Florence, Corey Lynch, Andy Zeng, Oscar A Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Conference on Robot Learning , 2022.

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

Aditya Ganapathi, Pete Florence, Jake Varley, Kaylee Burns, Ken Goldberg, and Andy Zeng. Implicit kinematic policies: Unifying joint and cartesian action spaces in end-to-end robot learning. In International Conference on Robotics and Automation , 2022.

Abdoulaye O Ly and Moulay Akhloufi. Learning to drive by imitation: An overview of deep behavior cloning methods. IEEE Transactions on Intelligent Vehicles , 2020.

Jack Harmer, Linus Gisslén, Jorge del Val, Henrik Holst, Joakim Bergdahl, Tom Olsson, Kristoffer Sjöö, and Magnus Nordin. Imitation learning with concurrent actions in 3d games. In IEEE Conference on Computational Intelligence and Games , 2018.

Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. In International Joint Conference on Artificial Intelligence , 2018.

Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In International Conference on Artificial Intelligence and Statistics , 2011.

Faraz Torabi, Garrett Warnell, and Peter Stone. Generative adversarial imitation from observation. ICML , 2019.

Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in Neural Information Processing Systems , 2014.

Minshuo Chen, Yizhou Wang, Tianyi Liu, Zhuoran Yang, Xingguo Li, Zhaoran Wang, and Tuo Zhao. On computation and generalization of generative adversarial imitation learning. In International Conference on Learning Representations , 2020.

Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adverserial inverse reinforcement learning. In International Conference on Learning Representations , 2018.

Youngwoon Lee, Andrew Szot, Shao-Hua Sun, and Joseph J. Lim. Generalizable imitation learning from observation via inferring goal proximity. In Neural Information Processing Systems , 2021.

Tim Pearce, Tabish Rashid, Anssi Kanervisto, David Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam Devlin. Imitating human behaviour with diffusion models. In International Conference on Learning Representations , 2023.

Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. In Robotics: Science and Systems , 2023.

Moritz Reuss, Maximilian Li, Xiaogang Jia, and Rudolf Lioutikov. Goal-conditioned imitation learning using score-based diffusion policies. arXiv preprint arXiv:2304.02532 , 2023.

Felipe Codevilla, Eder Santana, Antonio M López, and Adrien Gaidon. Exploring the limitations of behavior cloning for autonomous driving. In International Conference on Computer Vision , 2019.

Lingguang Wang, Carlos Fernandez, and Christoph Stiller. High-level decision making for automated highway driving via behavior cloning. IEEE Transactions on Intelligent Vehicles , 2022.

Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. In Neural Information Processing Systems , 2019.

Yang Song and Diederik P Kingma. How to train your energy-based models. arXiv preprint arXiv:2101.03288 , 2021.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representations , 2014.

Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International Conference on Machine Learning , 2015.

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

465

466

Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In International Conference on Learning Representations , 2017.

Qitian Wu, Rui Gao, and Hongyuan Zha. Bridging explicit and implicit deep generative models via neural stein estimators. In Neural Information Processing Systems , 2021.

Gabriel Loaiza-Ganem, Brendan Leigh Ross, Jesse C Cresswell, and Anthony L Caterini. Diagnosing and fixing manifold overfitting in deep generative models. Transactions on Machine Learning Research , 2022.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning . PMLR, 2015.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning , 2021.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Neural Information Processing Systems , 2021.

AJain J Ho. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems , 2020.

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.

Matthias Plappert, Marcin Andrychowicz, Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, et al. Multi-goal reinforcement learning: Challenging robotics environments and request for research. arXiv preprint arXiv:1802.09464 , 2018.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of International Conference on Machine Learning (ICML) , 2018.

Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.

Ilya Kostrikov. Pytorch implementations of reinforcement learning algorithms. https://github. com/ikostrikov/pytorch-a2c-ppo-acktr-gail , 2018.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new algorithm for behavior cloning (BC) where the BC learning objective is modified with a diffusion modeling loss that models the joint state-action distribution of the expert data. The paper demonstrates the benefits of modeling both the conditional probability and joint probability of the expert distribution. The results of the proposed method have been compared with baselines modeling both conditional probability and joint probability distributions and have been supported with appropriate ablation studies.\n\nSTRENGTHS:\n- The paper is well-written and easy to understand. \n- The problem has been well motivated. The authors claim that while solely modeling the conditional probability p(a|s) of the expert data struggle to generalize to unseen states during training, solely modeling the joint probability can improve generalization to unseen states at the cost of high inference time. Further, though generative models provide encouraging results in modeling stochastic and multimodal behaviors, they struggle with manifold overfitting (shown in Fig. 4(c)). Accordingly, the paper proposes a loss function that combines modeling the conditional and joint probabilities of the expert data while using a state-of-the-art generative model (diffusion model) to benefit from their superior behavior modeling capabilities (as shown in [1] which has also been cited in the paper).\n- An interesting feature of the proposed method is that since it only uses a diffusion model to guide the learning of conditional probability, it is able to circumvent the issue of high inference time associated with the noising and denoising procedure in diffusion models (mentioned in Sec. 5.3).\n- The primary results have been compared with baselines modeling both conditional probability and joint probability distributions. Table 1 shows that DBC outperforms BC (conditional probability), Implicit BC (joint probability), and Diffusion policy (a state-of-the-art generative model) on all tasks except the Walker. However, since DBC performs comparably to the best performing baseline (BC) (the authors provide some justification for this in Sec. 5.3 Locomotion).\n- The paper provides interesting insights on modeling high dimensional action spaces, inference efficiency, generalization to unseen goals, and manifold overfitting. Additional ablation studies to justify design choices such as the choice of the generative model, loss coefficient value, and effect of the normalization term have also been provided.\n\n[1] Chi, Cheng, et al. \"Diffusion policy: Visuomotor policy learning via action diffusion.\" arXiv preprint arXiv:2303.04137 (2023).\n\nWEAKNESSES:\n- The limitations of the proposed approach can only be found in the limitations. Though I believe that is fine given space limitations, it would be nice to mention it briefly in the main paper for completeness.\n- Since the paper focuses on imitation learning from expert data, it would be interesting to also add *Action Chunking with Transformers* (ACT) [1] as a baseline. ACT also uses a generative model (a conditional variational autoencoder) as a policy and given the reported results, it would be interesting to see how DBC compares with ACT. I understand that the paper does compare with a VAE. However, ACT with its transformer encoder and decoder would be an interesting baseline to compare with.\n- I am curious about how the performance of the proposed method scales with the dataset size. Specifically, since diffusion models seem to be “data hungry” (based on results from the computer vision and NLP community), it would be interesting how this method scales with dataset size. Since this method solely uses a diffusion model to guide the learning of policy, I believe a study of differences in the performance of DBC and the baselines with different amounts of training data might give interesting insights.\n\n[1] Zhao, Tony Z., et al. \"Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware.\" arXiv preprint arXiv:2304.13705 (2023).\n\nQUESTIONS:\nIt would be great if the authors could address the points mentioned in the “Weaknesses” section.\n\nLIMITATIONS:\nYes. However, the limitations and societal impact of the work can only be found in the appendix. It would be nice to either have a brief mention or a reference to the appendix section in the main paper for completeness.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe submission proposes an imitation learning method optimizing a loss that is a weighted sum of a behavioral cloning loss and a loss based on a diffusion model. The diffusion model loss penalizes the policy for generating actions that are unlikely under the diffusion model, which is pre-trained to maximize the likelihood of joint state-action pairs sampled from the data.\n\nSTRENGTHS:\nOriginality:\n\nI am not aware of any previous work combining a behavioral cloning loss with a diffusion model loss.\n\nQuality:\n\nI think there is one interesting \"key trick\" that seems to be required for the method to work, which is using the diffusion-model-loss incurred by the expert demonstration as a baseline. I can see how if the diffusion model is wrong somewhere, it could inappropriately penalize the policy for doing something reasonable. Only penalizing the loss in excess of the expert's loss mitigates the effect of errors in the diffusion model.\n\nThe paper's strongest point is probably the experimental results in section 5.3. A reasonably diverse set of tasks were evaluated, and the proposed method performs significantly better than the baselines in most cases. I also appreciated that the experiments were targeted towards testing specific hypotheses—e.g., the results in section 5.5, which validate the hypotheses that the BC loss is necessary, and that the diffusion model is the best choice of EBM for this purpose.\n\nSignificance:\n\nThe relative simplicity of the approach, along with positive experimental results in non-trivial tasks, could encourage more research into combining BC with EBM-type losses.\n\nWEAKNESSES:\nOriginality:\n\nAlthough the specific combination of BC and DDPMs is novel, there is some precedent for combining BC losses with other losses—for example [A], which combines a BC loss with the GAIL loss.\n\n[A] Jena et al. Augmenting GAIL with BC for sample efficient imitation learning. https://arxiv.org/pdf/2001.07798.pdf\n\nClarity:\n\nSome key aspects of the submission's presentation could be improved. In particular, the mathematical notation is lacking in key places. Section 4.2.2 is missing notation to describe where we are taking various expectations, and with respect to which distributions. For example, should eq. 5 read $E_{(s,a) \\sim D, a \\sim \\pi(s)} \\max(..., 0)$ or $\\max (E_{(s,a) \\sim D, a \\sim \\pi(s)} \\dots, 0)$? Are any expectations taken with respect to $s$ sampled from the DM? I believe the answer is no, but I had to check the algorithm description in the appendix to be sure. The answers to these questions should be obvious from the math, but this is unfortunately not the case.\n\nAnother point of confusion is in section 4.3, which mentions \"jointly optimizing the proposed diffusion model loss.\" I generally interpret this to mean simultaneously optimizing the parameters of all the models, but this is not the case. I again had to check the algorithm description in the appendix to verify that the method is strictly a two-phase method, where the DM is pretrained and fixed before optimizing eq. 6.\n\nQuality:\n\nOne of the paper's main weaknesses is confusion as to the motivation for the method and why the method seems to work. Some of the explanations offered by the paper include:\n1. Diffusion models / EBMs generalize better because they model the joint probability p(s,a) as opposed to the conditional probability p(a|s) (lines 52-54).\n2. Combining BC with an EBM helps because EBMs alone suffer from \"manifold overfitting.\" (lines 127-128)\n\nBoth of these explanations seem a bit tenuous. As for the first statement (EBMs generalize better because they model the joint probability), it is simultaneously vague, debatable, and it is unclear how that property would benefit the method, even if it were true. It is stated (line 113) that \"These methods demonstrate superior generalization performance in diverse domains,\" but this statement seems too vague to be useful. Even assuming that modeling the joint probability produces better generalization in some sense, how does that translate to better results in the proposed method? No plausible mechanism is given for this. \n\nFor example, I could imagine that if one were to sample joint state-action configurations from a good model, and then train behavioral cloning using those state-action pairs, then that could conceivably perform better than training BC on raw data, because sampling from generated states may serve as a form of data augmentation. However, this is not what the method does, according to the algorithm description in the appendix—states are sampled strictly from the data distribution.\n\nAs for the statement about manifold overfitting, it is true that this is theoretically an issue with EBMs—however, there is a trivial solution for this that works fairly well, which is to perturb the training examples with a small amount of random noise to break any manifold constraints in the data. So, it is not necessary to add a BC loss to EBMs in order to solve the manifold overfitting problem.\n\nThat said, I do believe it is plausible that combining the BC and DM losses helps, but not for any of the reasons above. A more plausible explanation for why the method works is that the loss optimizes (ignoring eq. 4)  an approximation to \"forward\" KL divergence plus \"reverse\" KL divergence—since the DM loss is a lower bound on the log-likelihood of the \"data,\" where the \"data\" in this case consists of samples from the model. The forward KL divergence encourages mode coverage, at the expense of occasionally producing bad samples, whereas the reverse KL has the opposite properties—it produces \"good\" samples, at the expense of losing modes. Therefore, these losses have complementary characteristics that benefit each other when combined. However, one has to be careful not to inadvertently penalize good samples in the reverse KL term, which is probably why adding the human baseline in equation 5, helps.\n\nNotably, imitation learning based on reverse KL divergence and more general f-divergences has been previously explored, e.g., in [B]. This would make one of the methods suggested in [B] a fair comparison in experiments.\n\nI noted a few other issues with the experiments. As stated earlier, I am suspicious of the idea that it is essential that the DM model the joint probability as opposed to the conditional probability—it seems plausible to me that the method would work just as well if the DM were trained as a conditional model (since the method doesn't use states sampled from the DM), so it would be interesting to test this experimentally. \n\nAs also mentioned previously, I believe the manifold problem can be trivially solved by adding a small amount of random noise to the training examples. Doing so may significantly affect the results in Table 2—i.e., adding noise would probably boost the performance of the \"without BC\" column. Similar comments apply to the \"manifold overfitting\" results in section 5.4. I think this is an important experiment to run, because it may weaken one of the core claims, which is that adding BC to EBMs helps because it helps address the manifold overfitting problem.\n\n[B] Ke, Liyiming, et al. \"Imitation learning as f-divergence minimization.\" Algorithmic Foundations of Robotics XIV: Proceedings of the Fourteenth Workshop on the Algorithmic Foundations of Robotics 14. Springer International Publishing, 2021.\n\nQUESTIONS:\nHow do the results change (esp. the manifold overfitting experiments) when the data is perturbed with a small amount of random noise?\n\nHave you tried comparing the method to BC with some sort of simple data augmentation?\n\nHave you considered comparing to another method that optimizes another type of f-divergence (as exemplified by [B])?\n\nLIMITATIONS:\nI didn't see much discussion in the way of limitations. What about the efficiency and stability of training?", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a method to augment a behavior cloning (BC) agent with additional diffusion loss. The goal is to leverage the conditional probability learned by the BC loss and the joint probability learned by the diffusion loss. The diffusion loss is calculated using the prediction error of a pre-trained diffusion model on both the agent's and expert's state-action pairs. Experimental results demonstrate that this method outperforms baseline methods, achieving improved performance.\n\nSTRENGTHS:\n1. Effective Experimental Results: The paper provides experimental results that showcase the superior performance of DBC compared to several baselines, including BC, Implicit BC, and Diffusion Policy.\n2. Leveraging the Advantage of Offline Learning: The proposed method operates purely in an offline setting, eliminating the need for interaction with the environment. This reduction in training complexity is an advantage.\n\nWEAKNESSES:\n1. Sensitivity to the $\\lambda$ Parameter: The paper introduces a $\\lambda$ parameter for BC loss and diffusion model loss balancing, but the analysis is limited to the maze environment. It is crucial to investigate and analyze the performance of the algorithm across various environments when the $\\lambda$ parameter is modified.\n2. Lack of Theoretical Analysis: The paper lacks a theoretical analysis of the proposed diffusion model loss. Providing a theoretical foundation for the diffusion loss would enhance the paper's scientific rigor.\n\nQUESTIONS:\n1. Mechanism of Knowledge Transfer: The paper does not sufficiently explain how the diffusion model loss leads to agent's improved generalizability to unseen states. While the modeling the joint distribution of state-action pair can extrapolate/interpolate actions better than original BC on unseen states [1], it remains unclear how the knowledge generated by modeling the joint distribution is transferred to the agent. Since DBC trains the agent with the diffusion model loss on  **expert states only**, it is essential to provide a theoretical proof or analysis to address this question. The authors should clarify the mechanism behind the transfer of knowledge from the diffusion model to the agent and explain how this transfer leads to improved performance on **unseen states**.\n2. Redundancy of Diffusion Model Loss: The paper describes the proposed diffusion model loss as a measure of how well a state-action pair $(s^e,\\hat{a})$ fits the expert state-action pair distribution, where $s^e$ is the state of $(s^e,a^e)$ sampled from the expert's state-action distribution and $\\hat{a}\\sim \\pi(s^e)$ is the action predicted by agent. Additionally, the BC loss $||\\hat{a}-a^e||^2$ is used to ensure that the agent's prediction $\\hat{a}$ is as close to $a^e$ as possible. Theoretically, the state-action pair $(s^e,a^e)$ can fully minimize the diffusion model loss since it already exists in the expert state-action pair distribution. Consequently, minimizing the BC loss should also minimize the diffusion model loss on all training data (which is the expert data), potentially rendering the diffusion model loss redundant. The authors should address this observation and clarify the necessity of the diffusion model loss despite its potential redundancy.\n3. Maximizing Joint Probability vs. Conditional Probability: The paper claims that the agent benefits from the diffusion model loss because it models the joint probability $p(s,a)$. However, maximizing the joint probability reduces to maximizing the conditional probability when the state $s$ is fixed, since the states in training data cannot be modified by the agent. The authors should clarify this discrepancy and provide further insights.\n4. Performance in Other Locomotion Tasks: The paper's locomotion task evaluation focuses solely on the Walker2d environment. To establish the generalizability of the proposed method, it is crucial to conduct experiments on other well-established benchmarks, including Hopper, HalfCheetah, Ant, and Humanoid. By evaluating the performance across these mainstream locomotion tasks, the authors can provide a more comprehensive assessment and demonstrate the effectiveness of the proposed method in diverse and challenging environments. \n5. Minor Concern: Conducting experiments with only three random seeds may not provide sufficient statistical significance. Increasing the number of random seeds would strengthen the experimental evaluation.\n\n[1] Florence, Pete, et al. \"Implicit behavioral cloning.\" Conference on Robot Learning. PMLR, 2022.\n\nLIMITATIONS:\n1. Limited Task Scope: The experimental evaluation is confined to specific tasks, raising questions about DBC's performance in other types of tasks or more complex environments.\n2. Offline Learning Restriction: As an offline imitation learning algorithm, DBC lacks the capability to acquire knowledge through online interaction with the environment.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a method for guiding behavior cloning via state-action joint distribution learning. They train a diffusion model to maximize the log-likelihood of state and action pairs in conjunction with an imitation learning model that learns to mimic expert actions given state observations. They combine the two objectives to obtain a policy that predicts actions that fit the expert joint probability captured by a diffusion model. The authors demonstrate experiments on four continuous control domains, and show ablations to demonstrate the effects of their design choices.\n\nSTRENGTHS:\n1) The paper is well-structured and easy to follow. The experimental setup is clearly described with relevant details. The diffusion model for learning a joint distribution over state-action pairs is novel and is presented concisely.\n2) The considered tasks for this method are challenging and results on 4 out of 5 continuous control tasks show improvements over the considered baselines.\n\nWEAKNESSES:\n1) Intro: The authors state that implicit behavioral cloning (IBC) learns a joint distribution of state and action p(s,a). This is incorrect. IBC learns the joint “energy” E(s,a) but the learned distributions are still conditionals p(a|s). This is evident from the fact that IBC is trained to maximize the log-likelihood of actions in the dataset, and minimize those of sampled negative actions, given state inputs, and never trained to generate states or minimize the energy of negatively-sampled states. This makes a major claim of this paper incorrect.\n2) Section 3.1: Talking about reinforcement learning as a preliminary seems very absurd and misleading for a subsection on imitation learning, and also when the paper has nothing to do with learning from rewards.\n3) The core approach of this paper does not seem technically sound to me. Learning two distributions over the same random variables (state and action) seems to bring inconsistency to the probabilistic model where at least one of the two distributions is bound to be approximate.\n\nQUESTIONS:\n1) According to the authors, what are some design choices that are critical to learning the joint over two variables of different modalities?\n2) Table 2 shows the performance in Maze when using different generative models for guiding policy learning. Can the authors comment on the ranking of all methods in this table, why GANs performed the worst and Diffusion Models the best?\n\nLIMITATIONS:\nThe authors have addressed the limitations of their method in the supplementary submission.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a novel approach in the field of imitation learning. The authors address the challenge of learning from expert demonstrations without access to reward signals from the environment. They propose a framework called Diffusion Model-Augmented Behavioral Cloning (DBC) that combines the benefits of modeling both the conditional and joint probability of the expert distribution. The authors demonstrate the effectiveness of DBC through extensive experiments in various continuous control tasks, including navigation, robot arm manipulation, dexterous manipulation, and locomotion. However, there are certain strengths, weaknesses, limitations, and questions that need to be addressed regarding this work.\n\nSTRENGTHS:\n1) Extensive experiments: The article presents a wide range of experiments conducted on diverse tasks, including navigation, robot arm manipulation, dexterous manipulation, and locomotion. This comprehensive evaluation demonstrates the effectiveness of DBC in various scenarios.\n2) Novel approach: The proposed DBC framework offers a promising approach to imitation learning. By combining behavioral cloning with a diffusion model, the authors achieve more stable training compared to methods that combine behavioral cloning with GANs, such as GAIL.\n3) Clear comparison: The authors compare DBC with existing baselines, providing a clear understanding of its advantages over other methods in terms of performance and generalization.\n\nWEAKNESSES:\n1) Coefficient selection: The combination of the behavioral cloning loss and diffusion model loss in DBC relies on a simple addition of coefficients. However, the sensitivity of these coefficients to different environments should be further investigated. A more elegant approach and thorough ablation experiments on diverse tasks are needed to validate the coefficient selection process. Currently, this paper only provide ablation on this coefficient in a single environment (Maze).\n2) Lack of comprehensive comparison: Although DBC is compared with baselines, the article does not provide a comprehensive comparison with other state-of-the-art methods in the field of imitation learning. It would be beneficial to evaluate DBC against a wider range of approaches, including GAIL and other recent advancements, to gain a more comprehensive understanding of its relative performance. This would provide a clearer perspective on the strengths and weaknesses of DBC in comparison to alternative methods.\n\nQUESTIONS:\n1) How does the proposed DBC framework handle scenarios with complex and high-dimensional state spaces? Are there any specific limitations or challenges encountered in such cases?\n2) Are there any plans to extend the evaluation of the diffusion model to other environments beyond the MAZE setting? How does the diffusion model compare to other generative models in terms of performance and efficiency?\n\nLIMITATIONS:\nWhile the article highlights improved generalization performance, there is limited analysis or discussion on the factors contributing to this improvement. A deeper analysis of the generalization capabilities and limitations of DBC would enhance the understanding of its strengths and weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}