Fix for new tensorflow probability version (#39)

araffin · Mar 31, 2024 · 46dcd7f · 46dcd7f
1 parent db6120b
commit 46dcd7f
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 8 deletions.
diff --git a/sbx/common/distributions.py b/sbx/common/distributions.py
@@ -1,9 +1,8 @@
 from typing import Any, Optional
 
 import jax.numpy as jnp
-import tensorflow_probability
+import tensorflow_probability.substrates.jax as tfp
 
-tfp = tensorflow_probability.substrates.jax
 tfd = tfp.distributions
 
 

diff --git a/sbx/ppo/policies.py b/sbx/ppo/policies.py
@@ -7,15 +7,14 @@
 import jax.numpy as jnp
 import numpy as np
 import optax
-import tensorflow_probability
+import tensorflow_probability.substrates.jax as tfp
 from flax.linen.initializers import constant
 from flax.training.train_state import TrainState
 from gymnasium import spaces
 from stable_baselines3.common.type_aliases import Schedule
 
 from sbx.common.policies import BaseJaxPolicy, Flatten
 
-tfp = tensorflow_probability.substrates.jax
 tfd = tfp.distributions
 
 

diff --git a/sbx/sac/policies.py b/sbx/sac/policies.py
@@ -5,7 +5,7 @@
 import jax.numpy as jnp
 import numpy as np
 import optax
-import tensorflow_probability
+import tensorflow_probability.substrates.jax as tfp
 from flax.training.train_state import TrainState
 from gymnasium import spaces
 from stable_baselines3.common.type_aliases import Schedule
@@ -14,7 +14,6 @@
 from sbx.common.policies import BaseJaxPolicy, Flatten
 from sbx.common.type_aliases import RLTrainState
 
-tfp = tensorflow_probability.substrates.jax
 tfd = tfp.distributions
 
 

diff --git a/sbx/tqc/policies.py b/sbx/tqc/policies.py
@@ -5,7 +5,7 @@
 import jax.numpy as jnp
 import numpy as np
 import optax
-import tensorflow_probability
+import tensorflow_probability.substrates.jax as tfp
 from flax.training.train_state import TrainState
 from gymnasium import spaces
 from stable_baselines3.common.type_aliases import Schedule
@@ -14,7 +14,6 @@
 from sbx.common.policies import BaseJaxPolicy, Flatten
 from sbx.common.type_aliases import RLTrainState
 
-tfp = tensorflow_probability.substrates.jax
 tfd = tfp.distributions