[Rllib] MultiAgentEnvRunner in episodes mode calls connectors one time too many.

[dennismalmgren]

What happened + What you expected to happen

When running a MultiAgentEnvRunner in collect episodes-mode there's an early-out termination of a for loop in _sample:

                    # Also early-out if we reach the number of episodes within this
                    # for-loop.
                    if eps == num_episodes:
                        break

This causes the code to skip the creation of a new episode but it also causes done_episodes_to_run_env_to_module to be a subset of the episodes in episodes. Normally, they would be disjoint since episodes would have re-created episodes in those slots. This causes a problem because slightly later in the call, env-to-module connectors are called, first with done_episodes_to_run_env_to_module, then with episodes:

if done_episodes_to_run_env_to_module:
                    # Run the env-to-module connector pipeline for all done episodes.
                    # Note, this is needed to postprocess last-step data, e.g. if the
                    # user uses a connector that one-hot encodes observations.
                    # Note, this pipeline run is not timed as the number of episodes
                    # can differ from `num_envs_per_env_runner` and would bias time
                    # measurements.
                    self._env_to_module(
                        episodes=done_episodes_to_run_env_to_module,
                        explore=explore,
                        rl_module=self.module,
                        shared_data=shared_data,
                        metrics=None,
                    )
                self._cached_to_module = self._env_to_module(
                    episodes=episodes,
                    explore=explore,
                    rl_module=self.module,
                    shared_data=shared_data,
                    metrics=self.metrics,
                    metrics_prefix_key=(ENV_TO_MODULE_CONNECTOR,),
                )

This causes connectors to have to deal with information they have already processed, for the done episodes.

Versions / Dependencies

We use ray[rllib]==2.44.1 but the issue seems present on Github.

Reproduction script

See description above.

Issue Severity

None

[pseudo-rnd-thoughts]

Hi @dennismalmgren, sorry for taking so long to look at this issue.
If I understand the issue then your finding that the env-to-module connector for completed episodes isn't running equivalent to the number of episodes returned?

from ray.rllib.algorithms import PPOConfig
from ray.rllib.connectors.connector_v2 import ConnectorV2
from ray.rllib.env.multi_agent_env_runner import MultiAgentEnvRunner
from ray.rllib.env.multi_agent_episode import MultiAgentEpisode
from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
from ray.rllib.utils import override


class EpisodeTracker(ConnectorV2):

    def __init__(self, input_observation_space = None, input_action_space  = None, **kwargs):
        super().__init__(input_observation_space, input_action_space, **kwargs)

        self.episode_end_counter = 0

    @override(ConnectorV2)
    def __call__(self, *, rl_module, batch, episodes: list[MultiAgentEpisode], explore, shared_data, metrics, **kwargs):
        if all(e.is_done for e in episodes):
            self.episode_end_counter += len(episodes)
        return batch


config = (
    PPOConfig()
    .environment(
        MultiAgentCartPole,
        env_config={"num_agents": 1},
    )
    .multi_agent(
        policies={"p0"},
        policy_mapping_fn=(
            lambda aid, *args, **kwargs: "p0"
        ),
    )
    .env_runners(
        num_envs_per_env_runner=5,
        env_to_module_connector=EpisodeTracker,
    )
)
env_runner = MultiAgentEnvRunner(config=config)
print('episode_end_counter pre', env_runner._env_to_module.connectors[0].episode_end_counter)
sampled_episodes = env_runner.sample(num_episodes=4, random_actions=True)
print('episode_end_counter post', env_runner._env_to_module.connectors[0].episode_end_counter)

I wrote the below script to validate that but can't reproduce the problem. Are you sampling a number of timesteps or episodes?

Could you provide a more complete script to test with?

[dennismalmgren]

I didn't notice you getting back on this. Thanks for the script @pseudo-rnd-thoughts!

Here you go:

from ray.rllib.algorithms import PPOConfig
from ray.rllib.connectors.connector_v2 import ConnectorV2
from ray.rllib.env.multi_agent_env_runner import MultiAgentEnvRunner
from ray.rllib.env.multi_agent_episode import MultiAgentEpisode
from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
from ray.rllib.utils import override

class EpisodeTracker(ConnectorV2):
    def __init__(self, input_observation_space=None, input_action_space=None, **kwargs):
        super().__init__(input_observation_space, input_action_space, **kwargs)

        self.episode_end_counter = 0
        self.episodes_encountered_list = list()
        self.episodes_encountered_set = set()

    @override(ConnectorV2)
    def __call__(self, *, rl_module, batch, episodes: list[MultiAgentEpisode], explore, shared_data, metrics, **kwargs):
        if all(e.is_done for e in episodes):
            self.episode_end_counter += len(episodes)
            for episode in episodes:
                self.episodes_encountered_list.append(episode.id_)
                self.episodes_encountered_set.add(episode.id_)        
        return batch
    
config = (
    PPOConfig()
    .environment(
        MultiAgentCartPole,
        env_config={"num_agents": 1},
    )
    .multi_agent(
        policies={"p0"},
        policy_mapping_fn=(
            lambda aid, *args, **kwargs: "p0"
        )
    )
    .env_runners(
        num_envs_per_env_runner=1,
        env_to_module_connector=EpisodeTracker,
    )
)
env_runner = MultiAgentEnvRunner(config=config)
print('episode_end_counter pre', env_runner._env_to_module.connectors[0].episode_end_counter)
sampled_episodes = env_runner.sample(num_episodes=4, random_actions=True)
print('episode_end_counter post', env_runner._env_to_module.connectors[0].episode_end_counter)

print('episodes_end_encountered_list (len)', len(env_runner._env_to_module.connectors[0].episodes_encountered_list))
      
print('episodes_end_encountered_set (len)', len(env_runner._env_to_module.connectors[0].episodes_encountered_set))

Printout:

episode_end_counter pre 0
episode_end_counter post 5
episodes_end_encountered_list (len) 5
episodes_end_encountered_set (len) 4

As you can see, the connectors are run more than once for one episode. The prints will align once you run multiple envs per env runner because this environment has varying episode lengths and often the duplicate will be processed with an incomplete episode once you run multiple environments in parallell, so the all done-guard catches it.

This check exposes it more clearly.

from ray.rllib.algorithms import PPOConfig
from ray.rllib.connectors.connector_v2 import ConnectorV2
from ray.rllib.env.multi_agent_env_runner import MultiAgentEnvRunner
from ray.rllib.env.multi_agent_episode import MultiAgentEpisode
from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
from ray.rllib.utils import override

class EpisodeTracker(ConnectorV2):
    def __init__(self, input_observation_space=None, input_action_space=None, **kwargs):
        super().__init__(input_observation_space, input_action_space, **kwargs)

        self.episode_duplicates_encountered = 0
        self.episodes_encountered_set = set()

    @override(ConnectorV2)
    def __call__(self, *, rl_module, batch, episodes: list[MultiAgentEpisode], explore, shared_data, metrics, **kwargs):
        for e in episodes:
            if e.is_done:
                if e.id_ in self.episodes_encountered_set:
                    self.episode_duplicates_encountered += 1
                self.episodes_encountered_set.add(e.id_)
        return batch
    
config = (
    PPOConfig()
    .environment(
        MultiAgentCartPole,
        env_config={"num_agents": 1},
    )
    .multi_agent(
        policies={"p0"},
        policy_mapping_fn=(
            lambda aid, *args, **kwargs: "p0"
        )
    )
    .env_runners(
        num_envs_per_env_runner=5,
        env_to_module_connector=EpisodeTracker,
    )
)
env_runner = MultiAgentEnvRunner(config=config)
print('episode_duplicates_encountered pre', env_runner._env_to_module.connectors[0].episode_duplicates_encountered)
sampled_episodes = env_runner.sample(num_episodes=4, random_actions=True)
print('episode_duplicates_encountered post', env_runner._env_to_module.connectors[0].episode_duplicates_encountered)

Printout:

episode_duplicates_encountered pre 0
episode_duplicates_encountered post 1

[pseudo-rnd-thoughts]

Thanks for the updated comment @dennismalmgren and the detailed testing script. I've created a PR #58931