There is a lot of detail there, or in any RL implementation for that matter. Taking a quick look, there are a number of quirks and ? in your impl. Without digging deep into slicing and off by one level detail, a few thoughts:


It’s a great learning experience to start from scratch, but there is a lot to get right, and just a small mistake to get wrong and nothing will work. I’d spend a little more time looking at other (working) reference impl to guide you.


You’re calling zero grad once for your whole inner loop. That is likely why you get the retain_graph warning. Doing one zero and multiple backward, and one step is a way of accumulating your gradients before an update. What you’re doing is all kinds of crazy  You should generally do the following within the loop:


optimizer.zero_grad()
loss.backward()
# most RL algorithms also do a nn.utils.clip_grad_norm_(model.parameters(), self.max_grad_norm)
optimizer.step()



For PPO specifically, your notion of ‘batches’ and samples in general is not typical. Usually you build up more trajectories from your simulator and then sample larger batches from them multiple times (PPO epochs).


From outer loop to outer loop the states (observations) aren’t being rolled over. You leave next state in states[steps+1], but you don’t bring that last observation into states[0] for the start of the next. Unless I’m missing something you’re losing all continuity in your observations from one outer loop to another, given that the environment episodes don’t necessarily line up with your rollout sampling.