Denys88/rl_games

wrapper flatten issue

Closed this issue · 0 comments

sry no time for PR but in common/wrappers.py it should be in BatchedFrameStack

    def _get_ob(self):
        assert len(self.frames) == self.k
        if self.transpose:
            frames = np.transpose(self.frames, (1, 2, 0))
        else:
            if self.flatten:
                frames = np.array(self.frames)
                shape = np.shape(frames)
                frames = np.transpose(frames, (1, 0, 2))
                frames = np.reshape(frames, (shape[1], shape[0] * shape[2]))
            else:
                frames = np.transpose(self.frames, (1, 0, 2))
        return frames