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.

Sign up for GitHub

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

Jump to bottom

[RLlib] Moving sampling coordination for batch_mode=complete_episodes to synchronous_parallel_sample. #46321

Merged
merged 5 commits into from
Jul 4, 2024
+10 −15
Merged

[RLlib] Moving sampling coordination for batch_mode=complete_episodes to synchronous_parallel_sample. #46321

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7e5ec32
Modified sampling for 'batch_mode=complete_episodes' as this was a) n…
simonsays1980 Jun 28, 2024
7a9184b
Added a comment in regard to @sven1977's review.
simonsays1980 Jun 28, 2024
41fc727
Merge branch 'master' into fix-complete-episodes-train-batch-size-in-…
simonsays1980 Jul 3, 2024
af85b4a
Removed an assertion from a callback test that was not passing b/c wi…
simonsays1980 Jul 3, 2024
8169d7b
Remvoed local mode from debugging the test:
simonsays1980 Jul 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 0 additions & 3 deletions rllib/algorithms/tests/test_callbacks_on_env_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,9 +179,6 @@ def test_episode_and_sample_callbacks_batch_mode_complete_episodes(self):

# Train one iteration.
algo.train()
# We must have had exactly one `sample()` call on our EnvRunner.
if not multi_agent:
self.assertEqual(callback_obj.counts["sample"], 1)
# We should have had at least one episode start.
self.assertGreater(callback_obj.counts["start"], 0)
# Episode starts must be exact same as episode ends (b/c we always complete
Expand Down
22 changes: 10 additions & 12 deletions rllib/env/single_agent_env_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,19 +197,17 @@ def sample(
explore=explore,
random_actions=random_actions,
)
# For complete episodes mode, sample as long as the number of timesteps
# done is smaller than the `train_batch_size`.
# For complete episodes mode, sample a single episode and
# leave coordination of sampling to `synchronous_parallel_sample`.
Copy link
Contributor

Choose a reason for hiding this comment

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

I very much like this!

Can we add a small TODO comment here that this logic, currently handled by synchronous_parallel_sample will eventually be moved fully into EnvRunnerGroup? So from the algo, you would do:

if self.config.batch_mode == "complete_episodes"
    self.env_runner_group.sample(num_timesteps=[batch size], complete_episodes=True)

something like this ^. Don't have to do this in this PR!

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Awesome. I would love this move!

# TODO (simon, sven): The coordination will eventually move
# to `EnvRunnerGroup` in the future. So from the algorithm one
# would do `EnvRunnerGroup.sample()`.
else:
total = 0
samples = []
while total < self.config.train_batch_size:
episodes = self._sample_episodes(
num_episodes=self.num_envs,
explore=explore,
random_actions=random_actions,
)
total += sum(len(e) for e in episodes)
samples.extend(episodes)
samples = self._sample_episodes(
num_episodes=1,
explore=explore,
random_actions=random_actions,
)

# Make the `on_sample_end` callback.
self._callbacks.on_sample_end(
Expand Down
Loading