[PettingZoo] Trouble running multiple MARL environments in parallel

N00bcak: <p>Hi all,</p>
<p>I’m trying to use Multi-Agent SAC to train a <a href="https://pettingzoo.farama.org/environments/sisl/multiwalker/" rel="noopener nofollow ugc">MultiWalker-v9</a> model.</p>
<p>Since I am new to MARL AND TorchRL, I loosely followed the MADDPG/MAPPO tutorials as well as the <a href="https://github.com/pytorch/rl/tree/main/sota-implementations/sac" rel="noopener nofollow ugc">SOTA implementation</a> and just swapped out the parts as necessary.</p>
<p><strong>NOTE</strong>: This is <em>not</em> a question about PettingZoo’s Parallel API, which works fine for me. I am <em>trying to run multiple copies of the PettingZoo environment</em> to speed up data collection.</p>
<p>This is my code:</p>
<pre data-code-wrap="python"><code class="lang-python">from copy import deepcopy
import tqdm
import numpy as np
from gymnasium.spaces import Box

import logging

import torch
from torch import nn

from torchrl.data.replay_buffers import TensorDictReplayBuffer
from torchrl.data.replay_buffers.samplers import RandomSampler
from torchrl.data.replay_buffers.storages import LazyMemmapStorage 
from torchrl.envs import (
    check_env_specs,
    PettingZooEnv, 
    ParallelEnv,
    GymEnv
)
from torchrl.modules import TanhNormal, AdditiveGaussianWrapper, ProbabilisticActor
from torchrl.modules.models import (
    MLP
)
from torchrl.modules.models.multiagent import (
    MultiAgentMLP,
    MultiAgentNetBase
)
from torchrl.collectors import SyncDataCollector, MultiSyncDataCollector
from torchrl.objectives import SACLoss, SoftUpdate

from tensordict import TensorDict
from tensordict.nn import TensorDictModule, TensorDictSequential, NormalParamExtractor
from torchrl.envs.utils import ExplorationType, set_exploration_type
from torchrl.envs import EnvCreator, TransformedEnv

import multiprocessing as mp

class SMACCNet(MultiAgentNetBase): 
    '''
    https://pytorch.org/rl/main/_modules/torchrl/modules/models/multiagent.html 
    This is just a more limited version of MultiAgentMLP. 
    '''
    def __init__(self, 
                n_agent_inputs: int | None,
                n_agent_outputs: int,
                n_agents: int,
                centralised: bool,
                share_params: bool,
                device = None,
                activation_class = nn.Tanh,
                **kwargs):

        self.n_agents = n_agents
        self.n_agent_inputs = n_agent_inputs
        self.n_agent_outputs = n_agent_outputs
        self.share_params = share_params
        self.centralised = centralised
        self.activation_class = activation_class

        super().__init__(
            n_agents=n_agents,
            centralised=centralised,
            share_params=share_params,
            agent_dim=-2,
            device = device,
            **kwargs,
        )
    
    def _pre_forward_check(self, inputs):
        if inputs.shape[-2] != self.n_agents:
            raise ValueError(
                f"Multi-agent network expected input with shape[-2]={self.n_agents},"
                f" but got {inputs.shape}"
            )
        # If the model is centralized, agents have full observability
        if self.centralised:
            inputs = inputs.flatten(-2, -1)
        return inputs

    def _build_single_net(self, *, device, **kwargs):
        n_agent_inputs = self.n_agent_inputs
        if self.centralised and n_agent_inputs is not None:
            n_agent_inputs = self.n_agent_inputs * self.n_agents
        
        # Note to self: This is where you change the model architecture.
        model = nn.Sequential(
            nn.Linear(n_agent_inputs, 400),
            self.activation_class(),
            nn.Linear(400, 300),
            self.activation_class(),
            nn.Linear(300, self.n_agent_outputs)
        ).to(device)

        return model

# ripped from StackOverflow
class TqdmLoggingHandler(logging.Handler):
    def __init__(self, level=logging.NOTSET):
        super().__init__(level)

    def emit(self, record):
        try:
            msg = self.format(record)
            tqdm.tqdm.write(msg)
            self.flush()
        except Exception:
            self.handleError(record)  

# Main Function
if __name__ == "__main__":
    logging.basicConfig(level = logging.INFO)
    logger = logging.getLogger(__name__)
    logger.propagate = False
    logger.addHandler(TqdmLoggingHandler())

    mp.set_start_method("spawn", force = True)
    
    NUM_CRITICS = 2
    EXPLORATION_STEPS = 100 #30000
    MAX_EPISODE_STEPS = 10 #2000
    DEVICE = "cuda"
    REPLAY_BUFFER_SIZE = 1e3 #5e5
    VALUE_GAMMA = 0.99
    BATCH_SIZE = 32 #256
    EPS = 1e-7
    LR = 3e-4
    UPDATE_STEPS_PER_BATCH = 1 #750
    WARMUP_STEPS = 10 #int(1e5)
    TRAIN_TIMESTEPS = 5000 #int(2.5e6)
    EVAL_INTERVAL = 1 #int(5e5 // EXPLORATION_STEPS)

    SEED = 42
    torch.manual_seed(SEED)
    np.random.seed(SEED)
    def env_fn(mode, parallel = True):
        base_env = PettingZooEnv(task = "multiwalker_v9", 
                                    parallel = True,
                                    seed = 42,
                                    n_walkers = 3, 
                                    shared_reward = False, 
                                    max_cycles = 1000, 
                                    render_mode = mode, 
                                    device = "cpu"
                                )

        if parallel:
            env = ParallelEnv(num_workers = 4,  # noqa: E731
                                        create_env_fn = EnvCreator(lambda: base_env), 
                                        device = "cpu",
                                        mp_start_method = "spawn",
                                        serial_for_single = True
                                    )
        else:
            env = base_env # noqa: E731

        return lambda: env

    train_env = env_fn(None)()
    if train_env.is_closed:
        train_env.start()
    eval_env = env_fn("rgb_array", parallel = False)()

    check_env_specs(train_env)

    obs_dim = train_env.full_observation_spec["walker", "observation"].shape[-1]
    action_dim = train_env.full_action_spec["walker", "action"].shape[-1]

    # NOTE: When using a ParallelEnv, you must call the attributes you want to access.
    # e.g. train_env.group_map(), not train_env.group_map
    agent_grpings = train_env.group_map()[0]
    num_agents = len(agent_grpings["walker"])
    # num_agents = len(train_env.group_map["walker"])

    # Construct the raw policy network (which acts in a decentralised manner).
    # NOTE: SAC is a probabilistic policy; here we ask that the network outputs the mean AND std
    # of the action space.
    policy_net = nn.Sequential(
                        SMACCNet(n_agent_inputs = obs_dim,
                          n_agent_outputs = 2 * action_dim, 
                          n_agents = num_agents,
                          centralised = False, # i.e. Not a joint policy.
                          share_params = True, # But agents act from the same playbook. This becomes less of a problem with communication.
                          device = DEVICE,
                          activation_class = nn.Tanh, 
                        ),
                        NormalParamExtractor(),
                    )
    
    # Construct the raw critic networks (which ARE in fact centralized).
    # As these are Q-functions, we need to add the action space too.
    # Apparently TorchRL already takes care of the multiple critic networks for you.
    critic_net = SMACCNet(n_agent_inputs = obs_dim + action_dim,
                          n_agent_outputs = 1,
                          n_agents = num_agents,
                          centralised = True, # i.e. A joint policy.
                          share_params = True,
                          device = DEVICE,
                          activation_class = nn.Tanh, 
                        )

    # TorchRL's pipelines are VERY HEAVILY focused on performing transformations on dictionaries.
    policy_net_td_module = TensorDictModule(module = policy_net,
                                            in_keys = [("walker", "observation")],
                                            # NOTE: These outputs must match with the parameter names of the 
                                            # distribution you are using!
                                            out_keys = [("walker", "loc"), ("walker", "scale")]
                                        )

    # I must confess I just copied this from the MADDPG setup.
    # However, if we are interested in using the same module code for the policy and critic networks,
    # this is probably a better solution than subclassing the MLP and saying you accept 2 inputs. 
    obs_act_module = TensorDictModule(lambda obs, act: torch.cat([obs, act], dim = -1),
                                        in_keys = [("walker", "observation"), ("walker", "action")],
                                        out_keys = [("walker", "obs_act")]
                                    )
    critic_net_td_module = TensorDictModule(module = critic_net,
                                            in_keys = [("walker", "obs_act")],
                                            out_keys = [("walker", "state_action_value")]
                                        )

    # Attach our raw policy network to a probabilistic actor
    policy_actor = ProbabilisticActor(
        module = policy_net_td_module,
        spec = train_env.full_action_spec["walker", "action"],
        in_keys = [("walker", "loc"), ("walker", "scale")],
        out_keys = [("walker", "action")],
        distribution_class = TanhNormal,
        distribution_kwargs = {
            "min": train_env.full_action_spec["walker", "action"].space.low,
            "max": train_env.full_action_spec["walker", "action"].space.high,
        },
        return_log_prob = False,
    )
    # Hopefully there are no lazy layers, but just in case, wake them up.
    fake_td = train_env.fake_tensordict().to(DEVICE)
    policy_actor(fake_td)

    # Add some extra noise to SAC.
    # In case you ask "Why are you adding noise when there is already noise sampled by the ProbabilisticActor",
    # Doing so allows us to ensure a non-negligble *minimum* action noise for our model, which can help it explore better.
    dora = AdditiveGaussianWrapper(
        policy = policy_actor,
        action_key = ("walker", "action"),
        sigma_init = 0.3,
        sigma_end = 0.1,
        annealing_num_steps = 20000
    )

    critic_actor = TensorDictSequential(
                            obs_act_module, critic_net_td_module
                        ) 
                    
    collector = SyncDataCollector( 
                    train_env,
                    policy = dora, # the explora
                    frames_per_batch = EXPLORATION_STEPS,
                    max_frames_per_traj = MAX_EPISODE_STEPS,
                    total_frames = TRAIN_TIMESTEPS,
                    device = "cpu",
                    reset_at_each_iter = False,
                )

    replay_buffer = TensorDictReplayBuffer(
        storage = LazyMemmapStorage(
            REPLAY_BUFFER_SIZE, device = "cpu",
        ),  # We will store up to memory_size multi-agent transitions
        sampler = RandomSampler(),
        batch_size = BATCH_SIZE,  # We will sample batches of this size
    )

    sac_loss = SACLoss(policy_actor, 
                        qvalue_network = critic_actor, 
                        num_qvalue_nets = 2,
                        loss_function = "l2",
                        delay_qvalue = True,
                        alpha_init = 0.1
                        )

    # Apparently TorchRL maintain their own default keybunch for the loss calculation
    # To avoid key mismatch, we must change them.
    sac_loss.set_keys(
        action = ("walker", "action"),
        state_action_value = ("walker", "state_action_value"),
        reward = ("walker", "reward"),
        done = ("walker", "done"),
        terminated = ("walker", "terminated"),
    )
    
    # TO be sure, we ARE using the Q-value-only formulation. Unfortunately, I don't know if this is necessary or not.
    # Let's leave it in for now.
    # My hunch is yes...? Because the gamma value is not yet supplied but nevertheless required for Q-learning updates.
    # https://pytorch.org/rl/stable/reference/generated/torchrl.objectives.SACLoss.html 
    # TODO: FIND OUT WHAT IS GOING ON HERE
    sac_loss.make_value_estimator(gamma = VALUE_GAMMA)

    polyak_updater = SoftUpdate(sac_loss, tau = 0.005) 

    critic_params = list(sac_loss.qvalue_network_params.flatten_keys().values())
    actor_params = list(sac_loss.actor_network_params.flatten_keys().values())

    optimizer_actor = torch.optim.Adam(
        actor_params,
        lr = LR,
        # weight_decay = 1e-4,
        eps = EPS,
    )
    optimizer_critic = torch.optim.Adam(
        critic_params,
        lr = LR,
        eps = EPS,
    )
    optimizer_alpha = torch.optim.Adam(
        [sac_loss.log_alpha],
        lr = 3e-4,
    )

    num_frames = 0
    pbar = tqdm.tqdm(total = TRAIN_TIMESTEPS)
    total_frames = 0
    ep_info = {'rewards': [], 'lengths': []}

    # Now everything is in place, write the training loop...
    breakpoint()
    for i, tensordict in enumerate(collector):

        collector.update_policy_weights_()

        pbar.update(tensordict.numel())

        tensordict = tensordict.reshape(-1)
        current_frames = tensordict.numel()
        # Add to replay buffer
        replay_buffer.extend(tensordict.cpu())
        total_frames += current_frames

        # Optimization steps
        if total_frames &gt;= WARMUP_STEPS:
            losses = TensorDict({}, batch_size=[UPDATE_STEPS_PER_BATCH])
            for i in range(UPDATE_STEPS_PER_BATCH):
                # Sample from replay buffer
                sampled_tensordict = replay_buffer.sample()
                if sampled_tensordict.device != DEVICE:
                    sampled_tensordict = sampled_tensordict.to(
                        DEVICE, non_blocking=True
                    )
                else:
                    sampled_tensordict = sampled_tensordict.clone()

                try:
                    # Compute loss
                    loss_td = sac_loss(sampled_tensordict)
                except KeyError:
                    raise Exception(f"Check {sampled_tensordict}\n{obs_act_module(sampled_tensordict)['walker', 'obs_act']}")

                actor_loss = loss_td["loss_actor"]
                q_loss = loss_td["loss_qvalue"]
                alpha_loss = loss_td["loss_alpha"]

                # Update actor
                optimizer_actor.zero_grad()
                actor_loss.backward()
                optimizer_actor.step()

                # Update critic
                optimizer_critic.zero_grad()
                q_loss.backward()
                optimizer_critic.step()

                # Update alpha
                optimizer_alpha.zero_grad()
                alpha_loss.backward()
                optimizer_alpha.step()

                losses[i] = loss_td.select(
                    "loss_actor", "loss_qvalue", "loss_alpha"
                ).detach()

                # Update qnet_target params
                polyak_updater.step()

        if not ((i + 1) % EVAL_INTERVAL):
            with set_exploration_type(ExplorationType.MODE), torch.no_grad():
                eval_rollout = eval_env.rollout(
                    MAX_EPISODE_STEPS,
                    policy_actor,
                    auto_cast_to_device=True,
                    break_when_any_done=True,
                )

                eval_reward = eval_rollout["next", "walker", "reward"].sum(-2).mean().item()
                logger.info(f"Mean Reward: {eval_reward}")
                # logger.info("SUCCESS")
            ep_reward_list = []
    collector.shutdown()
    train_env.close()
</code></pre>
<p>Running this code without <code>ParallelEnv</code> proceeds without errors. However, with <code>ParallelEnv</code>:</p>
<pre data-code-wrap="Traceback"><code class="lang-Traceback">  File "/home/n00bcak/Desktop/&lt;path_to_file&gt;/smacc.py", line 433, in &lt;module&gt;
    for i, tensordict in enumerate(collector):
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/torchrl/collectors/collectors.py", line 888, in iterator
    tensordict_out = self.rollout()
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/torchrl/_utils.py", line 480, in unpack_rref_and_invoke_function
    return func(self, *args, **kwargs)
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/torchrl/collectors/collectors.py", line 1007, in rollout
    env_output, env_next_output = self.env.step_and_maybe_reset(env_input)
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/torchrl/envs/common.py", line 2750, in step_and_maybe_reset
    tensordict_ = self.maybe_reset(tensordict_)
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/torchrl/envs/common.py", line 2795, in maybe_reset
    tensordict = self.reset(tensordict)
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/torchrl/envs/common.py", line 2120, in reset
    tensordict_reset = self._reset(tensordict, **kwargs)
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/torchrl/envs/transforms/transforms.py", line 809, in _reset
    tensordict_reset = self.base_env._reset(tensordict, **kwargs)
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/torchrl/envs/batched_envs.py", line 60, in decorated_fun
    return fun(self, *args, **kwargs)
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/torchrl/envs/batched_envs.py", line 1518, in _reset
    self.shared_tensordicts[i].apply_(
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/tensordict/base.py", line 4289, in apply_
    return self.apply(fn, *others, inplace=True, **kwargs)
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/tensordict/base.py", line 4387, in apply
    result = self._apply_nest(
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/tensordict/_td.py", line 988, in _apply_nest
    item_trsf = item._apply_nest(
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/tensordict/_td.py", line 1013, in _apply_nest
    item_trsf = fn(item, *_others)
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/lib/python3.10/site-packages/torchrl/envs/batched_envs.py", line 1515, in tentative_update
    val.copy_(other, non_blocking=self.non_blocking)
RuntimeError: output with shape [3, 1] doesn't match the broadcast shape [3, 3]
</code></pre>
<p>I have failed to pinpoint the source of the error after some time, although it appears to be due to a malformed tensordict which emerges from somewhere within torchRL’s source code.</p>
<p>Any insights on this issue are greatly appreciated.</p>
<p>P.S. I have alternatively tried to use <code>MultiSyncDataCollector</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">    collector = MultiSyncDataCollector(
                    [EnvCreator(lambda: train_env)] * 8, # train_env now ONLY uses PettingZoo's Parallel API
                    policy = dora.to('cpu'), # the explora
                    frames_per_batch = EXPLORATION_STEPS,
                    max_frames_per_traj = MAX_EPISODE_STEPS,
                    total_frames = TRAIN_TIMESTEPS,
                    device = "cpu",
                    policy_device = "cpu",
                    storing_device = "cpu",
                    env_device = "cpu",
                    reset_at_each_iter = False
                )
</code></pre>
<p>Which results in the following error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/home/n00bcak/Desktop/&lt;path_to_script&gt;/smacc.py", line 355, in &lt;module&gt;
    collector = MultiSyncDataCollector(
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/python3.10/site-packages/torchrl/collectors/collectors.py", line 1514, in __init__
    self._run_processes()
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/python3.10/site-packages/torchrl/collectors/collectors.py", line 1672, in _run_processes
    proc.start()
  File "/usr/lib/python3.10/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/usr/lib/python3.10/multiprocessing/context.py", line 224, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "/usr/lib/python3.10/multiprocessing/context.py", line 288, in _Popen
    return Popen(process_obj)
  File "/usr/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/usr/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/usr/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/usr/lib/python3.10/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/python3.10/site-packages/torch/multiprocessing/reductions.py", line 568, in reduce_storage
    fd, size = storage._share_fd_cpu_()
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/python3.10/site-packages/torch/storage.py", line 304, in wrapper
    return fn(self, *args, **kwargs)
  File "/home/n00bcak/Desktop/&lt;path_to_venv&gt;/python3.10/site-packages/torch/storage.py", line 374, in _share_fd_cpu_
    return super()._share_fd_cpu_(*args, **kwargs)
RuntimeError: _share_fd_: only available on CPU
</code></pre>
<p>Any insights on this bug are appreciated as well!</p>

matteobettini: <aside class="quote no-group" data-username="N00bcak" data-post="1" data-topic="203706">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/n00bcak/48/70407_2.png" class="avatar"> N00bcak:</div>
<blockquote>
<p><code>create_env_fn = EnvCreator(lambda: base_env)</code></p>
</blockquote>
</aside>
<p>At a first glance the problem could be stemming from here, the lambda should create the env, not just point to the existing one, otherwise multiple envs won’t be created.</p>

N00bcak: <p>Thank you. I can’t believe I didn’t notice that when debugging ^^;</p>
<p>However, changing the env_fn to:</p>
<pre data-code-wrap="python"><code class="lang-python">    def env_fn(mode, parallel = True):
        def base_env_fn():
            return PettingZooEnv(task = "multiwalker_v9", 
                                    parallel = True,
                                    seed = 42,
                                    n_walkers = 3, 
                                    shared_reward = False, 
                                    max_cycles = 1000, 
                                    render_mode = mode, 
                                    device = "cpu"
                                )

        if parallel:
            env = lambda: ParallelEnv(num_workers = 4,  # noqa: E731
                                        create_env_fn = EnvCreator(base_env_fn), 
                                        device = "cpu",
                                        mp_start_method = "spawn",
                                        serial_for_single = True
                                    )
        else:
            env = base_env_fn # noqa: E731

        return env
</code></pre>
<p>and the data collector to:</p>
<pre data-code-wrap="python"><code class="lang-python">    collector = SyncDataCollector( 
                    env_fn(None),
                    policy = dora,
                    frames_per_batch = EXPLORATION_STEPS,
                    max_frames_per_traj = MAX_EPISODE_STEPS,
                    total_frames = TRAIN_TIMESTEPS,
                    device = "cpu",
                    reset_at_each_iter = False,
                )
</code></pre>
<p>results in the same bug. I do not believe this is the cause of the RuntimeError.</p>
<p>P.S. Calling <code>EnvCreator(env_fn(None))</code> (which wraps a <code>ParallelEnv</code>) results in the following error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/home/n00bcak/Desktop/drones_go_brr/scripts/smacc.py", line 344, in &lt;module&gt;
    collector = SyncDataCollector( 
  File "/home/n00bcak/Desktop/programming/venvs/thales/lib/python3.10/site-packages/torchrl/collectors/collectors.py", line 455, in __init__
    env = create_env_fn(**create_env_kwargs)
  File "/home/n00bcak/Desktop/programming/venvs/thales/lib/python3.10/site-packages/torchrl/envs/env_creator.py", line 142, in __call__
    env.load_state_dict(self._transform_state_dict, strict=False)
  File "/home/n00bcak/Desktop/programming/venvs/thales/lib/python3.10/site-packages/torchrl/envs/batched_envs.py", line 60, in decorated_fun
    return fun(self, *args, **kwargs)
TypeError: ParallelEnv.load_state_dict() got an unexpected keyword argument 'strict'
</code></pre>
<p>This <em>seems</em> to be an unexpected bug, and not wrapping the function in <code>EnvCreator</code> resolves the error.</p>

matteobettini: <p>I see.</p>
<p>I have simplified your code to the core bug.<br>
It seems a bug in the resetting of ParallelEnv</p>
<pre data-code-wrap="python"><code class="lang-python">from torchrl.envs import PettingZooEnv, ParallelEnv, SerialEnv
from torchrl.collectors import SyncDataCollector

if __name__ == "__main__":

    def base_env_fn():
        return PettingZooEnv(
            task="multiwalker_v9",
            parallel=True,
            seed=42,
            n_walkers=3,
            shared_reward=False,
            max_cycles=1000,
            render_mode=None,
            device="cpu",
        )

    collector = SyncDataCollector(
        lambda: ParallelEnv(
            num_workers=4,  # noqa: E731
            create_env_fn=base_env_fn,
            device="cpu",
        ),
        policy=None,  # the explora
        frames_per_batch=100,
        max_frames_per_traj=50,
        total_frames=200,
        device="cpu",
        reset_at_each_iter=False,
    )
    try:
        for i, tensordict in enumerate(collector):
            print("Reached")
    finally:
        collector.shutdown()
</code></pre>
<p>I suggest opening an issue in the torchrl repo with this script and a description and we can find a solution from there</p>

vmoens: <p>I believe this will fix it</p><aside class="onebox githubpullrequest" data-onebox-src="https://github.com/pytorch/rl/pull/2181">
  <header class="source">

      <a href="https://github.com/pytorch/rl/pull/2181" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/rl</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">



    <div class="github-icon-container" title="Pull Request">
      <svg width="60" height="60" class="github-icon" viewBox="0 0 12 16" aria-hidden="true"><path fill-rule="evenodd" d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
    </div>

  <div class="github-info-container">



      <h4>
        <a href="https://github.com/pytorch/rl/pull/2181" target="_blank" rel="noopener nofollow ugc">[BugFix] Fix "_reset" assignment in batched envs</a>
      </h4>

    <div class="branches">
      <code>pytorch:main</code> ← <code>pytorch:fix-reset-batched-envs</code>
    </div>

      <div class="github-info">
        <div class="date">
          opened <span class="discourse-local-date" data-format="ll" data-date="2024-05-30" data-time="15:41:08" data-timezone="UTC">03:41PM - 30 May 24 UTC</span>
        </div>

        <div class="user">
          <a href="https://github.com/vmoens" target="_blank" rel="noopener nofollow ugc">
            <img alt="vmoens" src="https://avatars.githubusercontent.com/u/25529882?v=4" class="onebox-avatar-inline" width="20" height="20">
            vmoens
          </a>
        </div>

        <div class="lines" title="1 commits changed 2 files with 37 additions and 6 deletions">
          <a href="https://github.com/pytorch/rl/pull/2181/files" target="_blank" rel="noopener nofollow ugc">
            <span class="added">+37</span>
            <span class="removed">-6</span>
          </a>
        </div>
      </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">Solves a bug reported [here](https://discuss.pytorch.org/t/pettingzoo-trouble-ru<span class="show-more-container"><a href="https://github.com/pytorch/rl/pull/2181" target="_blank" rel="noopener nofollow ugc" class="show-more">…</a></span><span class="excerpt hidden">nning-multiple-marl-environments-in-parallel/203706/4?u=matteobettini)

```python
from torchrl.envs import PettingZooEnv, ParallelEnv, SerialEnv
from torchrl.collectors import SyncDataCollector

if __name__ == "__main__":

    def base_env_fn():
        return PettingZooEnv(
            task="multiwalker_v9",
            parallel=True,
            seed=42,
            n_walkers=3,
            shared_reward=False,
            max_cycles=1000,
            render_mode=None,
            device="cpu",
        )

    collector = SyncDataCollector(
        lambda: ParallelEnv(
            num_workers=4,  # noqa: E731
            create_env_fn=base_env_fn,
            device="cpu",
        ),
        policy=None,  # the explora
        frames_per_batch=100,
        max_frames_per_traj=50,
        total_frames=200,
        device="cpu",
        reset_at_each_iter=False,
    )
    try:
        for i, tensordict in enumerate(collector):
            print("Reached")
    finally:
        collector.shutdown()
```</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


N00bcak: <p>SFLR, I have opened <a href="https://github.com/pytorch/rl/issues/2179" rel="noopener nofollow ugc">github issues</a> <a href="https://github.com/pytorch/rl/issues/2180" rel="noopener nofollow ugc">relating to the mentioned bugs</a>.</p>
<p>In each I have pinpointed a possible source of the bug from what I can see.</p>

N00bcak: <p>I would agree with <a class="mention" href="/u/matteobettini">@matteobettini</a> on the root cause of the bug:</p>
<aside class="onebox githubpullrequest" data-onebox-src="https://github.com/pytorch/rl/pull/2181#issuecomment-2140964074">
  <header class="source">

      <a href="https://github.com/pytorch/rl/pull/2181#issuecomment-2140964074" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/rl</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">

    <div class="github-icon-container" title="Comment">
      <svg width="60" height="60" class="github-icon" viewBox="0 0 16 16" aria-hidden="true"><path fill-rule="evenodd" d="M1.5 2.75a.25.25 0 01.25-.25h8.5a.25.25 0 01.25.25v5.5a.25.25 0 01-.25.25h-3.5a.75.75 0 00-.53.22L3.5 11.44V9.25a.75.75 0 00-.75-.75h-1a.25.25 0 01-.25-.25v-5.5zM1.75 1A1.75 1.75 0 000 2.75v5.5C0 9.216.784 10 1.75 10H2v1.543a1.457 1.457 0 002.487 1.03L7.061 10h3.189A1.75 1.75 0 0012 8.25v-5.5A1.75 1.75 0 0010.25 1h-8.5zM14.5 4.75a.25.25 0 00-.25-.25h-.5a.75.75 0 110-1.5h.5c.966 0 1.75.784 1.75 1.75v5.5A1.75 1.75 0 0114.25 12H14v1.543a1.457 1.457 0 01-2.487 1.03L9.22 12.28a.75.75 0 111.06-1.06l2.22 2.22v-2.19a.75.75 0 01.75-.75h1a.25.25 0 00.25-.25v-5.5z"></path></svg>
    </div>



  <div class="github-info-container">

      <h4>
        Comment by
        <a href="https://github.com/matteobettini" target="_blank" rel="noopener nofollow ugc">
          <img alt="matteobettini" src="https://avatars.githubusercontent.com/u/55539777?u=c7a3d898fcdc8d32d7190063ab04336f3ec1d9e1&amp;v=4" class="onebox-avatar-inline" width="20" height="20">
          matteobettini
        </a> - <a href="https://github.com/pytorch/rl/pull/2181#issuecomment-2140964074" target="_blank" rel="noopener nofollow ugc">[BugFix] Fix "_reset" assignment in batched envs</a>
      </h4>



    <div class="branches">
      <code>pytorch:main</code> ← <code>pytorch:fix-reset-batched-envs</code>
    </div>

  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">I do not think the bug is in the changed code. 

"_reset" in torchrl has alway<span class="show-more-container"><a href="https://github.com/pytorch/rl/pull/2181" target="_blank" rel="noopener nofollow ugc" class="show-more">…</a></span><span class="excerpt hidden">s the same shape as "done" (https://github.com/pytorch/rl/blob/d59b8105cd7437cccb7c480f15ce6caf67f587c6/torchrl/envs/common.py#L2784). It keeps the singleton when there is the singleton.

You can check this in any env. In vmas for example, "_reset" has shape `(n_envs,1)`, just like done.

The strategy in `metadata_from_env` was correct I think. The bug regarding the different `_reset` shapes must be stemmng from elswhere</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p><a href="https://github.com/pytorch/rl/issues/2179" rel="noopener nofollow ugc">Stepping through the code</a> shows that that is the exact block where the _reset key emerges misshapen, as the wrong batch_size is taken as reference.</p>