feat(binpack): improve speed of the environment

[dluo96]

Is your feature request related to a problem? Please describe

There is a potential for speedup in the BinPack environment's step method. If the computation is quite sparse, using jax.lax.map instead of jax.vmap may speed the environment up when lots of EMSs are not alive.

Describe the solution you'd like

• POC to be done with the timer.
• Faster version of BinPack's step method.

[djbyrne]

How come map will be faster than vmap?

[clement-bonnet]

Operations done on the EMSs could be conditionally done on whether the EMS exist or not (using map instead of vmap over the EMSs). This would be nice when most EMSs are non-existent. However, this operation is not parallelisable so vmap-ing the whole step function (that includes map over EMSs) will lead to inefficiency because the computation is heterogeneous. A similar solution to the VmapAutoResetWrapper could be done for the step function by splitting the step into 2 components: homogeneous computation and heterogeneous computation, vmap-ing the former and map-ing the latter. But I don't see this as being easy to implement and even having a high potential. So, I think this issue could be closed for now.