diff --git a/src/lerobot/envs/configs.py b/src/lerobot/envs/configs.py index 27260070e..c9c43757d 100644 --- a/src/lerobot/envs/configs.py +++ b/src/lerobot/envs/configs.py @@ -109,12 +109,17 @@ class EnvConfig(draccus.ChoiceRegistry, abc.ABC): def _make_one(): return gym.make(self.gym_id, disable_env_checker=self.disable_env_checker, **self.gym_kwargs) + extra_kwargs: dict = {} + if env_cls is gym.vector.AsyncVectorEnv: + extra_kwargs["context"] = "forkserver" try: from gymnasium.vector import AutoresetMode - vec = env_cls([_make_one for _ in range(n_envs)], autoreset_mode=AutoresetMode.SAME_STEP) + vec = env_cls( + [_make_one for _ in range(n_envs)], autoreset_mode=AutoresetMode.SAME_STEP, **extra_kwargs + ) except ImportError: - vec = env_cls([_make_one for _ in range(n_envs)]) + vec = env_cls([_make_one for _ in range(n_envs)], **extra_kwargs) return {self.type: {0: vec}} def get_env_processors(self): diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py index 77a74d60e..4a8d3ab72 100644 --- a/tests/policies/test_policies.py +++ b/tests/policies/test_policies.py @@ -31,7 +31,7 @@ from lerobot.datasets.factory import make_dataset from lerobot.datasets.feature_utils import dataset_to_policy_features from lerobot.datasets.utils import cycle from lerobot.envs.factory import make_env, make_env_config -from lerobot.envs.utils import preprocess_observation +from lerobot.envs.utils import close_envs, preprocess_observation from lerobot.optim.factory import make_optimizer_and_scheduler from lerobot.policies.act.configuration_act import ACTConfig from lerobot.policies.act.modeling_act import ACTTemporalEnsembler @@ -224,6 +224,8 @@ def test_policy(ds_repo_id, env_name, env_kwargs, policy_name, policy_kwargs): # Test step through policy env.step(action) + close_envs(envs) + # TODO(rcadene, aliberts): This test is quite end-to-end. Move this test in test_optimizer? def test_act_backbone_lr():