# Buggy code
def Concurrently(
    ops: List[LocalIterator],
    *,
    mode="round_robin",
    output_indexes=None,
    round_robin_weights=None,
):
    """Operator that runs the given parent iterators concurrently.

    Args:
        mode (str): One of {'round_robin', 'async'}.
            - In 'round_robin' mode, we alternate between pulling items from
              each parent iterator in order deterministically.
            - In 'async' mode, we pull from each parent iterator as fast as
              they are produced. This is non-deterministic.
        output_indexes (list): If specified, only output results from the
            given ops. For example, if output_indexes=[0], only results from
            the first op in ops will be returned.
        round_robin_weights (list): List of weights to use for round robin
            mode. For example, [2, 1] will cause the iterator to pull twice
            as many items from the first iterator as the second. [2, 1, *] will
            cause as many items to be pulled as possible from the third
            iterator without blocking. This is only allowed in round robin
            mode.

        >>> sim_op = ParallelRollouts(...).for_each(...)
        >>> replay_op = LocalReplay(...).for_each(...)
        >>> combined_op = Concurrently([sim_op, replay_op], mode="async")
    """

    if len(ops) < 2:
        raise ValueError("Should specify at least 2 ops.")
    if mode == "round_robin":
        deterministic = True
    elif mode == "async":
        deterministic = False
        if round_robin_weights:
            raise ValueError("round_robin_weights cannot be specified in async mode")
    else:
        raise ValueError("Unknown mode {}".format(mode))
    if round_robin_weights and all(r == "*" for r in round_robin_weights):
        raise ValueError("Cannot specify all round robin weights = *")

    if output_indexes:
        for i in output_indexes:
            assert i in range(len(ops)), ("Index out of range", i)

        def tag(op, i):
            return op.for_each(lambda x: (i, x))

        ops = [tag(op, i) for i, op in enumerate(ops)]

    output = ops[0].union(
        *ops[1:], deterministic=deterministic, round_robin_weights=round_robin_weights
    )

    if output_indexes:
        output = output.filter(lambda tup: tup[0] in output_indexes).for_each(
            lambda tup: tup[1]
        )

    return output


# 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 Concurrently(
    ops: List[LocalIterator],
    *,
    mode="round_robin",
    output_indexes=None,
    round_robin_weights=None,
):
    """Operator that runs the given parent iterators concurrently.

    Args:
        mode (str): One of 'round_robin', 'async'. In 'round_robin' mode,
            we alternate between pulling items from each parent iterator in
            order deterministically. In 'async' mode, we pull from each parent
            iterator as fast as they are produced. This is non-deterministic.
        output_indexes (list): If specified, only output results from the
            given ops. For example, if ``output_indexes=[0]``, only results
            from the first op in ops will be returned.
        round_robin_weights (list): List of weights to use for round robin
            mode. For example, ``[2, 1]`` will cause the iterator to pull twice
            as many items from the first iterator as the second. ``[2, 1, *]``
            will cause as many items to be pulled as possible from the third
            iterator without blocking. This is only allowed in round robin
            mode.

    Examples:
        >>> sim_op = ParallelRollouts(...).for_each(...)
        >>> replay_op = LocalReplay(...).for_each(...)
        >>> combined_op = Concurrently([sim_op, replay_op], mode="async")
    """

    if len(ops) < 2:
        raise ValueError("Should specify at least 2 ops.")
    if mode == "round_robin":
        deterministic = True
    elif mode == "async":
        deterministic = False
        if round_robin_weights:
            raise ValueError("round_robin_weights cannot be specified in async mode")
    else:
        raise ValueError("Unknown mode {}".format(mode))
    if round_robin_weights and all(r == "*" for r in round_robin_weights):
        raise ValueError("Cannot specify all round robin weights = *")

    if output_indexes:
        for i in output_indexes:
            assert i in range(len(ops)), ("Index out of range", i)

        def tag(op, i):
            return op.for_each(lambda x: (i, x))

        ops = [tag(op, i) for i, op in enumerate(ops)]

    output = ops[0].union(
        *ops[1:], deterministic=deterministic, round_robin_weights=round_robin_weights
    )

    if output_indexes:
        output = output.filter(lambda tup: tup[0] in output_indexes).for_each(
            lambda tup: tup[1]
        )

    return output