instadeepai / jumanji

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

feat(specs): make environment specs managed attributes #220

Closed aar65537 closed 5 months ago

aar65537 commented 7 months ago

closes #98

Edit: now also closes #221

sash-a commented 6 months ago

Sorry for the radio silence, I still want to get this in! What do you think about using cached_properties instead of properties as the spec should never change?

aar65537 commented 6 months ago

I think cached_properties should work. The only issue I foresee is you might have to access the properties in the __init__ method to ensure the arrays are initialized outside of a jitted function.

sash-a commented 6 months ago

I think cached_properties should work. The only issue I foresee is you might have to access the properties in the __init__ method to ensure the arrays are initialized outside of a jitted function.

True, but we know what the cached properties are, so we could call them in the base env classes init?

sash-a commented 6 months ago

Also just a heads up, we're making a push to get sliding tile puzzle and flatpack envs merged before the ICLR camera ready deadline (about a week). Once this is done we'll get to this PR :smile:

aar65537 commented 6 months ago

Thanks for the heads up. I changed all of the environments to use cached_property. I think it is more succinct than my original implementation. Thank you for the suggestion.

aar65537 commented 5 months ago

I've implemented the changes for the SlidingTilePuzzle and cleaned up the last references to action_spec() in the docs. That should be all the environments.