Skip to content

Commit

Permalink
add flatten dict wrapper when observation space is dict (hack, don't …
Browse files Browse the repository at this point in the history 
…merge)
  • Loading branch information
@pzhokhov
pzhokhov committed Nov 17, 2018
1 parent d3fed18 commit 28b0567
Showing 1 changed file with 2 additions and 0 deletions.
2 changes: 2 additions & 0 deletions baselines/common/cmd_util.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,8 @@ def make_env(env_id, env_type, subrank=0, seed=None, reward_scale=1.0, gamestate
env = retro_wrappers.make_retro(game=env_id, max_episode_steps=10000, use_restricted_actions=retro.Actions.DISCRETE, state=gamestate)
else:
env = gym.make(env_id)
if isinstance(env.observation_space, gym.spaces.Dict):
env = gym.wrappers.FlattenDictWrapper(env, dict_keys=['observation', 'desired_goal'])

env.seed(seed + subrank if seed is not None else None)
env = Monitor(env,
Expand Down

2 comments on commit 28b0567

@WuXinyang2012
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for your work and kind reply.

@gautams3
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would be more generic that specifying keys yourself
env = gym.wrappers.FlattenDictWrapper(env, dict_keys=env.observation_space.spaces.keys())

Please sign in to comment.