luchris429 / purejaxrl

Really Fast End-to-End Jax RL Implementations
Apache License 2.0
615 stars 53 forks source link

observation masking wrapper #8

Open esraaelelimy opened 9 months ago

esraaelelimy commented 9 months ago

Using ppo_rnn.py with cartpole might be a bit misleading since cartpole is a fully observable environment. I made this wrapper, which masks elements from the observation vector to create partially observable environments. This could be applied to environments used with ppo_rnn to make them partially observable.

To use this wrapper, you would need to add the following line: env = MaskedObservationWrapper(env,config = {'obs_idx':[0,2],'mu':0.0,'sigma':0.1}) The obs_idx list indicates which indices will be masked from the observation vector. A noise will also be added to the remaining elements of the observation vector to make the task harder.

luchris429 commented 9 months ago

Thanks! Do you think it might make more sense to just instead have it evaluate on one of the relevant bsuite or bernoulli_bandit environments from Gymnax?

esraaelelimy commented 9 months ago

There are a few partially observable options in Gymnax, like memory chain from bsuite. The idea of masking elements of observations can allow experimenting with a wide range of environments, though. We can apply the wrapper to Brax environments as well or any of the classical control environments. Basically, something similar to this benchmarking paper could be done with this simple wrapper