🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/1921

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (14 Unrelated Failures)

As of commit 030d0dd with merge base 799f939 ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of CPU Benchmark Tests

Total Benchmarks: 89. Improved: $\large\color{#35bf28}6$. Worsened: $\large\color{#d91a1a}5$.

$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of GPU Benchmark Tests

Total Benchmarks: 92. Improved: $\large\color{#35bf28}5$. Worsened: $\large\color{#d91a1a}9$.

It is faster, we ran multiple benchmarks on this.
It is faster in low regimes if you consider the backward pass (which is much slower when you build multiple graphs) and when you populate your optimizers with many more params (since you'll be calling step and zero grad on many more tensors with ops executed in python loops).

https://gist.github.com/vmoens/4b6037896a6a0ad347e91877ade354ae

@kfu02 @matteobettini for context: the problem with unitilialized params is that if you have modules with lazy, non initialized params it is usually assumed that you can pass them to an optimizer and the optmizer will know that it must wait until they are initialized to do smth with them.

Currently this isn't supported by from_modules which will only create dense params. We could create lazy params but I'm not super duper sure I see how that will work with vmap...

It is faster, we ran multiple benchmarks on this.

That is cool! I would still run the full multiagent training script to check those 2 things

@kfu02 @matteobettini for context: the problem with unitilialized params is that if you have modules with lazy, non initialized params it is usually assumed that you can pass them to an optimizer and the optmizer will know that it must wait until they are initialized to do smth with them.

Since lazy modules was not a feature of these classes in the first place, why don’t we just leave it that way?

i personally would prefer this to a complex solution to support it

It was a feature of CNN so I just made it uniform

@kfu02 wanna give a shot at MA-RNNs with this once it's merged? Or should I draft a PR?

@kfu02 wanna give a shot at MA-RNNs with this once it's merged? Or should I draft a PR?

Yes! I will put up a draft by the end of the week!

I ran some benchmarks on the mappo_ippo example for the MLP. With MAPPO and non-sharinmg params for 3 agents.

It seems to work! Do not see any regression in this case

All metrics match (no perf improvement tho)

Further to #1957 i ran some tests with MASAC and non-shared parameters on the same task for 4 agents and the results look good!

Almost half training time!