From 096f4d9cf02a18b282970527450a47fb7c8b8ce0 Mon Sep 17 00:00:00 2001 From: Peter Zhokhov Date: Mon, 1 Apr 2019 15:47:13 -0700 Subject: [PATCH] neaten up stacking logic in mujoco_dset in gail --- baselines/gail/dataset/mujoco_dset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/baselines/gail/dataset/mujoco_dset.py b/baselines/gail/dataset/mujoco_dset.py index 2ada872..f5e9c27 100644 --- a/baselines/gail/dataset/mujoco_dset.py +++ b/baselines/gail/dataset/mujoco_dset.py @@ -50,7 +50,7 @@ class Mujoco_Dset(object): # obs, acs: shape (N, L, ) + S where N = # episodes, L = episode length # and S is the environment observation/action space. # Flatten to (N * L, prod(S)) - if len(obs.shape[2:]) != 0: + if len(obs.shape) > 2: self.obs = np.reshape(obs, [-1, np.prod(obs.shape[2:])]) self.acs = np.reshape(acs, [-1, np.prod(acs.shape[2:])]) else: