nicklashansen / tdmpc2

Code for "TD-MPC2: Scalable, Robust World Models for Continuous Control"
https://www.tdmpc2.com
MIT License
272 stars 49 forks source link

[Feature] Faster replay buffer + support pixel observations #10

Closed nicklashansen closed 6 months ago

nicklashansen commented 6 months ago

This PR adds support for pixel observations in DMControl tasks, and replaces the previous replay buffer with a faster implementation, resulting in approx 1.5x faster training times for state observations and 8x for pixel observations. This PR also adds other minor quality-of-life improvements such as the ability to only install dependencies for environments that one wishes to use (rather than all 4 environments).