store numpy.array instead of DeviceArray in trajectories
store numpy.array instead of DeviceArray in trajectories

This speeds up trajectory collection because extract_info_at_index() in
t2t/.../trjaectory.py now index numpy.array instead of DeviceArray and
it's about 10000x times faster.

I've benchmarked using command:
python trax/rl_trainer.py   --config_file=trax/rl/configs/ppo_acrobot.gin   --train_batch_size=32   --output_dir=/tmp/ppo_acrobot   --alsologtostderr

Epoch time goes down from 84s to 9s per iteration.

google/trax,google/trax