# Buggy code
def ParallelRollouts(
    workers: WorkerSet, *, mode="bulk_sync", num_async=1
) -> LocalIterator[SampleBatch]:
    """Operator to collect experiences in parallel from rollout workers.

    If there are no remote workers, experiences will be collected serially from
    the local worker instance instead.

    Args:
        workers (WorkerSet): set of rollout workers to use.
        mode (str): One of {'async', 'bulk_sync', 'raw'}.
            - In 'async' mode, batches are returned as soon as they are
              computed by rollout workers with no order guarantees.
            - In 'bulk_sync' mode, we collect one batch from each worker
              and concatenate them together into a large batch to return.
            - In 'raw' mode, the ParallelIterator object is returned directly
              and the caller is responsible for implementing gather and
              updating the timesteps counter.
        num_async (int): In async mode, the max number of async
            requests in flight per actor.

    Returns:
        A local iterator over experiences collected in parallel.

    Examples:
        >>> rollouts = ParallelRollouts(workers, mode="async")
        >>> batch = next(rollouts)
        >>> print(batch.count)
        50  # config.rollout_fragment_length

        >>> rollouts = ParallelRollouts(workers, mode="bulk_sync")
        >>> batch = next(rollouts)
        >>> print(batch.count)
        200  # config.rollout_fragment_length * config.num_workers

    Updates the STEPS_SAMPLED_COUNTER counter in the local iterator context.
    """

    # Ensure workers are initially in sync.
    workers.sync_weights()

    def report_timesteps(batch):
        metrics = _get_shared_metrics()
        metrics.counters[STEPS_SAMPLED_COUNTER] += batch.count
        return batch

    if not workers.remote_workers():
        # Handle the serial sampling case.
        def sampler(_):
            while True:
                yield workers.local_worker().sample()

        return LocalIterator(sampler, SharedMetrics()).for_each(report_timesteps)

    # Create a parallel iterator over generated experiences.
    rollouts = from_actors(workers.remote_workers())

    if mode == "bulk_sync":
        return (
            rollouts.batch_across_shards()
            .for_each(lambda batches: SampleBatch.concat_samples(batches))
            .for_each(report_timesteps)
        )
    elif mode == "async":
        return rollouts.gather_async(num_async=num_async).for_each(report_timesteps)
    elif mode == "raw":
        return rollouts
    else:
        raise ValueError(
            "mode must be one of 'bulk_sync', 'async', 'raw', got '{}'".format(mode)
        )


# Traceback (KeyError)
Traceback (most recent call last):
File "/home/enes/ws/code/arl/mt/test/consume_experiences.py", line 16, in <module>
results = trainer.train()
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/rllib/agents/trainer.py", line 522, in train
raise e
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/rllib/agents/trainer.py", line 508, in train
result = Trainable.train(self)
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/tune/trainable.py", line 332, in train
result = self.step()
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/rllib/agents/trainer_template.py", line 110, in step
res = next(self.train_exec_impl)
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/util/iter.py", line 758, in __next__
return next(self.built_iterator)
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/util/iter.py", line 785, in apply_foreach
for item in it:
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/util/iter.py", line 845, in apply_filter
for item in it:
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/util/iter.py", line 845, in apply_filter
for item in it:
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/util/iter.py", line 785, in apply_foreach
for item in it:
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/util/iter.py", line 845, in apply_filter
for item in it:
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/util/iter.py", line 1078, in build_union
item = next(it)
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/util/iter.py", line 758, in __next__
return next(self.built_iterator)
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/util/iter.py", line 785, in apply_foreach
for item in it:
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/util/iter.py", line 785, in apply_foreach
for item in it:
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/util/iter.py", line 785, in apply_foreach
for item in it:
[Previous line repeated 2 more times]
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/rllib/execution/replay_ops.py", line 89, in gen_replay
item = local_buffer.replay()
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/rllib/execution/replay_buffer.py", line 331, in replay
beta=self.prioritized_replay_beta)
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/rllib/execution/replay_buffer.py", line 173, in sample
batch = self._encode_sample(idxes)
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/rllib/execution/replay_buffer.py", line 64, in _encode_sample
out = SampleBatch.concat_samples([self._storage[i] for i in idxes])
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/rllib/policy/sample_batch.py", line 93, in concat_samples
out[k] = concat_aligned([s[k] for s in samples])
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/rllib/policy/sample_batch.py", line 93, in <listcomp>
out[k] = concat_aligned([s[k] for s in samples])
File "/home/enes/ws/envs/rlws/lib/python3.7/site-packages/ray/rllib/policy/sample_batch.py", line 294, in __getitem__
return self.data[key]
KeyError: 'action_logp'

# Fixed code
def ParallelRollouts(
    workers: WorkerSet, *, mode="bulk_sync", num_async=1
) -> LocalIterator[SampleBatch]:
    """Operator to collect experiences in parallel from rollout workers.

    If there are no remote workers, experiences will be collected serially from
    the local worker instance instead.

    Args:
        workers (WorkerSet): set of rollout workers to use.
        mode (str): One of 'async', 'bulk_sync', 'raw'. In 'async' mode,
            batches are returned as soon as they are computed by rollout
            workers with no order guarantees. In 'bulk_sync' mode, we collect
            one batch from each worker and concatenate them together into a
            large batch to return. In 'raw' mode, the ParallelIterator object
            is returned directly and the caller is responsible for implementing
            gather and updating the timesteps counter.
        num_async (int): In async mode, the max number of async
            requests in flight per actor.

    Returns:
        A local iterator over experiences collected in parallel.

    Examples:
        >>> rollouts = ParallelRollouts(workers, mode="async")
        >>> batch = next(rollouts)
        >>> print(batch.count)
        50  # config.rollout_fragment_length

        >>> rollouts = ParallelRollouts(workers, mode="bulk_sync")
        >>> batch = next(rollouts)
        >>> print(batch.count)
        200  # config.rollout_fragment_length * config.num_workers

    Updates the STEPS_SAMPLED_COUNTER counter in the local iterator context.
    """

    # Ensure workers are initially in sync.
    workers.sync_weights()

    def report_timesteps(batch):
        metrics = _get_shared_metrics()
        metrics.counters[STEPS_SAMPLED_COUNTER] += batch.count
        return batch

    if not workers.remote_workers():
        # Handle the serial sampling case.
        def sampler(_):
            while True:
                yield workers.local_worker().sample()

        return LocalIterator(sampler, SharedMetrics()).for_each(report_timesteps)

    # Create a parallel iterator over generated experiences.
    rollouts = from_actors(workers.remote_workers())

    if mode == "bulk_sync":
        return (
            rollouts.batch_across_shards()
            .for_each(lambda batches: SampleBatch.concat_samples(batches))
            .for_each(report_timesteps)
        )
    elif mode == "async":
        return rollouts.gather_async(num_async=num_async).for_each(report_timesteps)
    elif mode == "raw":
        return rollouts
    else:
        raise ValueError(
            "mode must be one of 'bulk_sync', 'async', 'raw', got '{}'".format(mode)
        )