-
Notifications
You must be signed in to change notification settings - Fork 388
[Feature, Example] A3C Atari Implementation for TorchRL #3001
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
base: main
Are you sure you want to change the base?
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3001
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 9 Unrelated FailuresAs of commit a6eb18d with merge base 3ebe93d ( NEW FAILURE - The following job has failed:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
BROKEN TRUNK - The following jobs failed but was present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This all looks pretty good!
Could you share a (couple of) learning curve?
Another thing to do before landing is to add it to the sota-implementations CI run:
https://github.com/pytorch/rl/blob/main/.github/unittest/linux_sota/scripts/test_sota.py
Make sure the config passed there is as much barebone as we can - we just want to run the script for a couple of collection / optim iters and make sure it runs without error (not that it properly trains).
We also need to add it to the sota-check runs
Thanks @vmoens . I'll add the required changes as well as some training curves. |
@vmoens, I have added the required scripts as well. Not getting enough resources and time for hyperparam tuning to generate a proper training curve. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM, just a minor comment on the logger!
logger = get_logger( | ||
cfg.logger.backend, | ||
logger_name="a3c", | ||
experiment_name=exp_name, | ||
wandb_kwargs={ | ||
"config": dict(cfg), | ||
"project": cfg.logger.project_name, | ||
"group": cfg.logger.group_name, | ||
}, | ||
) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What I usually see is that the logger is only passed to the first worker.
Another thing is that you may want to assume that the logger isn't serializable and should be instantiated locally within the worker.
num_workers = cfg.multiprocessing.num_workers | ||
|
||
if num_workers is None: | ||
num_workers = mp.cpu_count() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
we should have way fewer workers - I think we need users to tell us how many.
data_reshape = data.reshape(-1) | ||
losses = [] | ||
|
||
mini_batches = data_reshape.split(self.mini_batch_size) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
To shuffle things a bit I usually rely on a replay buffer instance rather than just splitting the data
for local_param, global_param in zip( | ||
self.local_actor.parameters(), self.global_actor.parameters() | ||
): | ||
global_param._grad = local_param.grad | ||
|
||
for local_param, global_param in zip( | ||
self.local_critic.parameters(), self.global_critic.parameters() | ||
): | ||
global_param._grad = local_param.grad | ||
|
||
gn = torch.nn.utils.clip_grad_norm_( | ||
self.loss_module.parameters(), max_norm=max_grad_norm | ||
) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
can you explain what we do here? What do we use the _grad for?
torch.set_float32_matmul_precision("high") | ||
|
||
|
||
class SharedAdam(torch.optim.Adam): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
shouldn't we move this to the utils file?
d95de87
to
a6eb18d
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I made a few edits.
Can you explain the way the params are shared and updated? I'm not sure I see the logic
Description
Describe your changes in detail.
This PR adds an implementation of the Asynchronous Advantage Actor-Critic (A3C) algorithm for Atari environments in the torchrl/sota-implementations directory. The main files added are:
a3c_atari.py: Contains the A3C worker class, shared optimizer, and main training loop using multiprocessing.
utils_atari.py: Provides utility functions for environment creation, model construction, and evaluation, adapted for Atari tasks.
config_atari.yaml: Configuration file for hyperparameters, environment settings, and logging.
The implementation leverages TorchRL's collectors, objectives, and logging utilities, and is designed to be modular and extensible for research and benchmarking. Some of the utils functions are also borrowed from a2c_atari.
Motivation and Context
This change is required to provide a strong, reproducible baseline for A3C on Atari environments using TorchRL. It enables researchers and practitioners to benchmark and compare reinforcement learning algorithms within the TorchRL ecosystem. The implementation follows best practices for distributed RL and is compatible with TorchRL's API.
This PR solves the issue: #1755
Types of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Go over all the following points, and put an
x
in all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!