Adversarial algorithm matching original paper's implementation

[taufeeque9]

Description

This PR updates the adversarial algorithm by training the discriminator between collecting the rollouts of the generator and training the generator. This matches the reference implementation provided in Algorithm 1 of the AIRL paper.

The modification is done by implementing the TrainDiscriminatorCallback, which is called to train the discriminator after collecting rollouts through the callback.on_rollout_end(). The callback first stores the latest rollout in the replay buffer, which is then used to train the discriminator. Once the discriminator is trained, the callback updates the generator's rollout/replay buffer by updating the rewards using the latest discriminator.

Note that we must also update the advantages and returns in the rollout buffer of the on-policy algorithms upon updating the rewards. This is tricky to do since information like value and done on the last observations of the rollouts is not stored in the rollout buffer. These are obtained in this PR by using the original advantages and rewards. A test of whether it produces correct values needs to be added.

Testing

All the tests for adversarial algorithms run successfully.

• The performance of this implementation remains to be compared to the other re-implementation Change adversarial algorithms to collect rollouts first #731 and the master branch's algorithm.

[codecov]

Codecov Report

Merging #770 (5c23650) into master (19c7f35) will increase coverage by 0.03%.
Report is 2 commits behind head on master.
The diff coverage is 98.68%.

@@            Coverage Diff             @@
##           master     #770      +/-   ##
==========================================
+ Coverage   96.33%   96.37%   +0.03%     
==========================================
  Files          93       93              
  Lines        8789     8846      +57     
==========================================
+ Hits         8467     8525      +58     
+ Misses        322      321       -1     

Files Changed	Coverage Δ
src/imitation/algorithms/adversarial/common.py	97.68% <98.14%> (+0.85%)	⬆️
src/imitation/policies/replay_buffer_wrapper.py	100.00% <100.00%> (ø)
tests/algorithms/test_adversarial.py	100.00% <100.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[ernestum]

Just some nits. Throughout review will follow.

[ernestum]

Easier to read if you do:

ext_obs = np.concatenate([obs[1:], obs[-1:]], axis=0)

[ernestum]

same as above: pull buffer.episode_starts in this line.

[ernestum]

Why not define the gen_callback here like this:

self.gen_callback: List[callbacks.BaseCallback] = [self.disc_trainer_callback]

and then just append it down in the else block?

And while you are at it rename it to use a plural because it is actually more than one callback. E.g. self.gen_callbacks.