instadeepai / jumanji

🕹️ A diverse suite of scalable reinforcement learning environments in JAX
https://instadeepai.github.io/jumanji
Apache License 2.0
645 stars 80 forks source link

Support for Multi-Dimensional Discrete Spec #257

Closed thomashirtz closed 2 weeks ago

thomashirtz commented 2 weeks ago

Is your feature request related to a problem? Please describe.

Currently, DiscreteArray only supports scalar values, making it challenging to define multi-dimensional discrete action spaces with specified shapes.

Describe the solution you'd like

A discrete spec similar to BoundedArray that accepts a shape parameter and allows each element within the specified shape to take on discrete values. This would enable easier handling of grid-based or multi-dimensional discrete action spaces.

Misc

clement-bonnet commented 2 weeks ago

Hi Thomas, Would the MultiDiscreteArray spec work for your use case? See e.g. how it is used in BinPack: https://github.com/instadeepai/jumanji/blob/main/jumanji/environments/packing/bin_pack/env.py#L264 Let me know if this helps, Clément

thomashirtz commented 2 weeks ago

Oh, I didn't see that you use the shape of num_values, I thought it was a scalar. Yes that would fit my need I believe Thanks for the quick help !