[Feature Request] Having configurable `frames_per_batch_worker` on `_MultiDataCollector`

[Ghelfi]

Motivation

Currently, the frames_per_batch parameter is uniformly applied across all workers, which can lead to suboptimal utilisation of computational resources when dealing with heterogeneous environments per worker. One might want to build a batch with different fraction of the batch's frames per worker to account for specifics in the environments.

Solution

I propose that _MultiDataCollector exposes both frames_per_batch: int | None and frames_per_batch_worker: Sequence[int] | None. These should be mutually exclusive and in the case were frames_per_batch_worker is provided, frames_per_batchcan be computed as sum(frames_per_batch_worker).

I'll propose a PR in this direction.

Alternatives

An alternative would be to implement a new Collector handling this. This will lead to a load of code duplication for little to gain.

Additional context

Checklist

• I have checked that there is no similar issue in the repo (required)

[vmoens]

Alternatively:

frames_per_batch: int | Sequence[int] 

to avoid multiplying the kwargs?

If the async collector is used as an iterator, it should be ok because we won't stack stuff, but with the sync one it may be a problem no (shapes of the collectors won't match)?

[Ghelfi]

We can group the argument under same name to avoid kwargs multiplying for sure.

As far I made it work, this is precisely to make the output of each collector stack.
Imagine we have a MultiSyncDataCollector with 2 processes. The first process carries a ParallelEnv with 4 envs, the second process 2 envs.
You cannot concat because the number of frames_per_batch_worker is set for all the workers. It only works if all the envs in each worker produces the same number of frames.
The goal of the proposed solution is the (as an opt-in) release this constraint an enable different envs in each processes and still be able to concat because you can set the frames_per_batch_worker per worker.

[vmoens]

Ok that looks good! Looks also like you have a prototype?
How difficult do you expect this to be (not having tried it, it looks simpler than it actually is for sure 😆 )

[Ghelfi]

I have opened a PR here.
I'll do the arg packing.