Improve PPO readability

[alan-cooney]

Adds mostly comments and renames a few variables to be more descriptive. Also deletes two unused methods.

I've run an example with wandb to check this doesn't change anything unexpectedly - https://wandb.ai/alancooney/trlx/runs/aew4nu69 (same as before the changes)

Co-authored-by: @jezgillen

[alan-cooney]

@reciprocated to help with documenting this code, are you able to confirm what for _ in range(self.n_updates_per_batch): is intended to do in the base trainer?

trlx/trlx/trainer/accelerate_base_trainer.py

Lines 445 to 449 in 84a0711

	for _ in range(self.config.train.epochs):
	for batch in self.train_dataloader:
	for _ in range(self.n_updates_per_batch):
	forward_time = time()

Git blame suggests this is your code. My understanding is it runs the same batch multiple times - is that as expected? That is to say it lets the user repeat a batch, rather than e.g. having multiple steps e.g. turns in a text-based game (which I don't think the library supports).

[maxreciprocate]

Hey there! n_updates_per_batch let's perform multiple gradient updates on the same batch of data, this is a common PPO implementation detail. Here is an excerpt from PPO's paper abstract [1] as reference:

Whereas standard policy gradient methods perform one gradient update per data sample, we propose a novel objective function that enables multiple epochs of minibatch updates

[1] https://arxiv.org/abs/1707.06347

[alan-cooney]

Whereas standard policy gradient methods perform one gradient update per data sample, we propose a novel objective function that enables multiple epochs of minibatch updates

Awesome thanks for the quick reply! I'll add this in to the comments

[maxreciprocate]

Splendid, what's left is to merge changes from the main and do a pre-commit run and I've also left a small suggestion to remain consistent in naming throughout. Thanks for the effort @alan-cooney @jezgillen

[LouisCastricato]

Let's get this merged

[alan-cooney]

Thanks for the review @reciprocated

Note I've had to resolve a tonne of merge conflicts, so it's worth checking those are in line with expectations (they're primarily with your commits).

Given the size of this PR please can you take a look as soon as possible & merge, otherwise we'll likely get other merge conflicts as well.

[maxreciprocate]

Marvellous, merging!

https://wandb.ai/sorry/trlx/reports/Improve-PPO-readability-210---VmlldzozNDg3ODEx

[alan-cooney]

Thanks for the quick review!