Put back normalization

Farama-Foundation · Feb 19, 2024 · 513c7be · 513c7be
1 parent f0a092f
commit 513c7be
Show file tree

Hide file tree

Showing 2 changed files with 0 additions and 2 deletions.
diff --git a/momaland/learning/cooperative_momappo/continuous_momappo.py b/momaland/learning/cooperative_momappo/continuous_momappo.py
@@ -552,7 +552,6 @@ def _env_step(runner_state):
     critic = Critic(net_arch=args.critic_net_arch, activation=args.activation)
     vmapped_get_value = vmap(critic.apply, in_axes=(None, 0))
     critic.apply = jax.jit(critic.apply)
-    actor.apply = jax.jit(actor.apply)
 
     if args.track:
         exp_name = args.exp_name

diff --git a/momaland/learning/cooperative_momappo/discrete_momappo.py b/momaland/learning/cooperative_momappo/discrete_momappo.py
@@ -548,7 +548,6 @@ def _env_step(runner_state):
     critic = Critic(net_arch=args.critic_net_arch, activation=args.activation)
     vmapped_get_value = vmap(critic.apply, in_axes=(None, 0))
     critic.apply = jax.jit(critic.apply)
-    actor.apply = jax.jit(actor.apply)
 
     if args.track:
         exp_name = args.exp_name