Skip to content
This repository was archived by the owner on Oct 31, 2025. It is now read-only.

store numpy.array instead of DeviceArray in trajectories#386

Merged
trax-robot merged 1 commit intogoogle:masterfrom
pkol:rl_speedup
Mar 19, 2020
Merged

store numpy.array instead of DeviceArray in trajectories#386
trax-robot merged 1 commit intogoogle:masterfrom
pkol:rl_speedup

Conversation

@pkol
Copy link
Copy Markdown
Contributor

@pkol pkol commented Mar 18, 2020

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.

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.
@afrozenator
Copy link
Copy Markdown
Contributor

This is great, thanks @pkol !

@afrozenator afrozenator added the ready to pull Added when the PR is ready to be merged. label Mar 19, 2020
@trax-robot trax-robot merged commit c239229 into google:master Mar 19, 2020
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

cla: yes ready to pull Added when the PR is ready to be merged.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants