instadeepai / jumanji

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

fix: reward and discount spec not in wrapper #206

Closed sash-a closed 1 year ago

sash-a commented 1 year ago

The reward_spec and discount_spec methods are not in Wrapper. As such when wrapping an environment with a reward/discount spec that is not the default (shape=()), the reward and discount specs are set to the default unless they are specified again in the wrapper.

For example

class MyWrapper(Wrapper):
    def reset(self, key):
        do_some_stuff()
        return self._env.reset(key)

env = Connector()
wrapped = MyWrapper(env)

env.reward_spec() == wrapped.reward_spec()  # <-- false