Fix indexing LazyFrames

[wwiiiii]

In current version, indexing i-th element of LazyFrame return wrong result.
It stores the stacked frame in (height, width, channel) format but using first axis -which is height- when indexing, therefore return value has (width, channel) shape.

---

from baselines.common.atari_wrappers import make_atari, wrap_deepmind

env = wrap_deepmind(make_atari('PongNoFrameskip-v4'), frame_stack=True)

state = env.reset()
print(type(state))
print(state[0].shape)

The result of above snippet is

<class 'baselines.common.atari_wrappers.LazyFrames'>
(84, 4)

for now, but it would be

<class 'baselines.common.atari_wrappers.LazyFrames'>
(84, 84)

after the fix.

[pzhokhov]

see discussion here: #907 - the snippet above should now change into

from baselines.common.atari_wrappers import make_atari, wrap_deepmind
env = wrap_deepmind(make_atari('PongNoFrameskip-v4'), frame_stack=True)
state = env.reset()
print(type(state))
print(state.frame(0).shape)

this way, the behaviour of LazyFrames is consistent with numpy arrays.