(
env_per_collector,
transform_state_dict,
)
| 575 | |
| 576 | |
| 577 | def parallel_env_constructor( |
| 578 | env_per_collector, |
| 579 | transform_state_dict, |
| 580 | ): |
| 581 | if env_per_collector == 1: |
| 582 | |
| 583 | def make_t_env(): |
| 584 | env = make_transformed_env(make_env()) |
| 585 | env.transform[2].init_stats(3) |
| 586 | env.transform[2].loc.copy_(transform_state_dict["loc"]) |
| 587 | env.transform[2].scale.copy_(transform_state_dict["scale"]) |
| 588 | return env |
| 589 | |
| 590 | env_creator = EnvCreator(make_t_env) |
| 591 | return env_creator |
| 592 | |
| 593 | parallel_env = ParallelEnv( |
| 594 | num_workers=env_per_collector, |
| 595 | create_env_fn=EnvCreator(lambda: make_env()), |
| 596 | create_env_kwargs=None, |
| 597 | pin_memory=False, |
| 598 | ) |
| 599 | env = make_transformed_env(parallel_env) |
| 600 | # we call `init_stats` for a limited number of steps, just to instantiate |
| 601 | # the lazy buffers. |
| 602 | env.transform[2].init_stats(3, cat_dim=1, reduce_dim=[0, 1]) |
| 603 | env.transform[2].load_state_dict(transform_state_dict) |
| 604 | return env |
| 605 | |
| 606 | |
| 607 | # The backend can be ``gym`` or ``dm_control`` |
no test coverage detected