d1a05a0dd2
* Baselines for Tensorflow 2.0. Please do note that: 1. ACER, ACKTR, GAIL is still under development by external contributors. 2. HER is still under development by tanzheny@google.com. * Some cleanup. * Addressing some comments.
11 lines
321 B
Python
11 lines
321 B
Python
from .vec_env import VecEnvObservationWrapper
|
|
|
|
|
|
class VecExtractDictObs(VecEnvObservationWrapper):
|
|
def __init__(self, venv, key):
|
|
self.key = key
|
|
super().__init__(venv=venv,
|
|
observation_space=venv.observation_space.spaces[self.key])
|
|
|
|
def process(self, obs):
|
|
return obs[self.key] |