Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adversarial algorithm matching original paper's implementation #770

Draft
wants to merge 56 commits into
base: master
Choose a base branch
from

Conversation

taufeeque9
Copy link
Collaborator

@taufeeque9 taufeeque9 commented Aug 10, 2023

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.

taufeeque9 and others added 30 commits January 5, 2023 01:49
This change just made some error messages go away indicating the missing imitation.algorithms.dagger.ExponentialBetaSchedule but it did not fix the root cause.
@codecov
Copy link

codecov bot commented Aug 10, 2023

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 ernestum self-assigned this Aug 28, 2023
Copy link
Collaborator

@ernestum ernestum left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just some nits. Throughout review will follow.

assert buffer.actions is not None
obs = buffer.observations
next_obs = obs[1:]
next_obs = np.concatenate([next_obs, obs[-1:]], axis=0) # last obs not available
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Easier to read if you do:

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

next_obs = np.concatenate([next_obs, obs[-1:]], axis=0) # last obs not available
actions = buffer.actions
dones = buffer.episode_starts
dones = np.roll(dones, -1, axis=0)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@@ -222,16 +279,22 @@ def __init__(

self.venv_buffering = wrappers.BufferingWrapper(self.venv)

self.disc_trainer_callback = TrainDiscriminatorCallback(self)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants