fix: change default float format from tf32 to fp32 (#12)

Author: typoverflow

examples/online/main_dmc_offpolicy.py

@@ -2,6 +2,7 @@
 
 import gymnasium as gym
 import hydra
+import jax
 import jax.numpy as jnp
 import numpy as np
 import omegaconf
@@ -17,6 +18,8 @@
 from flowrl.utils.logger import CompositeLogger
 from flowrl.utils.misc import set_seed_everywhere
 
+jax.config.update("jax_default_matmul_precision", "float32")
+
 SUPPORTED_AGENTS: Dict[str, BaseAgent] = {
     "sac": SACAgent,
     "td3": TD3Agent,

examples/online/main_mujoco_offpolicy.py

@@ -3,6 +3,7 @@
 import gymnasium as gym
 import gymnasium_robotics
 import hydra
+import jax
 import numpy as np
 import omegaconf
 import wandb
@@ -16,6 +17,8 @@
 from flowrl.utils.logger import CompositeLogger
 from flowrl.utils.misc import set_seed_everywhere
 
+jax.config.update("jax_default_matmul_precision", "float32")
+
 SUPPORTED_AGENTS: Dict[str, BaseAgent] = {
     "sac": SACAgent,
     "td3": TD3Agent,

examples/online/main_mujoco_onpolicy.py

@@ -4,6 +4,7 @@
 import gymnasium as gym
 import gymnasium_robotics
 import hydra
+import jax
 import jax.numpy as jnp
 import numpy as np
 import wandb
@@ -16,6 +17,8 @@
 from flowrl.utils.logger import CompositeLogger
 from flowrl.utils.misc import set_seed_everywhere
 
+jax.config.update("jax_default_matmul_precision", "float32")
+
 SUPPORTED_AGENTS: Dict[str, BaseAgent] = {
     "ppo": PPOAgent,
 }