[Feature] Use ObservationNorm.init_stats for stats computation in example scripts #715
Conversation
facebook-github-bot
added
the
CLA Signed
vmoens
left a comment
vmoens
left a comment
There was a problem hiding this comment.
I think that here, if we spawn multiple processes, each env on each process will run its own init_env_steps and hence they will all have a different set of summary stats.
We should compute the stats in the main process like we did before (but using the method you provided), then pass these stats to each env.
…tionNorm transforms. Add checks in init_stats to ensure proper initialization
vmoens
left a comment
vmoens
left a comment
There was a problem hiding this comment.
LGTM
Do you think things would be simpler with state dict?
We could
(1) create a dummy env, compute stats
(2) get the state dict of this dummy env
(3) load the state dict on every env created subsequently
LMK what you think
torchrl/trainers/helpers/envs.py
Outdated
| raise AttributeError("init_env_steps missing from arguments.") | ||
|
|
||
| if ( | ||
| type(proof_environment.transform) != Compose |
There was a problem hiding this comment.
Why not isinstance?
Why equality and not is not?
torchrl/trainers/helpers/envs.py
Outdated
|
|
||
| if ( | ||
| type(proof_environment.transform) != Compose | ||
| and type(proof_environment.transform) != ObservationNorm |
torchrl/trainers/helpers/envs.py
Outdated
| ) | ||
|
|
||
| obs_norm_transforms = [] | ||
| if type(proof_environment.transform) == Compose: |
torchrl/trainers/helpers/envs.py
Outdated
| obs_norm_transforms.append((0, proof_environment.transform)) | ||
|
|
||
| stats = [] | ||
| for (idx, transform) in obs_norm_transforms: |
There was a problem hiding this comment.
Upon reflection we could simply take the state dict of the transforms and load it no?
Wouldn't that be simpler than this?
If loc and scale are buffers it could simplify things a bit
romainjln
force-pushed
the
init-stats-models
branch
from
acf16b8 to
17b1b15
Compare
Codecov Report
@@ Coverage Diff @@
## main #715 +/- ##
==========================================
- Coverage 88.71% 88.66% -0.05%
==========================================
Files 120 120
Lines 20240 20386 +146
==========================================
+ Hits 17955 18075 +120
- Misses 2285 2311 +26
Flags with carried forward coverage won't be shown. Click here to find out more.
📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more |
torchrl/trainers/helpers/envs.py
Outdated
|
|
||
|
|
||
| def retrieve_observation_norms_state_dict(proof_environment: TransformedEnv): | ||
| """Traverse the transforms of the environment and retrieve the ObservationNorm state dicts. |
There was a problem hiding this comment.
Maybe Traverses + retrieves?
Also for code objects we can use :obj:ObservationNorm for a cleaner rendering in the doc
torchrl/trainers/helpers/envs.py
Outdated
| num_iter: int = 1000, | ||
| key: Union[str, Tuple[str, ...]] = None, | ||
| ): | ||
| """Calling init_stats on all uninitialised ObservationNorms transform of a TransformedEnv. |
There was a problem hiding this comment.
Maybe
Calls :obj:`ObservationNorm.init_stats` on all uninitialized :obj:`ObservationNorm` instances of a :obj:`TransformedEnv`.
torchrl/trainers/helpers/envs.py
Outdated
| ): | ||
| """Calling init_stats on all uninitialised ObservationNorms transform of a TransformedEnv. | ||
| If an ObservationNorm already has non-null loc or stats, it will be skipped. |
There was a problem hiding this comment.
If an :obj:`ObservationNorm.init_stats` already has non-null loc or stats, a call to :obj:`initialize_observation_norm_transforms` will be a no-op.
Similarly, if the transformed environment does not contain any ObservationNorm, a call to this function will have no effect.
3 participants
Description
Refactor examples to use
ObservationNorm.init_statsinstead ofget_stats_random_rolloutMotivation and Context
close #699
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
xin all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!