Title: The RL Perceptron: Generalisation Dynamics of Policy Learning in High Dimensions
Venue: NeurIPS 2023
Accepted: no
Keywords: Statistical physics of learning

generalisation models

reinforcement learning

Reinforce Policy gradient
Abstract: Reinforcement learning (RL) algorithms have proven transformative in a range of domains. To tackle real-world domains, these systems often use neural networks to learn policies directly from pixels or other high-dimensional sensory input. By contrast, much theory of RL has focused on discrete state spaces or worst-case analysis, and fundamental questions remain about the dynamics of policy learning in high-dimensional settings. Here, we propose a solvable high-dimensional model of RL that can capture a variety of learning protocols, and derive its typical dynamics as a set of closed-form ordinary differential equations (ODEs). We derive optimal schedules for the learning rates and task difficulty—analogous to annealing schemes and curricula during training in RL, and show that the model exhibits rich behaviour, including delayed learning under sparse rewards; a variety of learning regimes depending on reward baselines; and a speed-accuracy trade-off driven by reward stringency. Experiments on a variant of the Procgen game “Bossfight” also show such a speed-accuracy trade-off in practice. Together, these results take a step towards closing the gap between theory and practice in high-dimensional RL.

Abstract

Reinforcement learning (RL) algorithms have proven transformative in a range of domains. To tackle real-world domains, these systems often use neural networks to learn policies directly from pixels or other high-dimensional sensory input. By contrast, much theory of RL has focused on discrete state spaces or worst-case analysis, and fundamental questions remain about the dynamics of policy learning in high-dimensional settings. Here, we propose a solvable high-dimensional model of RL that can capture a variety of learning protocols, and derive its typical dynamics as a set of closed-form ordinary differential equations (ODEs). We derive optimal schedules for the learning rates and task difficulty-analogous to annealing schemes and curricula during training in RL-and show that the model exhibits rich behaviour, including delayed learning under sparse rewards; a variety of learning regimes depending on reward baselines; and a speed-accuracy trade-off driven by reward stringency. Experiments on a variant of the Procgen game 'Bossfight' also show such a speed-accuracy trade-off in practice. Together, these results take a step towards closing the gap between theory and practice in high-dimensional RL.

Recent years have seen rapid progress in Reinforcement Learning (RL): algorithmic and engineering 16 breakthroughs led to super-human performance in a variety of domains, for example complex games 17 like Go [Silver et al., 2016, Mnih et al., 2015]. Despite these practical successes, our theoretical 18 understanding of RL for high-dimensional problems requiring non-linear function approximation 19 is still limited. While comprehensive theoretical results exist for tabular RL, where the state and 20 action spaces are discrete and small enough for value functions to be represented directly, the curse 21 of dimensionality limits these methods to low-dimensional problems. The lack of a clear notion of 22 similarity between discrete states further means that tabular methods do not address the core question 23 of generalisation: how are values and policies extended to unseen states and across seen states [Kirk 24 et al., 2023]? As a consequence, much of this theoretical work is far from the current practice of RL, 25 which increasingly relies on deep neural networks to approximate and generalise value functions, 26 policies and other building blocks of RL. Moreover, while RL theory has often addressed 'worst-case' 27 performance and convergence behaviour, the typical behaviour has received comparatively little 28 attention (cf. further related work below). Meanwhile, a growing sub-field of deep learning theory 29 has employed tools from statistical mechanics to analyse various supervised learning paradigms 30 in the average-case, see Seung et al. [1992], Engel and Van den Broeck [2001], Carleo et al. [2019], 31 Bahri et al. [2020], Gabrié et al. [2023] for classical and recent reviews. While this approach has 32 recently been extended to curriculum learning [Saglietti et al., 2022], continual learning [Asanuma 33 et al., 2021, Lee et al., 2021, 2022], few-shot learning [Sorscher et al., 2022] and transfer learning 34 [Lampinen and Ganguli, 2018, Dhifallah and Lu, 2021, Gerace et al., 2022], RL has not been 35 analysed yet using statistical mechanics-a gap we address here by studying the high-dimensional 36 generalisation dynamics of a simple neural network trained on a reinforcement learning task. 37

(a)

Input

x

Teacher

w

∗

Student

w

y

∗

y

L

(

y, y

∗

)

y

∗

1

y

∗

2

y

∗

3

y

∗

T

y

1

y

2

y

3

y

T

Φ(

{

y

}

,

{

y

∗

}

)

x

1

x

2

x

3

x

T

w

∗

w

∗

w

∗

w

∗

w

w

w

w

s

1

s

2

s

3

s

T

(b)

(c)

0

2

4

6

0

0

.

2

0

.

4

0

.

6

0

.

8

Time,

t

Expected Reward

·

10

6

ODE

0

1

2

Q

·

10

7

0

2

4

6

0

2

4

t

R

·

10

3

·

10

6

The RL perceptron: In the classic teacher-student model of supervised learning [Gardner and 38 Derrida, 1989, Seung et al., 1992], a neural network called the student is trained on inputs x whose 39 labels y ∗ are given by another neural network called the teacher (see fig. 1a). The goal of the 40 student is to learn the function represented by the teacher from samples ( x , y ∗ ) . In RL, agents face 41 a sequential decision-making task in which a sequence of correct intermediate choices is required 42 to successfully complete an episode. We translate this process into the RL perceptron, a solvable 43 model for a high-dimensional, sequential policy learning task shown in fig. 1b. The student with 44 weights w takes a sequence of T choices over an episode. The correct choices are governed by 45 the same teacher network w ∗ , i.e. the same underlying rule throughout every time-step of every 46 episode. Crucially, unlike in the supervised learning setting, the student does not observe the correct 47 choice for each input; instead, it receives a reward which depends on whether earlier decisions are 48 correct. For instance, the student could receive a reward only if all T choices are correct, and no 49 reward otherwise-a learning signal that is considerably less informative than in supervised learning. 50 In addition to introducing the RL perceptron, our main contributions are as follows: 51

52

53

54

We derive an asymptotically exact set of Ordinary Differential Equations (ODEs) that describe the typical learning dynamics of policy gradient RL agents by building on classic work by Saad and Solla [1995], Biehl and Schwarze [1995], see section 2.1.

55

56

57

58

59

60

61

62

63

64

65

66

We use these ODEs to characterize learning behaviour in a diverse range of scenarios:

-We explore several sparse delayed reward schemes and investigate the impact of negative rewards (section 2.2)

-We derive optimal learning rate schedules and episode length curricula, and recover annealing strategies typically used in practice (section 2.3)

-At fixed learning rates, we identify ranges of learning rates for which learning is 'easy,' and 'hybrid-hard'-possibly causing a critical slowing down in the dynamics (section 2.4)

-We identify a speed-accuracy trade-off driven by reward stringency (section 2.5)

Finally we demonstrate that a similar speed-accuracy trade-off exists in simulations of highdimensional policy learning from pixels using the procgen environment 'Bossfight' [Cobbe et al., 2019], see section 3.

Further related work 67

68

69

70

71

72

73

74

Sample complexity in RL. An important line of work in the theory of RL focuses on the sample complexity and other learnability measures for specific classes of models such as tabular RL [Azar et al., 2017, Zhang et al., 2020b], state aggregation [Dong et al., 2019], various forms of MDPs [Jin et al., 2020, Yang and Wang, 2019, Modi et al., 2020, Ayoub et al., 2020, Du et al., 2019a, Zhang et al., 2022], reactive POMDPs [Krishnamurthy et al., 2016], and FLAMBE [Agarwal et al., 2020]. Here, we are instead concerned with the learning dynamics: how do reward rates, episode length, etc. influence the speed of learning and the final performance of the model.

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

Statistical learning theory for RL aims at finding complexity measures analogous to the Rademacher complexity or VC dimension from statistical learning theory for supervised learning Bartlett and Mendelson [2002], Vapnik and Chervonenkis [2015]. Proposals include the Bellman Rank Jiang et al. [2017], or the Eluder dimension [Russo and Van Roy, 2013] and its generalisations [Jin et al., 2021]. This approach focuses on worst-case analysis, which typically differs significantly from practice (at least in supervised learning [Zhang et al., 2021]). Furthermore, complexity measures for RL are generally more suitable for value-based methods; policy gradient methods have received less attention despite their prevalence in practice Bhandari and Russo [2019], Agarwal et al. [2021]. We focus instead on average-case dynamics of policy-gradient methods.

Dynamics of learning. A series of recent papers considered the dynamics of temporal-difference learning and policy gradient in the limit of wide two-layer neural networks Cai et al. [2019], Zhang et al. [2020a], Agazzi and Lu [2021, 2022]. These works focus on one of two 'wide' limits: either the neural tangent kernel [Jacot et al., 2018, Du et al., 2019b] or 'lazy' regime [Chizat et al., 2019], where the network behaves like an effective kernel machine and does not learn data-dependent features, which is key for efficient generalisation in high-dimensions. In our setting, the success of the student crucially relies on learning the weight vector of the teacher, which is hard for lazy methods [Ghorbani et al., 2019, 2020, Chizat and Bach, 2020, Refinetti et al., 2021]. The other 'wide' regime is the mean-field limit of interacting particles, akin to Mei et al. [2018], Chizat and Bach [2018], Rotskoff and Vanden-Eijnden [2018], where learning dynamics are captured by a non-linear partial differential equation. While this elegant description allows them to establish global convergence properties, it is hard to solve in practice. The ODE description we derive here instead will allow us to describe a series of effects in the following sections.

1 The RL Perceptron: setup and learning algorithm

We study the simplest possible student network, a perceptron with weight vector w that takes in high-dimensional inputs x ∈ R D and outputs y ( x ) = sgn ( w ⊺ x ) . We interpret the outputs y ( x ) as decisions, for example whether to go left or right in an environment. Because the student makes choices in response to high-dimensional inputs, it is analogous to a policy network. To train the network, we therefore consider a policy gradient learning update analogous to the REINFORCE algorithm [Sutton et al., 2000] that is adapted to the perceptron. At every timestep t during the µ th episode of length T , the agent occupies some state s t in the environment, receives an observation x µ t conditioned on s t , and takes an action y µ t = sgn ( w µ ⊺ x µ t ) , with t = 1 , . . . , T . The correct choice for each input is given by a fixed perceptron teacher with weights w ∗ . The crucial point is that the student does not have access to all the correct choices; it only receives a reward at the end of the episode if it completes the episode successfully, for example by making the correct decision at all times. If it does not succeed, it may receive a penalty; we will see in section 2.4 that receiving penalties is not always beneficial. In our setup, this translates into a weight update at the end of the µ th episode that is given by

where I is an indicator function and Φ is the criterion that determines whether the episode was 112 completed successfully-for instance, I (Φ) = ∏ T t θ ( y t y ∗ t ) (where θ is the step function) if the student 113 has to get every decision right in order to receive a reward. The update is general in the sense that the 114 term proportional to the learning rate η 1 > 0 prescribes the reward update for the fulfillment of the 115 condition, while the term proportional to η 2 ≥ 0 gives us the possibility to add a a penalty or negative 116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

(a)

10

0

10

2

10

4

10

6

10

8

0

0

.

5

1

Time,

t

Normalised Overlap,

R/

Q

η

2

.

0

.

005

.

01

(b)

10

0

10

2

10

4

10

6

10

8

0

0

.

5

1

Time,

t

Normalised Overlap,

R/

√

Q

T

0

= 8

T

0

= 7

T

0

= 6

T

0

= 5

T

0

= 4

T

0

= 3

T

0

= 2

(c)

10

0

10

2

10

4

10

6

10

8

0

0

.

5

1

Time,

t

Normalised Overlap,

R/

√

Q

r

b

.

0

.

005

.

01

√

reward should the student not succeed. Note that in the case of T = 1 , η 2 = 0 , and I (Φ) = θ ( yy ∗ ) , the learning rule updates the weight only if the student is correct on a given sample. It can thus be seen as the 'opposite' of the famous perceptron learning rule of supervised learning [Rosenblatt, 1962], where weights are only updated if the student is wrong. For a more in-detail discussion of the relation between the weight update in eq. (1) and the REINFORCE algorithm, see appendix A.

2 Theoretical Results

2.1 A set of dynamical equations captures the learning dynamics of an RL perceptron exactly

̸

The goal of the student during training is to emulate the teacher as closely as possible; or in other words, have a small number of disagreements with the teacher y ( x ) = y ∗ ( x ) . The generalisation error is given by the average number of disagreements

where the average ⟨·⟩ is taken over the inputs x , and we have introduced the scalar pre-activations for the student and the teacher, λ ≡ w · x / √ D and ν ≡ w ∗ · x / √ D , respectively. We can therefore transform the high-dimensional average over the inputs x into a low-dimensional average over the pre-activations ( λ, ν ) . The average in eq. (2) can be carried out by noting that the tuple ( λ, ν ) follow a jointly Gaussian distribution with means ⟨ λ ⟩ = ⟨ ν ⟩ = 0 and covariances

These covariances, or overlaps as they are sometimes called in the literature, have a simple interpreta132 tion. The overlap S is simply the length of the weight vector of the teacher; in the high-dimensional 133 limit D →∞ , S → 1 . Likewise, the overlap Q gives the length of the student weight vector; however, 134 this is a quantity that will vary during training. For example, when starting from small initial weights, 135 Q will be small, and grow throughout training. Lastly, the 'alignment' R quantifies the correlation 136 between the student and the teacher weight vector. At the beginning of training, R ≈ 0 , as both the 137 teacher and the initial condition of the student are drawn at random. As the student starts learning, 138 the overlap R increases. Evaluating the Gaussian average in eq. (2) shows that the generalisation 139 error is then a function of the normalised overlap ρ = R/ √ Q , and given by 140

The crucial point here is that we have reduced the description of the high-dimensional learning 141 problem from the D parameters of the student weight w to two time-evolving quantities, Q and R . 142 We now discuss how to analyse their dynamics. 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

182

183

184

185

186

187

188

189

190

191

192

The dynamics of order parameters. At any given point during training, the value of the order parameters determines the test error via eq. (4). But how do the order parameters evolve during training with the update rule eq. (1)? We followed the approach of Kinzel and Ruján [1990], Saad and Solla [1995], Biehl and Schwarze [1995] to derive a set of dynamical equations that describe the dynamics of the student in the high-dimensional limit where the input dimension goes to infinity. We give explicit dynamics for different reward conditions Φ , namely requiring all decisions correct in an episode of length T ; requiring n or more decisions correct in an episode of length T ; and receiving reward for each correct response. Due to the length of these expressions, we report the generic expression of the updates in the supplementary material in appendix B. Below, we state a version of the equations for the specific reward condition where the agent must survive until the end of an episode to receive a reward, I (Φ) = ∏ T t θ ( y t y ∗ t ) . The ODEs for the order parameters then read

where α ≡ µ/D serves as a continuous time variable in the limit D → ∞ (not to be confused with t which counts episode steps), and P = ( 1 -cos -1 ( R/ √ Q ) /π ) is the probability of a single correct decision. While our derivation of the equations follow heuristics from statistical physics, we anticipate that their asymptotic correctness in the limit D →∞ can be established rigorously using the techniques of Goldt et al. [2019], Veiga et al. [2022], Arnaboldi et al. [2023]. We illustrate the accuracy of these equations already in finite dimensions ( D = 900 ) in fig. 1c, where we show the expected reward, as well as the overlaps R and Q , of a student as measured during a simulation and from integration of the dynamical equations (solid and dotted lines, respectively).

The derivation of the dynamical equations that govern the learning dynamics of the RL perceptron are our first main result. Equipped with this tool, we now analyse several phenomena exhibited by the RL perceptron through a detailed study of these equations.

2.2 Learning protocols

The RL perceptron allows for the characterization of different RL protocols by adapting the reward condition Φ . We considered the following three settings:

Vanilla: The dynamics in the 'standard' case without penalty, η 2 = 0 , is shown in fig. 5a and fig. 5b. Rewards are sparsest in this protocol, and as a result we observe a characteristic initial plateau in expected reward followed by a rapid jump. The length of this plateau increases with T , consistent with the notion that sparse rewards make exploration hard and slow learning [Bellemare et al., 2016]. Plateaus during learning, which arise from saddle points in the loss landscape, have also been studied for (deep) neural networks in the supervised setting [Saad and Solla, 1995, Dauphin et al., 2014], but do not arise in the supervised perceptron. Hence the RL setting can qualitatively change the learning trajectory. The benefit of withholding penalties is that while slower, the perceptron reaches the highest level of expected reward in this case. This is a first example of a speed-accuracy trade-off that we will explore in more detail in section 2.5 and that we also found in our experiments with Bossfight in section 3.

Penalty: The initial plateau can be reduced by providing a penalty or negative reward ( η 2 > 0 ) when the student fails in the task. This change provides weight updates much earlier in training and thus accelerates the escape from the plateau. The dynamics under this protocol are shown in fig. 2a. It is clear the penalty provides an initial speed-up in learning, as expected if the agent were to be unaligned and more likely to commit an error. However, a high penalty can create additional sub-optimal fixed points in the dynamics leading to a low asymptotic performance (more on this in section 2.4). In the simulations, finite size effects occasionally permit escape from the sub-optimal fixed point and jumps to the optimal one, leading to a high variance in the results.

Subtask and breadcrumbs: The model is also able to capture the dynamics of more complicated protocols: fig. 3b shows learning under the protocol where a smaller sub-reward is received if the agent survives beyond a shorter duration T 0 < T , i.e. some reward is still received even if the agent does not survive for the entire episode. Another learning protocol we can capture is that of 'graded-breadcrumbs', where the agent receives a small reward r b for every correct decision made

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

(a)

0

0

.

5

1

1

.

5

2

0

0

.

5

1

Time,

t

ρ

Optimal

10

1

·

10

5

T

(b)

0

0

.

5

1

1

.

5

2

0

0

.

5

1

Time,

t

ρ

Optimal

10

0

·

10

5

η

1

(c)

10

0

10

2

10

4

10

6

10

-

4

10

-

2

10

0

10

2

Time,

t

η

10

0

10

1

10

2

10

3

10

4

T

η

= 0

.

1

η

= 1

η

= 10

T

= 3

T

= 8

T

= 13

in an episode, i.e. like the previous method some reward is still received even if the agent does not survive for the entire episode, these dynamics are captured in fig. 3c.

2.3 Optimal hyper-parameter schedules: make episodes longer and anneal your learning rate

Hyper-parameter schedules are crucial for successful training of RL agents. In our setup, the two most important hyper-parameters are the learning rates and the episode length. In the RL perceptron, we can derive optimal schedules for both hyper-parameters. For simplicity, here we report the results in the spherical case, where the length of the student vector is fixed at √ D (we discuss the unconstrained case in the appendix C), then Q ( α ) = 1 at all times and we only need to track the teacher-student overlap ρ = R/ √ Q , which quantifies the generalisation performance of the agent. Keeping the choice I (Φ) = ∏ T t =1 θ ( y t y ∗ t ) and turning off the penalty term ( η 2 = 0 ), we find that the teacher-student overlap is governed by the equation

The optimal schedules over episodes for T and η can then be found by maximising the change in overlap at each update, i.e. setting ∂ ∂T ( dρ dα ) and ∂ ∂η ( dρ dα ) to zero respectively. After some calculations, we find the optimal schedules to be

where ⌊·⌋ indicates the floor function.

Figure 3a shows the evolution of ρ under the optimal episode length schedule (dashed) compared to other constant episode lengths (green). Similarly, fig. 3b shows the evolution of ρ under the optimal learning rate schedule (dashed) compared to other constant learning rates (blue). The functional forms of T opt and η opt over time are shown in fig. 3c.

During learning the student seeks increasingly refined information to improve its expected reward. 212 This simple observation explains the monotonic increase of the optimal episode length and the 213 decrease in learning rates. Starting from the episode duration, we can observe that given the discrete 214 nature of the decisions, information obtained from the rewards simply pushes the decision boundary 215 towards a partition of the input space. This partition is determined by the episode length T and 216 correspond to a fraction 1 / 2 T of the entire input space. Therefore a positive reward conveys T bits of 217 information. At a fixed learning rate, when the student becomes proficient in the task it will not be 218 able to improve further the decision boundary, and will fluctuate around the optimal solution unless 219 longer episodes are provided. 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

(a)

(b)

(c)

0

0

.

2

0

.

4

0

.

6

0

.

8

1

0

0

.

2

0

.

4

0

.

6

0

.

8

1

η

2

ρ

fix

0

0

.

2

0

.

4

0

.

6

0

.

8

1

0

0

.

2

0

.

4

0

.

6

0

.

8

1

η

1

η

2

0

0

.

2

0

.

4

0

.

6

0

.

8

1

0

0

.

2

0

.

4

0

.

6

0

.

8

1

η

1

η

2

Hybrid-hard

Easy

Our analysis shows that a polynomial increase in the episode length gives the optimal performance in the RL perceptron, see fig. 3c (top); increasing T in the RL perceptron is akin to increasing task difficulty, and the polynomial scheduling of T opt specifies a curriculum. Curricula of increasing task difficulty are commonly used in RL to give convergence speed-ups and learn problems that otherwise would be too difficult to learn ab initio Narvekar et al. [2020]. Analogously, the fluctuations can be reduced by annealing the learning rate and averaging over a larger number of samples. Akin to work in RL literature studying adaptive step-sizes [Dabney, 2014, Pirotta et al., 2013], we find that annealing the learning rate during training is beneficial for greater speed and generalisation performance. For the RL perceptron, a polynomial decay in the learning rate gives optimal performance as shown in fig. 3c (bottom), consistent with work in the parallel area of high-dimensional non-convex optimization problems [d'Ascoli et al., 2022], and stochastic approximation algorithms in RL [Dalal et al., 2017].

2.4 Phase Space

With a non-zero penalty ( η 2 ), the generalisation performance of the agent can enter different regimes of learning. This is most clearly exemplified in the spherical case, where the number of fixed points of the ODE governing the dynamics of the overlap exist in distinct phases determined by the combination of reward and penalty. For the simplest case ( I (Φ) = ∏ T t ( y t y ∗ t ) ) these phases are shown in fig. 4. Figure 4a shows the fixed points achievable over a range of penalties for a fixed η 1 = 1 (obtained from a numerical solution of the ODE in ρ ). There are two distinct regions: 1) Easy , where there is a unique fixed point and the algorithm naturally converges to this optimal ρ fix from a random initialisation, 2) a Hybrid-hard region (given the analogy with results from inference problems Ricci-Tersenghi et al. [2019]), where there are two stable (1 good and 1 bad) fixed points, and 1 unstable fixed point, and either stable point is achievable depending on the initialisation of the student (orange). The 'hybrid-hard' region separates two easy regions with very distinct performance levels. In this region the algorithm with high probability converges to ρ fix with the worse performance level. These two regions are visualised in ( η 1 , η 2 ) space in fig. 4b for an episode length of T = 13 . The topology of these regions are also governed by episode length, with a sufficiently small T reducing the the area of the 'hybrid-hard' phase to zero, meaning there is always 1 stable fixed point which may not necessarily give 'good' generalisation. Figure 4c shows the phase plot for T = 8 , where the orange (hybrid-hard) has shrunk, this corresponds to the s-shaped curve in fig. 4a becoming flatter (closer to monotonic). Learning with η 2 This is not a peculiarity specific to the spherical case, indeed, we observe different regimes in the learning dynamics in the setting with unrestricted Q which we report in appendix C.

These phases show that at a fixed η 1 increasing η 2 will eventually lead to a first order phase transition, and the speed benefits gained from a non-zero η 2 will be nullified due to the transition into the hybrid-hard phase. In fact, when taking η 2 close to the transition point, instead of speeding up learning there is the presence of a critical slowing down, which we report in appendix C.

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

(a)

10

0

10

2

10

4

10

6

10

8

0

0

.

5

1

Time,

t

Expected Reward

(b)

10

0

10

2

10

4

10

6

10

8

0

0

.

5

1

Time,

t

Normalised Overlap,

R/

√

Q

T

12

9

6

3

(c)

10

0

10

2

10

4

10

6

10

8

0

0

.

5

1

Time,

t

Normalised Overlap,

R/

√

Q

n

= 7

n

= 8

n

= 9

n

= 10

n

= 11

n

= 12

n

= 13

A common problem with REINFORCE is high variance gradient estimates leading to bad performance [Marbach and Tsitsiklis, 2003, Schulman et al., 2015]. The reward ( η 1 ) and punishment ( η 2 ) magnitude alters the variance of the updates, and we show that the interplay between reward, penalty and reward-condition and their effect on performance can be probed within our model. This framework opens the possibility for studying phase transitions between learning regimes [Gamarnik et al., 2022].

2.5 Speed-accuracy trade-off

Figure 5c shows the evolution of normalised overlap ρ = R/ √ Q between the student and teacher obtained from simulations and from solving the ODEs in the case where n or more decisions must be correctly made in an episode of length T = 13 in order to receive a reward (with η 2 = 0 ). We observe a speed-accuracy trade-off, where decreasing n increases the initial speed of learning but leads to worse asymptotic performance; this alleviates the initial plateau in learning seen previously in fig. 5b at the cost of good generalisation. In essence, a lax reward function is probabilistically more achievable early in learning; but it rewards some fraction of incorrect decisions, leading to lower asymptotic accuracy. By contrast a stringent reward function slows learning but eventually produces a highly aligned student. For a given MDP, it is known that arbitrary shaping applied to the reward function will change the optimal policy (reduce asymptotic performance) [Ng et al., 1999]. Empirically, reward shaping has been shown to speed up learning and help overcome difficult exploration problems [Gullapalli and Barto, 1992]. Reconciling these results with the phenomena observed in our setting is an interesting avenue for future work.

3 Experiments

To verify that our theoretical framework captures qualitative features of more general settings, we 278 train agents from pixels on the Procgen [Cobbe et al., 2019] game 'Bossfight' (example frame, fig. 6a 279 (top)). To remain close to our theoretical setting, we consider a modified version of the game where 280 the agent cannot defeat the enemy and wins only if it survives for a given duration T . On each 281 timestep the agent has the binary choice of moving left/right and aims to dodge incoming projectiles. 282 We give the agent h lives, where the agent loses a life if struck by a projectile and continues an 283 episode if it has lives remaining. This reward structure reflects the sparse reward setup from our 284 theory and is analogous to requiring n out of T decisions to be correct within an episode. We further 285 add asteroids at the left and right boundaries of the playing field which destroy the agent on contact, 286 such that the agent cannot hide in the corners. Observations, shown in fig. 6a (bottom), are centred on 287 the agent and downsampled to size 35 × 64 with three colour channels, yielding a 6720 dimensional 288 input. The pixels corresponding to the agent are set to zero since these otherwise act as near-constant 289 bias inputs not present in our model. The agent is endowed with a shallow policy network with 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

315

316

317

318

(a)

(b)

2

4

20

40

60

Episode

%Successful Episodes

h

= 1

h

= 2

h

= 3

h

= 4

·

10

5

(c)

0

.

01

(Left)

-

0

.

01

(Right)

0

.

1

(Left)

-

0

.

1

(Right)

ω

ω

logistic output unit that indicates the probability of left or right action. The weights of the policy network are trained using the policy gradient update of eq. (1) under a pure random policy.

To study the speed-accuracy trade-off, we train agents with different numbers of lives. As seen in fig. 6b, we observe a clear speed-accuracy trade-off mediated by agent health consistent with our theoretical findings (c.f. fig. 3c). Figure 6c shows the final policy weights for agents trained with h = 1 and h = 4 . These show interpretable structure, roughly split into thirds vertically: the weights in the top third detect the position of the boss and centre the agent beneath it; this causes projectiles to arrive vertically rather than obliquely, making them easier to dodge. The weights in the middle third dodge projectiles. Finally, the weights in the bottom third avoid asteroids near the agent. Notably, the agent trained in the more stringent reward condition ( h = 1 ) places greater weight on dodging projectiles, showing the qualitative impact of reward on learned policy. Hence similar qualitative phenomena as in our theoretical model can arise in more general settings.

4 Concluding perspectives

The RL perceptron provides a framework to investigate high-dimensional policy gradient learning in RL for a range of plausible sparse reward structures. We derive closed ODEs that capture the averagecase learning dynamics in high-dimensional settings. The reduction of the high-dimensional learning dynamics to a low-dimensional set of differential equations permits a precise, quantitative analysis of learning behaviours: computing optimal hyper-parameter schedules, or tracing out phase diagrams of learnability. Our framework offers a starting point to explore additional settings that are closer to many real-world RL scenarios, such as those with conditional next states. Furthermore, the RL perceptron offers a means to study common training practices, including curricula; and more advanced algorithms, like actor-critic methods. We hope to extract more analytical insights from the ODEs, particularly on how initialization and learning rate influence an agent's learning regime. Our findings emphasize the intricate interplay of task, reward, architecture, and algorithm in modern RL systems.

References

Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural complexity and representation learning of low rank mdps. Advances in neural information processing systems , 33:20095-20107, 2020.

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

353

354

355

356

357

358

359

360

361

362

363

364

Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. The Journal of Machine Learning Research , 22(1):4431-4506, 2021.

Andrea Agazzi and Jianfeng Lu. Global optimality of softmax policy gradient with single hidden layer neural networks in the mean-field regime. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=bB2drc7DPuB .

Andrea Agazzi and Jianfeng Lu. Temporal-difference learning with nonlinear function approximation: lazy training and mean field regimes. In Joan Bruna, Jan Hesthaven, and Lenka Zdeborova, editors, Proceedings of the 2nd Mathematical and Scientific Machine Learning Conference , volume 145 of Proceedings of Machine Learning Research , pages 37-74. PMLR, 16-19 Aug 2022. URL https://proceedings.mlr.press/v145/agazzi22a.html .

Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro. From high-dimensional & mean-field dynamics to dimensionless odes: A unifying approach to sgd in two-layers networks. arXiv preprint arXiv:2302.05882 , 2023.

Haruka Asanuma, Shiro Takagi, Yoshihiro Nagano, Yuki Yoshida, Yasuhiko Igarashi, and Masato Okada. Statistical mechanical analysis of catastrophic forgetting in continual learning with teacher and student networks. Journal of the Physical Society of Japan , 90(10):104001, 2021.

Alex Ayoub, Zeyu Jia, Csaba Szepesvari, Mengdi Wang, and Lin Yang. Model-based reinforcement learning with value-targeted regression. In International Conference on Machine Learning , pages 463-474. PMLR, 2020.

Mohammad Gheshlaghi Azar, Ian Osband, and Rémi Munos. Minimax regret bounds for reinforcement learning. In International Conference on Machine Learning , pages 263-272. PMLR, 2017.

Yasaman Bahri, Jonathan Kadmon, Jeffrey Pennington, Sam S. Schoenholz, Jascha Sohl-Dickstein, and Surya Ganguli. Statistical mechanics of deep learning. Annual Review of Condensed Matter Physics , 11(1):501-528, 2020. doi: 10.1146/annurev-conmatphys-031119-050745. URL https: //doi.org/10.1146/annurev-conmatphys-031119-050745 .

Peter L Bartlett and Shahar Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research , 3(Nov):463-482, 2002.

Marc Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying count-based exploration and intrinsic motivation. Advances in neural information processing systems , 29, 2016.

Jalaj Bhandari and Daniel Russo. Global optimality guarantees for policy gradient methods. arXiv preprint arXiv:1906.01786 , 2019.

Michael Biehl and Holm Schwarze. Learning by on-line gradient descent. Journal of Physics A: Mathematical and general , 28(3):643, 1995.

Qi Cai, Zhuoran Yang, Jason D Lee, and Zhaoran Wang. Neural temporal-difference learning converges to global optima. Advances in Neural Information Processing Systems , 32, 2019.

Giuseppe Carleo, Ignacio Cirac, Kyle Cranmer, Laurent Daudet, Maria Schuld, Naftali Tishby, Leslie Vogt-Maranto, and Lenka Zdeborová. Machine learning and the physical sciences. Reviews of Modern Physics , 91(4):045002, 2019.

Chizat and F. Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. In Advances in Neural Information Processing Systems 31 , pages 3040-3050, 2018.

Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on Learning Theory , pages 1305-1338. PMLR, 2020.

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

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32 , pages 2933-2943. Curran Associates, Inc., 2019.

Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. arXiv preprint arXiv:1912.01588 , 2019.

William M. Dabney. Adaptive step-sizes for reinforcement learning. 2014.

Gal Dalal, Balázs Szörényi, Gugan Thoppe, and Shie Mannor. Concentration bounds for two timescale stochastic approximation with applications to reinforcement learning. CoRR , abs/1703.05376, 2017. URL http://arxiv.org/abs/1703.05376 .

Stéphane d'Ascoli, Maria Refinetti, and Giulio Biroli. Optimal learning rate schedules in highdimensional non-convex optimization problems, 2022.

Yann N Dauphin, Razvan Pascanu, Caglar Gulcehre, Kyunghyun Cho, Surya Ganguli, and Yoshua Bengio. Identifying and attacking the saddle point problem in high-dimensional non-convex optimization. Advances in neural information processing systems , 27, 2014.

Oussama Dhifallah and Yue M Lu. Phase transitions in transfer learning for high-dimensional perceptrons. Entropy , 23(4):400, 2021.

Shi Dong, Benjamin Van Roy, and Zhengyuan Zhou. Provably efficient reinforcement learning with aggregated states. arXiv preprint arXiv:1912.06366 , 2019.

Simon Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford. Provably efficient rl with rich observations via latent state decoding. In International Conference on Machine Learning , pages 1665-1674. PMLR, 2019a.

S.S. Du, X. Zhai, B. Poczos, and A. Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations , 2019b.

Andreas Engel and Christian Van den Broeck. Statistical mechanics of learning . Cambridge University Press, 2001.

Marylou Gabrié, Surya Ganguli, Carlo Lucibello, and Riccardo Zecchina. Neural networks: from the perceptron to deep nets. arXiv preprint arXiv:2304.06636 , 2023.

David Gamarnik, Cristopher Moore, and Lenka Zdeborová . Disordered systems insights on computational hardness. Journal of Statistical Mechanics: Theory and Experiment , 2022(11):114015, nov 2022. doi: 10.1088/1742-5468/ac9cc8. URL https://doi.org/10.1088%2F1742-5468% 2Fac9cc8 .

Elizabeth Gardner and Bernard Derrida. Three unfinished works on the optimal storage capacity of networks. Journal of Physics A: Mathematical and General , 22(12):1983, 1989.

Federica Gerace, Luca Saglietti, Stefano Sarao Mannelli, Andrew Saxe, and Lenka Zdeborová. Probing transfer learning with a model of synthetic correlated datasets. Machine Learning: Science and Technology , 3(1):015030, 2022.

Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limitations of lazy training of two-layers neural network. In Advances in Neural Information Processing Systems , volume 32, pages 9111-9121, 2019.

Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? In Advances in Neural Information Processing Systems , volume 33, 2020.

Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information processing systems , 32, 2019.

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

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

454

455

Vijaykumar Gullapalli and Andrew G Barto. Shaping as a method for accelerating reinforcement learning. In Proceedings of the 1992 IEEE international symposium on intelligent control , pages 554-559. IEEE, 1992.

Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems 32 , pages 8571-8580, 2018.

Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. Contextual decision processes with low bellman rank are pac-learnable. In International Conference on Machine Learning , pages 1704-1713. PMLR, 2017.

Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael I Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory , pages 2137-2143. PMLR, 2020.

Chi Jin, Qinghua Liu, and Sobhan Miryoosefi. Bellman eluder dimension: New rich classes of rl problems, and sample-efficient algorithms. Advances in neural information processing systems , 34:13406-13418, 2021.

Kinzel and P. Ruján. Improving a Network Generalization Ability by Selecting Examples. EPL (Europhysics Letters) , 13(5):473-477, 1990.

Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A Survey of Zero-shot Generalisation in Deep Reinforcement Learning. Journal of Artificial Intelligence Research , 76: 201-264, January 2023. ISSN 1076-9757. doi: 10.1613/jair.1.14174. URL https://jair.org/ index.php/jair/article/view/14174 .

Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Pac reinforcement learning with rich observations. Advances in Neural Information Processing Systems , 29, 2016.

Andrew K Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. arXiv preprint arXiv:1809.10374 , 2018.

Sebastian Lee, Sebastian Goldt, and Andrew Saxe. Continual learning in the teacher-student setup: Impact of task similarity. In International Conference on Machine Learning , pages 6109-6119. PMLR, 2021.

Sebastian Lee, Stefano Sarao Mannelli, Claudia Clopath, Sebastian Goldt, and Andrew Saxe. Maslow's hammer for catastrophic forgetting: Node re-use vs node activation. arXiv preprint arXiv:2205.09029 , 2022.

Peter Marbach and John N. Tsitsiklis. Approximate gradient methods in policy-space optimization of markov reward processes. Discrete Event Dynamic Systems , 13:111-148, 2003.

Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of twolayer neural networks. Proceedings of the National Academy of Sciences , 115(33):E7665-E7671, 2018.

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.

Aditya Modi, Nan Jiang, Ambuj Tewari, and Satinder Singh. Sample complexity of reinforcement learning using linearly combined model ensembles. In International Conference on Artificial Intelligence and Statistics , pages 2010-2020. PMLR, 2020.

Sanmit Narvekar, Bei Peng, Matteo Leonetti, Jivko Sinapov, Matthew E. Taylor, and Peter Stone. Curriculum learning for reinforcement learning domains: A framework and survey, 2020.

Andrew Y Ng, Daishi Harada, and Stuart Russell. Policy invariance under reward transformations: Theory and application to reward shaping. In Icml , volume 99, pages 278-287. Citeseer, 1999.

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

Matteo Pirotta, Marcello Restelli, and Luca Bascetta. Adaptive step-size for policy gradient methods. In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Weinberger, editors, Advances in Neural Information Processing Systems , volume 26. Curran Associates, Inc., 2013. URL https://proceedings.neurips.cc/paper_files/paper/2013/file/ f64eac11f2cd8f0efa196f8ad173178e-Paper.pdf .

Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborova. Classifying highdimensional gaussian mixtures: Where kernel methods fail and neural networks succeed. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 8936-8947. PMLR, 18-24 Jul 2021. URL https://proceedings.mlr.press/v139/refinetti21b.html .

Federico Ricci-Tersenghi, Guilhem Semerjian, and Lenka Zdeborová. Typology of phase transitions in bayesian inference problems. Physical Review E , 99(4):042109, 2019.

Rosenblatt. Principles ofNeurodynamics . Spartan, New York, 1962.

G.M. Rotskoff and E. Vanden-Eijnden. Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks. In Advances in Neural Information Processing Systems 31 , pages 7146-7155, 2018. URL http://arxiv.org/abs/1805.00915 .

Daniel Russo and Benjamin Van Roy. Eluder dimension and the sample complexity of optimistic exploration. Advances in Neural Information Processing Systems , 26, 2013.

David Saad and Sara A Solla. On-line learning in soft committee machines. Physical Review E , 52 (4):4225, 1995.

Luca Saglietti, Stefano Sarao Mannelli, and Andrew Saxe. An analytical theory of curriculum learning in teacher-student networks. Journal of Statistical Mechanics: Theory and Experiment , 2022(11):114014, 2022.

John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional continuous control using generalized advantage estimation, 2015. URL https://arxiv.org/ abs/1506.02438 .

Hyunjune Sebastian Seung, Haim Sompolinsky, and Naftali Tishby. Statistical mechanics of learning from examples. Physical review A , 45(8):6056, 1992.

David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature , 529(7587):484-489, 2016.

Ben Sorscher, Surya Ganguli, and Haim Sompolinsky. Neural representational geometry underlies few-shot concept learning. Proceedings of the National Academy of Sciences , 119(43): e2200800119, 2022. doi: 10.1073/pnas.2200800119. URL https://www.pnas.org/doi/abs/ 10.1073/pnas.2200800119 .

S. Sutton, D. Mcallester, S. Singh, and Y. Mansour. Policy gradient methods for reinforcement learning with function approximation. In Advances in Neural Information Processing Systems 12 , volume 12, pages 1057-1063. MIT Press, 2000.

Vladimir N Vapnik and A Ya Chervonenkis. On the uniform convergence of relative frequencies of events to their probabilities. Measures of complexity: festschrift for alexey chervonenkis , pages 11-30, 2015.

Rodrigo Veiga, Ludovic STEPHAN, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborova. Phase diagram of stochastic gradient descent in high-dimensional two-layer neural networks. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id= GL-3WEdNRM .

Lin Yang and Mengdi Wang. Sample-optimal parametric q-learning using linearly additive features. In International Conference on Machine Learning , pages 6995-7004. PMLR, 2019.

504

505

506

507

508

509

510

511

512

513

514

515

516

517

Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM , 64(3):107-115, 2021.

Xuezhou Zhang, Yuda Song, Masatoshi Uehara, Mengdi Wang, Alekh Agarwal, and Wen Sun. Efficient reinforcement learning in block mdps: A model-free representation learning approach. In International Conference on Machine Learning , pages 26517-26547. PMLR, 2022.

Yufeng Zhang, Qi Cai, Zhuoran Yang, Yongxin Chen, and Zhaoran Wang. Can temporal-difference and q-learning learn representation? a mean-field theory. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 19680-19692. Curran Associates, Inc., 2020a. URL https://proceedings.neurips. cc/paper_files/paper/2020/file/e3bc4e7f243ebc05d66a0568a3331966-Paper.pdf .

Zihan Zhang, Yuan Zhou, and Xiangyang Ji. Almost optimal model-free reinforcement learningvia reference-advantage decomposition. Advances in Neural Information Processing Systems , 33: 15198-15207, 2020b.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper focuses on developing theories for the learning dynamics of policy gradient reinforcement learning (RL) algorithms with a particular focus on high-dimensional latent feature space. As an early work along this direction, the authors study a binary-action environment setup for simplicity. The authors develop ODE-based learning dynamic equations that generalizes across diverse protocols, including different policy horizons, the existance of failure penalties, and the choices of dense / sparse rewards. The authors further develop theories for optimal learning rate, optimal horizon scheduling, and learnability with respect to these hyperparameters. The authors finally conduct an experiment on vision-based Procgen Bossfight environment and demonstrate that under more general settings, similar phenomena arise as their theoretical model developed under simpler setups.\n\nSTRENGTHS:\nOverall the paper is well-structured and well-written. Under the binary-action environment setup, the ODE dynamics model developed by the authors is shown to accurately describe diverse common learning setups. There are also extensive experiments and plots that illustrate the difference of learning dynamics under different environment and optimization parameters, which provide much insights for the readers.\n\nWEAKNESSES:\nIn the main paper, the latent feature space dimension $D$ is fixed to 900 except Procgen. It would be helpful if authors provide more analysis on the influence of latent dimension $D$ on the learning dynamics. Empirically, for which $D$ is author's proposed ODE-based learning dynamics equation still accurate?\n\nFor Fig. 6b, Plotting environments of different episode lengths by comparing their \"Number of episodes\" seems misleading. Authors claim that agents learn slower for environments with shorter episode lengths, but this is not accurate. If one compares the total time step of learning (num episodes * episode length), agents acturally learn faster on these shorter horizon environments.\n\nQUESTIONS:\nSee \"weakness\" section above.\n\nMinor typos: \n\n- Equation 2: looks like 1-eps_g is defined as disagreement instead of eps_g.\n\n- Line 194, Fig. 3c -> Fig. 2c\n\nLIMITATIONS:\nLimitations need to be explicitly addressed in the conclusion section, including (1) the simplicity of the problem setup studied by the paper (binary action scenarios); and (2) the paper's focus on shallow, one-layer neural network that takes high dimensional feature as input. For more general applications, neural networks typically consist of many layers, activations, and normalizations stacked on top of each other, so empirical analysis in these scenarios will be particularly helpful.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a model for solving high-dimensional problems in reinforcement learning (RL) referred to as the RL perceptron. The model is used as a framework for analyzing generalization dynamics of simple neural networks for RL tasks. More precisely, the model employs a student teacher design in which the student takes a sequence of choices and the correct choices are given via the teacher. However, the student does not have access to the correct choice at every step but rather only receives a signal at the end of each episode. As such, the model is studied as a sequential version of the perceptron algorithm. The work first derives a set of differential equations that capture the learning dynamics of the model. The dynamics are analyzed via the overlap of the weight vectors of the student and teacher respectively. The manuscript studies multiple reward settings: a vanilla setting that constitutes sparse rewards, a setting with penalty at every step and a setting where the agent is given small sub-rewards after a certain amount of time. Each of the proposed settings provides insights into how different reward functions lead to different solutions for optimal parameters when solving the system. The derivation of optimal hyperparameters demonstrates that annealing the learning rate and building a curriculum of episodes are crucial for optimal convergence. Then, it is shown that there exist phase transitions under different learning rates that can lead to convergence to sub-optimal minima and that there is a speed-accuracy tradeoff when varying reward functions. Lastly an experimental section provides insides in the practical properties of the algorithm that closely follow the analytical results on the speed-accuracy tradeoff.\n\nSTRENGTHS:\n1.) First, I would like to say that I enjoyed reading this paper. I think it is well written and well structured with a clear line of reasoning throughout the manuscript. The figures are illustrative of the analysis and the captions are sufficiently descriptive to understand the plots quickly.\n\n2.) The idea of employing a perceptron-like algorithm in order to understand the policy gradient system dynamics is novel to the best of my knowledge and the generated insights are interesting. I think this is a nice contribution as it offers a way to analyze high-dimensional RL systems in a different way than the commonly employed linear MDPs.\n\n3.) The manuscript is technically sound and the analysis provides a good understanding of the proposed algorithm and its inner workings. I did check the math in the appendix for crude errors and was unable to find any but I did not try to understand all the math in detail.\n\n4.) The manuscript provides a way to think about RL systems that is not common and I think it is likely going to be useful in understanding some parts of the systems that may have been hard to understand previously. As such, I think it is a decent contribution that can likely be built upon by others in the sub-field of policy gradient methods.\n\nWEAKNESSES:\na.) The connection to policy gradients was not immediately clear to me. It might make sense to move equation 9 in place of equation 1 and highlight the connection between the perceptron update rule and the REINFORCE algorithm in a brief sentence. I think it would be good to highlight that the update rule uses an approximation that is only accurate early in training.\n\nb.) The limitations of the model are addressed rather sparsely. I think the work would benefit from having a clearer picture of the weaknesses of the approach which would enable researchers to use it and improve upon it in the future.\n\nc.)  The model seems to be unable to solve the benchmark problem fully. However, it is hard to tell whether that is just a limitation of the model or whether the task is hard the way it is designed with the changes in the manuscript. Having a baseline performance line in the plot or giving a brief sentence of what the expected performance of a commonly used RL algorithm on the benchmark is would be very useful to determine the capabilities of the method.\n\nd.) I think one key thing that is missing from the paper and would make it a very strong contribution is to show the relationship between the proposed model and common deep RL methods. It would have been nice to have a direct comparison from the proposed method to a neural network approach using standard REINFORCE-like updates to see if there is a correlation produced from insights of the RL-perceptron with the behavior of the regular deep neural network. I do understand that space is limited though.\n\n----\n\nOverall, I think changes that could improve the manuscript would establish the model's connections to other research that people have done in the area. This could, for instance, include baseline performances, transfer of insights to other methods or any theoretical results that put the work into reference with commonly knows results.\n\nMinor clarity suggestions:\n* Line 250, there is a broken off sentence in there that should be removed.\n* Line 562 equation reference is missing.\n\nQUESTIONS:\nQ1: What is the relationship between the proposed algorithm and standard policy gradient approaches with deep neural networks? Are there any conclusions that can be transferred from the RL perceptron to standard deep RL algorithms?\n\nQ2: What is the expected performance on the benchmark task where we would consider the task solved? Is it possible to win 100% of the time? This seems to be crucial information in order to determine the practical utility of the algorithm.\n\nQ3: One question that I have is about the connection between the proposed method and the standard perceptron algorithm. Is there a reduction from the RL perceptron to the standard perceptron? \n\nSuppose the model is updated at the end of every epoch. I’m thinking something along the following lines. Let n be the number of correct guesses made by the algorithm and N be the number of total timesteps in a trajectory. Then N-n is the number of incorrect guesses. Assume the sparse reward model where we have to be correct on k guesses to see a reward. Now, take a full trajectory as an input to the algorithm. If k < n, label the trajectory positive, otherwise negative. Of course there is some detail missing that would include the teacher in the labeling process but you should get a rough idea. If such a reduction exists it would reduce the problem to binary classification of whether or not a certain action sequence will lead to success.\n\nLIMITATIONS:\nI do not believe there is any negative societal impact that needs to be addressed regarding this work. The limitations are addressed rather sparsely. As stated before, I believe that the manuscript would benefit from more structured limitations sections. One limitation that I see is that the model currently requires actions to be discrete and (possibly ?) binary. While for a first version of the model this is absolutely fine, this might be something that researchers can work on in the future. Another limitation of the model is that it seems to not be able to fully solve the suggested benchmark problems. Again, I don't believe this to be an issue for the manuscript as the goal is not to provide a state-of-the-art model but rather to make progress towards understanding the learning dynamics of deep RL systems. Yet, in the future a goal should be to have models that can be described analytically that also achieve comparably high performance on realistic tasks.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: SUMMARY:
The authors develop a set of differential equations that describe the learning dynamics in high-dimensional settings, allowing for a quantitative analysis of learning behaviors. This framework enables the computation of optimal hyper-parameter schedules and the visualization of phase diagrams for learnability. It also serves as a starting point for exploring RL scenarios closer to real-world situations, including those with conditional next states. The RL perceptron can be used to investigate various training practices, such as curricula, and advanced algorithms like actor-critic methods. The authors aim to gain analytical insights from the differential equations, particularly regarding how initialization and learning rate affect an agent's learning process. Overall, their research highlights the complex interplay between task, reward, architecture, and algorithm in modern RL systems.

STRENGTHS:
* The paper tackles an important problem of understanding high dimensional RL policies. 
* The method section seems theoritically sound.

WEAKNESSES:
Experiments are not convincing. For a paper investigating such an important problem, the paper should have shown results on multiple environments.

QUESTIONS:
I don't have any questions.

LIMITATIONS:
A considerable limitation of the paper is the lack of a comprehensive evaluation of the method. The paper provides only results on one experiment. I would suggest the authors conduct more thourough experiments in order to make the applicability of method more clear.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe work proposes a theoretical framework to study the average case learning behavior of deep RL policy gradient methods. The framework is based on ODEs that can describe the typical learning dynamics of PG RL agents. The framework is used in various settings to describe learning behaviors in these and a final experiment on training a policy-gradient agent in a ProcGen game is added to bridge the theory-practice gap. This final experiment verifies that a speed-accuracy trade-off exists in practice, similar as predicted by the theoretical framework.\n\nSTRENGTHS:\n* The work addresses an important aspect in deep RL research as most theoretical guarantees for RL are not well connected to the practical side of deep RL.\n* The work sets out to provide a theoretical framework from which to study deep RL methods.\n* It shows how such the theoretical framework can be used to understand the influence of\n  * delayed rewards and reward penalties\n  * learning rate schedules and episode lengths\n  * reward stringency\n\nWEAKNESSES:\n* The work was very difficult to follow for me. Due to the structure of the paper, many aspects seem to \"fall out of thin air\".\n* Without Appendix A it seems impossible to begin to understand Section 2 since more assumptions about the student teacher environment are given.\n* It often feels like the work requires extensive prior knowledge to be understandable.\n* Wording is often confusing. For example, in the beginning when talking about the reward for the RL-Perceptron case Fig 1 has a description about rewards that seems permissible for very dense rewards, whereas lines 48-50 talk only about extremely sparse rewards.\n\nOverall the paper seems very interesting and full of great ideas but due to a somewhat convoluted presentation and missing details that seem to be pushed to the appendix it falls short of clearly communicating these ideas.\n\nI might have missed something obvious, but to me it seems that the paper would first need fairly substantial rewriting to be easier to parse before it can be accepted.\n\nQUESTIONS:\nCould the framework be adapted to work for other RL agent types?\n\nLIMITATIONS:\nThe authors do not explicitly list limitations of their framework", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper discusses the application and theoretical understanding of Reinforcement Learning (RL) algorithms in high-dimensional settings. The authors propose a high-dimensional model of RL that can capture a variety of learning protocols and derive its dynamics as a set of closed-form ordinary differential equations.\n\nThe authors introduce the RL perceptron, a model for high-dimensional, sequential policy learning. In this model, a student network learns from a teacher network in a sequential decision-making task. The student does not observe the correct choice for each input; instead, it receives a reward that depends on whether earlier decisions are correct.\n\nThe authors derive an asymptotically exact set of ODEs that describe the typical learning dynamics of policy gradient RL agents. They use these ODEs to characterize learning behavior in a diverse range of scenarios, including exploring several sparse delayed reward schemes, deriving optimal learning rate schedules and episode length curricula, identifying ranges of learning rates for which learning is 'easy,' and 'hybrid-hard,' and identifying a speed-accuracy trade-off driven by reward stringency.\n\nThey also demonstrate that a similar speed-accuracy trade-off exists in simulations of high-dimensional policy learning from pixels using the procgen environment \"Bossfight\" and Atari \"Pong\". The authors aim to close the gap between theory and practice in high-dimensional RL.\nThe paper also discusses the sample complexity in RL, statistical learning theory for RL, and dynamics of learning, providing a comprehensive overview of the current state of RL theory and practice.\n\nSTRENGTHS:\nThis is a robust theoretical paper that enhances our comprehension of high-dimensional RL policy learning. The paper's claims are supported by high-level experimental evidence. A significant advantage of this paper is its use of the challenging procgen \"Bossfight\" environment, which is much closer to the real use cases compared to the preceding works. Additionally, the \"Pong\" game was analyzed in the supplementary materials to demonstrate the speed-accuracy tradeoff. Another strong aspect of the paper, and a significant advantage, is its focus on analyzing the average-case scenarios rather than the worst-case ones. The authors have also released the code to reproduce the results, further strengthening the paper's credibility.\n\nWEAKNESSES:\nTo the best of my knowledge, the paper does not exhibit any significant weaknesses.\n\nQUESTIONS:\n1) Do you expect any substantial changes in the main results and conclusions in the paper when bigger/deeper policies are used?\n2) What about environments with higher-dimensional action space?\n\nLIMITATIONS:\nLimitations were addressed reasonably well.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a new problem definition and class of solution methods for decision making. \n\n\nThe motivation put forward by the authors is that the current pool of solution methods do not have theoretical results that capture the neural functional classes used by practical applications in RL, leaving a gap between theory and practice. \n\nThe authors claim that the model of learning by reinforcement that they propose has sufficient flexibility to capture the same class of problems as the classic RL model, and that their problem definition, despite its generality, is solvable in closed form for higher dimensionalities (and even infinite), as opposed to the classic RL model, presumably formalised as POMDPs. In addition, they claim their model (unclear whether the authors mean the solution method or the problem, see detailed comments) behaves similarly w.r.t. hyperparameters, for which they provide optimal schedules and hypersensitivity plots. \n\nMoreover, the authors claim their proposed class of problems and solution methods exhibits in practice similar behaviour as predicted by theory, and closes a gap between theory and practice, providing as evidence empirical illustrations in a particular environment, specially designed, called “Bossfight”, which they engineer using a platform for a procedurally generated environments.\n\nSTRENGTHS:\n$\n\\textbf{Originality}$\n\nThe authors propose a new problem definition and class of solution methods for RL, particularly policy optimisation. They describe their model as a perceptron, which they call “the RL perceptron”, and analyse the learning dynamics of a particular solution method for action space of 2 discrete actions. They then define different problem instances via heuristics for various feedback signals that such a solution method can receive. The authors propose a new perspective on RL inspired by statistical mechanics and dynamical system theory.\n\n$\\textbf{Quality}$\n\nThe particular math and statistical expressions the authors derive was not checked in detail beyond  the first few simple update equations, which seem correct. The appendix was not verified for correctness. \nPlots also appear to be consistent with the experimental claims provided in text.\n\n\n$\\textbf{Clarity}$\n\nThe paper is at times clear, with some mysterious redefinitions, naming, and confusions (see in the next sections). Although the flow is cursive, the authors do not use the general breakdown of the paper and miss important sections like a “Background” section, which is placed in the introduction. It does not bother much, but clearly marking such section can help the reader understand better what is novel and what is known. \n\n\n$\\textbf{Significance}$\n\nThe main motivation of the paper is very important. Theoretical guarantees for solution methods with feasible at-scale practical implementations is highly desirable, since this informs us that they are generalizable to all problem instances, not just the particular settings in which they were tested in and shown good empirical results. Generalisation and theoretical guarantees in RL, particularly in policy-based methods with neural policy classes, which are the most successful algorithms used in practice, is a very important area of research.\n\nWEAKNESSES:\n$\\textbf{Motivation}$\n\nWhile the motivation of the paper is very important, i.e. insufficient theoretical guarantees for policy-gradient methods with neural function classes, PG have been shown to have global convergence beyond tabular fn classes which the authors claim to be the issue.  For log-linear policy classes see: Yuan  et al - “Linear Convergence of Natural Policy Gradient Methods with Log-Linear Policies” - and references within. For neural policy classes see “Neural Policy Gradient Methods: Global Optimality and Rates of Convergence” -- Wang et al and references within. Both of these show global convergence, the latter for a two-layer NN, including finite-sample guarantees/convergence rates and properties of actor-critic methods, with estimated critics, required for convergence. \n\n$\\textbf{Relation to RL and placement in context}$\n\nThe second main weakness of this paper is the lack of placement of their work in relation to the standard RL model used by the community through the formalism of MDPs or POMDPs, with accompanying class of solution methods. The paper would be significantly strengthened if they authors can clearly state what the limitations are with the previous class of problems described in RL via a reward signal and transition dynamics and in addition for POMDPs, how current definitions of high dimensional observation spaces proposed are better captured with their problem definition. It is unclear how the models proposed by this paper generalize beyond the log-linear policy class and action spaces of dimension 2, described within. The authors claim lack of theoretical work, yet previous work in RL has also analyzed distribution shift and generalization error, but with a different definition than the authors propose. It is unclear to me what the prior limitation in definition were and how this is a better way of capturing such quantities. Additional details in this respect would significantly strengthen the paper..\n\n$\\textbf{Purpose/Goal}$\n\nFurthermore, it is rather confusing the purpose of the paper, i.e. it seems the authors compare problems against each other, for the same solution method, instead of comparing solution methods that are general enough to work on every problem instance. It is also confusing the setting and problem definition. It appears that the problem is not learning by trial-and-error, and that the problem instances proposed need heuristic descriptions of reward signals, horizon sizes and task termination. It is unclear the setting in which we are in, whether that is undiscounted finite-horizon, or continuing learning (infinite horizon, average reward). The authors reference terms from RL related to this but never actually formalize the problem.\n\n$\\textbf{Empirical study/experimental illustration}$\n\nLastly, the experimental section is performed on a certain game designed in particular way, which is an illustration rather than a practical algorithmic implementation, akin to the solution methods employed in empirical RL. \n\nMore details/questions in the next sections about these points.\n\nI am happy to adjust my score if I have not correctly understood the paper, and the authors provide more details on how their work can be placed in the context of RL, which would help me understand the significance and impact of the result and analysis provided.\n\nQUESTIONS:\n$\\textbf{Confusion w.r.t. setting}$\n\nAt times it is unclear which setting the authors analyze? Is it the finite-horizon setting? There is a significant emphasis on horizon size $T$, but this in their formulation is a hyperparmeter, which a solution method optimizes, not part of the problem definition and or the objective an algorithm is trying to optimize.\n\n$\\textbf{Confusion w.r.t. the feedback signal}$\n\nAnother confusion is the feedback signal. \nThe authors propose different heuristics for defining a feedback signal, but it is unclear what their connection is with the current reward signal in the standard MDP/POMDP model. \nE.g. the number of times n, and the boolean task termination signal, the authors relate to sparsity, but it appears to also be related to the Markovian properties of the reward signal. \nIsn’t a scalar reward signal enough even for continuing problems? Since the reward rate can adequately describe the class of lifelong learning problems in RL. \n\nThe authors mention the reward rate in the related work only, but throughout the paper they have an objective which they call “generalization error” or later redefined as “test error”. Next, the reward is given by a teacher represented as a neural network, where does this come from? The paper would be considerably strengthened by a clear comparison with the current model and a description of the limitations of the current model that they are trying to solve.\n\nAt line 109 - they say receiving penalties is not always beneficial, but this should be a problem definition, not available for an agent to change.\nAt line 167-which RL protocols? finite/infinite horizons? discounted/undiscounted(continuing), partially Markovian reward signals of certain orders? \nLine 185 -”low asymptotic performance” — for a particular solution method?\nLine 186 - “finite-size effects” - unclear what this means. Changing the reward changes the objective and the RL problem is known to be non-convex anyway. It is unclear how this analysis of the stationary points relates to what is known for policy-gradient methods.\nLine 189 - “sub-reward” ?\n\n$\\textbf{Confusion on solution methods}$\n\nThe authors propose a particular model of policy gradients defined via a student network (or linear layer?) + a sigmoid activation function. It is unclear if their model/method extends to more general classes, particularly since in RL the class of solution methods extends to highly combinatorial action spaces or continuous action spaces, with log-linear, softmax, or neural function classes.\n\n$\\textbf{Confusing comparisons/hyperparameter optimization}$\n\n At times the authors compare different objective functions and say that a particular change in the reward or dynamics of the problem would yield better performance, which I find very confusing, since in general we are trying to compare the performance of different solution methods on the same problem/objective function, known or unknown. They then describe a tradeoff between problems, and how fast they are solved by the same solution method, or how well that solution method solves them (accuracy). They also appear to optimize over objectives, by doing sensitivity plots on the horizon size or the number of lives of an agent, and by comparing different feedback signals against each other claiming that one is better than another.\nLine 198 - “optimal schedules for both hyper-parameters” - isn’t this problem dependent? Why is the episode length a hyperparameter? Why is this under the control of a solution method/algorithm/agent?\n\n$\\textbf{Confusing referrals to exploration}$\n\n At times the authors mention exploration, but I am very confused how a shallow single layer neural network with a sigmoid function is dealing with exploration. It is even unclear how exploration is defined for the optimization problem described by the “generalization error” objective function they propose. Generalization over what? Observations? Tasks? Environments?\nLine 212 - “refined information” - wdym? Isn’t this changing the problem?\n\n$\\textbf{Confusion about phase transitions}$\n\nAre these transitions between different objectives? Because they change based on the problem instance.\nLine 253 “first-order phase transitions” - Not defined. I don't know what this means, I assume related to first-order stationary points of the objective function the solution method optimizes.\n\n$\\textbf{Confusion about speed-accuracy tradeoff}$\n\nIt is unclear the significance of this quantity. Doesn’t it actually tell you the speed with which different problems are solved by the same solution method that you chose? It is not a description of a solution method. \n\n$\\textbf{Clarity}$\n\nAt times the author use particular definitions without defining them, e.g., “reward stringency”, “order parameters”, “speed-accuracy” (speed of what? accuracy of what, w.r.t. what?).\n\n$\\textbf{Related work}$\n\nThe authors mention they are concerned with reward rate and episode length, but generally, the continuing setting using reward rates is the infinite horizon problem setting using average reward formulation (and thus reward rates) and has no episodes, by definition being continuing, never-ending, lifelong.\nIn the same section, the authors mention they focus on the average-case dynamics of policy-gradient methods, what does it mean “average-case”? There are reasons to use worst-case guarantees, they tell us the behaviour of an algorithm on all problem instances. Is the average case, an expectation over problems?\n\nThe authors make certain referrals to terms from other fields, like statistical mechanics or heuristics from statistical physics which may be unknown to the reader and unclear on their significance. If they are important, they should be defined.\n\n$\\textbf{Experimental section}$\n\nIsn’t the dimensionality of the observation space very small to be representative of “high dimensional” observation spaces?\nLine 292: what is a pure random policy? Does this mean this is off-policy learning from a behaviour policy that is random? Wouldn’t that mean that the gradient-update is biased?\nThe problem instance seems very particular. It is unclear how this particular problem instance reflects the generality of the claims set in the introductory contributions.\n\n$\\textbf{Miscellaneous}$\n\n\"D\" is not defined, though from context it appears to be the dimensionality of the observation space.\nTypo at line 250. Sentence doesn’t make any sense. “Learning with $\\eta_2 This is not a …”\nTypo at line 246 “the the”\n\nLIMITATIONS:\n$\\textbf{Lack of clear description of the limitations}$\n\nThe authors do not provide a clear description of the limitation of their model. It is unclear what exactly is “solvable” in closed form, is it any problem of any dimensionality with any kind of reward signal, including non Markovian? How general is the problem that is “solvable”? Does it capture all problems a POMDP would capture, and these problems are solvable? Any kind of additional information in this respect would be useful for the reader. It is also unclear what is “the average-case”. A short definition would be very useful.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-19"}