danijar / dreamerv2

Mastering Atari with Discrete World Models
https://danijar.com/dreamerv2
MIT License
900 stars 195 forks source link

Code For: "Algorithm 1: Straight-Through Gradients with Automatic Differentiation" #59

Closed lucasmgomez closed 6 months ago

lucasmgomez commented 7 months ago

I am having trouble finding where the code for the Algorithm 1 described in the paper. Could you please point me to it?

danijar commented 6 months ago

https://github.com/danijar/dreamerv3/blob/a1cf96c05fc9e1e1716ff918d1da75bd05d456bf/dreamerv3/jaxutils.py#L115