1 parent 67fb4b4 commit ae98625Copy full SHA for ae98625
1 file changed
dreamerv3/agent.py
@@ -2,6 +2,7 @@
2
import jax
3
import jax.numpy as jnp
4
import ruamel.yaml as yaml
5
+from pathlib import Path
6
tree_map = jax.tree_util.tree_map
7
sg = lambda x: tree_map(jax.lax.stop_gradient, x)
8
0 commit comments