Merge pull request #232 from LisaCoiffard/refactor/remove-jit-decorators

Remove jax.jit decorators

Author: LisaCoiffard

examples/aurora.ipynb

@@ -452,6 +452,7 @@
     "n_target = 1024\n",
     "\n",
     "previous_error = jnp.sum(repertoire.fitnesses != -jnp.inf) - n_target\n",
+    "container_size_control_fn = jax.jit(aurora.container_size_control)\n",
     "\n",
     "iteration = 0\n",
     "while iteration < max_iterations:\n",
@@ -480,7 +481,7 @@
     "        )\n",
     "\n",
     "    elif iteration % 2 == 0:\n",
-    "        repertoire, previous_error = aurora.container_size_control(\n",
+    "        repertoire, previous_error = container_size_control_fn(\n",
     "            repertoire,\n",
     "            target_size=n_target,\n",
     "            previous_error=previous_error,\n",

examples/cmaes.ipynb

@@ -208,18 +208,21 @@
     "covs = [(state.sigma**2) * state.cov_matrix]\n",
     "\n",
     "iteration_count = 0\n",
+    "sample_fn = jax.jit(cmaes.sample)\n",
+    "update_fn = jax.jit(cmaes.update)\n",
+    "stop_condition_fn = jax.jit(cmaes.stop_condition)\n",
     "for _ in range(num_iterations):\n",
     "    iteration_count += 1\n",
     "\n",
     "    # sample\n",
     "    key, subkey = jax.random.split(key)\n",
-    "    samples = cmaes.sample(state, subkey)\n",
+    "    samples = sample_fn(state, subkey)\n",
     "\n",
     "    # update\n",
-    "    state = cmaes.update(state, samples)\n",
+    "    state = update_fn(state, samples)\n",
     "\n",
     "    # check stop condition\n",
-    "    stop_condition = cmaes.stop_condition(state)\n",
+    "    stop_condition = stop_condition_fn(state)\n",
     "\n",
     "    if stop_condition:\n",
     "        break\n",

examples/mapelites_asktell.ipynb

@@ -327,20 +327,23 @@
     "except ImportError:\n",
     "    bar = range(num_iterations)\n",
     "\n",
+    "ask_fn = jax.jit(map_elites.ask)\n",
+    "tell_fn = jax.jit(map_elites.tell)\n",
+    "\n",
     "# Main loop\n",
     "for i in bar:\n",
     "    start_time = time.time()\n",
     "    key, subkey = jax.random.split(key)\n",
     "    # Generate solutions\n",
-    "    genotypes, extra_info = map_elites.ask(repertoire, emitter_state, subkey)\n",
+    "    genotypes, extra_info = ask_fn(repertoire, emitter_state, subkey)\n",
     "\n",
     "    # Evaluate solutions: get fitness, descriptor and extra scores.\n",
     "    # This is where custom evaluations on CPU or GPU can be added.\n",
     "    key, subkey = jax.random.split(key)\n",
     "    fitnesses, descriptors, extra_scores = scoring_fn(genotypes, subkey)\n",
     "\n",
     "    # Update MAP-Elites\n",
-    "    repertoire, emitter_state, current_metrics = map_elites.tell(\n",
+    "    repertoire, emitter_state, current_metrics = tell_fn(\n",
     "        genotypes=genotypes,\n",
     "        fitnesses=fitnesses,\n",
     "        descriptors=descriptors,\n",

examples/mees.ipynb

@@ -248,11 +248,11 @@
     "\n",
     "# Prepare the scoring functions for the offspring generated following\n",
     "# the approximated gradient (each of them is evaluated 30 times)\n",
-    "sampling_fn = functools.partial(\n",
+    "sampling_fn = jax.jit(functools.partial(\n",
     "    sampling,\n",
     "    scoring_fn=scoring_fn,\n",
     "    num_samples=30,\n",
-    ")\n",
+    "))\n",
     "\n",
     "# Get minimum reward value to make sure qd_score are positive\n",
     "reward_offset = environments.reward_offset[env_name]\n",
@@ -448,11 +448,8 @@
    "provenance": []
   },
   "gpuClass": "standard",
-  "interpreter": {
-   "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
-  },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": ".venv",
    "language": "python",
    "name": "python3"
   },
@@ -466,7 +463,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.11.10"
   }
  },
  "nbformat": 4,

examples/pga_aurora.ipynb

@@ -514,6 +514,7 @@
     "n_target = 1024\n",
     "\n",
     "previous_error = jnp.sum(repertoire.fitnesses != -jnp.inf) - n_target\n",
+    "container_size_control_fn = jax.jit(aurora.container_size_control)\n",
     "\n",
     "iteration = 0\n",
     "while iteration < max_iterations:\n",
@@ -542,7 +543,7 @@
     "        )\n",
     "\n",
     "    elif iteration % 2 == 0:\n",
-    "        repertoire, previous_error = aurora.container_size_control(\n",
+    "        repertoire, previous_error = container_size_control_fn(\n",
     "            repertoire,\n",
     "            target_size=n_target,\n",
     "            previous_error=previous_error,\n",

qdax/baselines/cmaes.py

@@ -3,7 +3,6 @@
 a CMA optimization script. Link to the paper: https://arxiv.org/abs/1604.00772
 """
 
-from functools import partial
 from typing import Callable, Optional, Tuple
 
 import flax
@@ -165,7 +164,6 @@ def init(self) -> CMAESState:
             invsqrt_cov=invsqrt_cov,
         )
 
-    @partial(jax.jit, static_argnames=("self",))
     def sample(self, cmaes_state: CMAESState, key: RNGKey) -> Genotype:
         """
         Sample a population.
@@ -186,7 +184,6 @@ def sample(self, cmaes_state: CMAESState, key: RNGKey) -> Genotype:
         )
         return samples
 
-    @partial(jax.jit, static_argnames=("self",))
     def update_state(
         self,
         cmaes_state: CMAESState,
@@ -198,7 +195,6 @@ def update_state(
             weights=self._weights,
         )
 
-    @partial(jax.jit, static_argnames=("self",))
     def update_state_with_mask(
         self, cmaes_state: CMAESState, sorted_candidates: Genotype, mask: Mask
     ) -> CMAESState:
@@ -217,7 +213,6 @@ def update_state_with_mask(
             weights=weights,
         )
 
-    @partial(jax.jit, static_argnames=("self",))
     def _update_state(
         self,
         cmaes_state: CMAESState,
@@ -332,7 +327,6 @@ def update_eigen(
 
         return cmaes_state
 
-    @partial(jax.jit, static_argnames=("self",))
     def update(self, cmaes_state: CMAESState, samples: Genotype) -> CMAESState:
         """Updates the distribution.
 
@@ -352,7 +346,6 @@ def update(self, cmaes_state: CMAESState, samples: Genotype) -> CMAESState:
 
         return new_state  # type: ignore
 
-    @partial(jax.jit, static_argnames=("self",))
     def stop_condition(self, cmaes_state: CMAESState) -> bool:
         """Determines if the current optimization path must be stopped.
 

qdax/baselines/dads.py

@@ -4,7 +4,6 @@
 """
 
 from dataclasses import dataclass
-from functools import partial
 from typing import Callable, Tuple
 
 import jax
@@ -191,7 +190,6 @@ def init(  # type: ignore
             steps=jnp.array(0),
         )
 
-    @partial(jax.jit, static_argnames=("self",))
     def _compute_diversity_reward(
         self, transition: QDTransition, training_state: DadsTrainingState
     ) -> Reward:
@@ -244,8 +242,7 @@ def _compute_diversity_reward(
 
         return reward
 
-    @partial(jax.jit, static_argnames=("self", "env", "deterministic", "evaluation"))
-    def play_step_fn(
+    def play_step_fn(  # type: ignore
         self,
         env_state: EnvState,
         training_state: DadsTrainingState,
@@ -339,14 +336,13 @@ def play_step_fn(
 
         return next_env_state, training_state, transition
 
-    @partial(jax.jit, static_argnames=("self", "play_step_fn", "env_batch_size"))
-    def eval_policy_fn(
+    def eval_policy_fn(  # type: ignore
         self,
         training_state: DadsTrainingState,
         eval_env_first_state: EnvState,
         play_step_fn: Callable[
-            [EnvState, Params, RNGKey],
-            Tuple[EnvState, Params, RNGKey, QDTransition],
+            [EnvState, Params],
+            Tuple[EnvState, Params, QDTransition],
         ],
         env_batch_size: int,
     ) -> Tuple[Reward, Reward, Reward, StateDescriptor]:
@@ -400,7 +396,6 @@ def eval_policy_fn(
 
         return true_return, true_returns, diversity_returns, transitions.state_desc
 
-    @partial(jax.jit, static_argnames=("self",))
     def _compute_reward(
         self, transition: QDTransition, training_state: DadsTrainingState
     ) -> Reward:
@@ -417,7 +412,6 @@ def _compute_reward(
             transition=transition, training_state=training_state
         )
 
-    @partial(jax.jit, static_argnames=("self",))
     def _update_dynamics(
         self, operand: Tuple[DadsTrainingState, QDTransition]
     ) -> Tuple[Params, float, optax.OptState]:
@@ -448,7 +442,6 @@ def _update_dynamics(
             dynamics_optimizer_state,
         )
 
-    @partial(jax.jit, static_argnames=("self",))
     def _not_update_dynamics(
         self, operand: Tuple[DadsTrainingState, QDTransition]
     ) -> Tuple[Params, float, optax.OptState]:
@@ -464,7 +457,6 @@ def _not_update_dynamics(
             training_state.dynamics_optimizer_state,
         )
 
-    @partial(jax.jit, static_argnames=("self",))
     def _update_networks(
         self,
         training_state: DadsTrainingState,
@@ -566,7 +558,6 @@ def _update_networks(
 
         return new_training_state, metrics
 
-    @partial(jax.jit, static_argnames=("self",))
     def update(
         self,
         training_state: DadsTrainingState,

qdax/baselines/dads_smerl.py

@@ -5,7 +5,6 @@
 """
 
 from dataclasses import dataclass
-from functools import partial
 from typing import Optional, Tuple
 
 import jax
@@ -40,8 +39,7 @@ def __init__(self, config: DadsSmerlConfig, action_size: int, descriptor_size: i
         super(DADSSMERL, self).__init__(config, action_size, descriptor_size)
         self._config: DadsSmerlConfig = config
 
-    @partial(jax.jit, static_argnames=("self",))
-    def _compute_reward(
+    def _compute_reward(  # type: ignore
         self,
         transition: QDTransition,
         training_state: DadsTrainingState,
@@ -74,7 +72,6 @@ def _compute_reward(
 
         return rewards
 
-    @partial(jax.jit, static_argnames=("self",))
     def update(
         self,
         training_state: DadsTrainingState,

qdax/baselines/diayn.py

@@ -4,7 +4,6 @@
 """
 
 from dataclasses import dataclass
-from functools import partial
 from typing import Callable, Tuple
 
 import jax
@@ -177,7 +176,6 @@ def init(  # type: ignore
             steps=jnp.array(0),
         )
 
-    @partial(jax.jit, static_argnames=("self", "add_log_p_z"))
     def _compute_diversity_reward(
         self,
         transition: QDTransition,
@@ -212,8 +210,7 @@ def _compute_diversity_reward(
             reward += jnp.log(self._config.num_skills)
         return reward
 
-    @partial(jax.jit, static_argnames=("self", "env", "deterministic"))
-    def play_step_fn(
+    def play_step_fn(  # type: ignore
         self,
         env_state: EnvState,
         training_state: DiaynTrainingState,
@@ -279,14 +276,13 @@ def play_step_fn(
 
         return next_env_state, training_state, transition
 
-    @partial(jax.jit, static_argnames=("self", "play_step_fn", "env_batch_size"))
-    def eval_policy_fn(
+    def eval_policy_fn(  # type: ignore
         self,
         training_state: DiaynTrainingState,
         eval_env_first_state: EnvState,
         play_step_fn: Callable[
-            [EnvState, Params, RNGKey],
-            Tuple[EnvState, Params, RNGKey, QDTransition],
+            [EnvState, Params],
+            Tuple[EnvState, Params, QDTransition],
         ],
         env_batch_size: int,
     ) -> Tuple[Reward, Reward, Reward, StateDescriptor]:
@@ -347,7 +343,6 @@ def eval_policy_fn(
             transitions.state_desc,
         )
 
-    @partial(jax.jit, static_argnames=("self",))
     def _compute_reward(
         self, transition: QDTransition, training_state: DiaynTrainingState
     ) -> Reward:
@@ -366,7 +361,6 @@ def _compute_reward(
             add_log_p_z=True,
         )
 
-    @partial(jax.jit, static_argnames=("self",))
     def _update_networks(
         self,
         training_state: DiaynTrainingState,
@@ -469,7 +463,6 @@ def _update_networks(
 
         return new_training_state, metrics
 
-    @partial(jax.jit, static_argnames=("self",))
     def update(
         self,
         training_state: DiaynTrainingState,

qdax/baselines/diayn_smerl.py

@@ -5,7 +5,6 @@
 """
 
 from dataclasses import dataclass
-from functools import partial
 from typing import Optional, Tuple
 
 import jax
@@ -39,8 +38,7 @@ def __init__(self, config: DiaynSmerlConfig, action_size: int):
         super(DIAYNSMERL, self).__init__(config, action_size)
         self._config: DiaynSmerlConfig = config
 
-    @partial(jax.jit, static_argnames=("self",))
-    def _compute_reward(
+    def _compute_reward(  # type: ignore
         self,
         transition: QDTransition,
         training_state: DiaynTrainingState,
@@ -81,7 +79,6 @@ def _compute_reward(
 
         return rewards
 
-    @partial(jax.jit, static_argnames=("self",))
     def update(
         self,
         training_state: DiaynTrainingState,