main files

.gitattributes

@@ -23,7 +23,6 @@
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
 *.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
@@ -33,3 +32,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+models/**/* filter=lfs diff=lfs merge=lfs -text
+evaluations/**/* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -26,23 +26,23 @@ The sparse model parameters were obtained with [EauDeQN](https://arxiv.org/pdf/2
 
 5 seeds are available for each configuration which makes a total of **750 available models** 📈.
 
-The [evaluate.ipynb](./evaluate.ipynb) notebook contains a minimal example to evaluate to model parameters 🧑‍🏫 It uses JAX 🚀 The hyperparameters used during training are reported in [config.json](./config.json) 🔧
+The [evaluate.ipynb](./evaluate.ipynb) notebook contains a minimal example to evaluate to model parameters 🧑‍🏫 It uses JAX 🚀 The hyperparameters used during training are reported in [online_config.json](./online_config.json) and [offline_config.json](./offline_config.json) 🔧
 
 The training code is available soon ⏳
 
-### Model performances 
-| <div style="width:300px; font-size: 30px; font-family:Serif; font-name:Times New Roman" > **EauDeDQN** and **EauDeCQL** achieve high sparsity while keeping performances high. <br> Published at [RLDM](https://arxiv.org/pdf/2503.01437)✨ </br> <div style="font-size: 16px"> <details> <summary id=games>List of Atari games</summary> *BeamRider, MsPacman, Qbert, Pong, Enduro, SpaceInvaders, Assault, CrazyClimber, Boxing, VideoPinball.* </details> </div> </div> | <img src="performances.png" alt="drawing" width="600px"/> |
+### Model sparsity & performances
+| <div style="width:300px; font-size: 30px; font-family:Serif; font-name:Times New Roman" > **EauDeDQN** and **EauDeCQL** achieve high sparsity while keeping performances high. <br> Published at [RLDM](https://arxiv.org/pdf/2503.01437)✨ </br> <div style="font-size: 16px"> <details> <summary id=games>List of Atari games</summary> *BeamRider, MsPacman, Qbert, Pong, Enduro, SpaceInvaders, Assault, CrazyClimber, Boxing, VideoPinball.* </details> </div> </div> | <img src="sparsities.png" alt="drawing" width="600px"/> |
 | :-: | :-: |
 
+The episodic returns and lenghts are available in the [evaluations](./evaluations/) folder 🔬
+
 ## User installation
 Python 3.10 is recommended. Create a Python virtual environment, activate it, update pip and install the package and its dependencies in editable mode:
 ```bash
 python3.10 -m venv env
 source env/bin/activate
-pip install --upgrade pip
-pip install numpy==1.23.5  # to avoid numpy==2.XX
+pip install --upgrade pip setuptools wheel
 pip install -r requirements.txt
-pip install --upgrade "jax[cuda12_pip]==0.4.13" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
 ## Citing `Eau De Q-Network`

atari.py

@@ -0,0 +1,127 @@
+"""
+The environment is inspired from https://github.com/google/dopamine/blob/master/dopamine/discrete_domains/atari_lib.py
+"""
+
+import ale_py
+from typing import Tuple, Dict
+from gymnasium.wrappers import RecordVideo
+import gymnasium as gym
+import numpy as np
+import jax
+import jax.numpy as jnp
+import cv2
+
+
+class AtariEnv:
+    def __init__(self, name: str) -> None:
+        self.name = name
+        self.state_height, self.state_width = (84, 84)
+        self.n_stacked_frames = 4
+        self.n_skipped_frames = 4
+
+        gym.register_envs(ale_py)  # To use ale with gym which speeds up step()
+        self.env = gym.make(
+            f"ALE/{self.name}-v5",
+            full_action_space=False,
+            frameskip=1,
+            repeat_action_probability=0.25,
+            max_num_frames_per_episode=100_000,
+            continuous=False,
+            continuous_action_threshold=0.0,
+            render_mode="rgb_array",
+        ).env
+
+        self.n_actions = self.env.action_space.n
+        self.original_state_height, self.original_state_width, _ = self.env.observation_space._shape
+        self.screen_buffer = [
+            np.empty((self.original_state_height, self.original_state_width), dtype=np.uint8),
+            np.empty((self.original_state_height, self.original_state_width), dtype=np.uint8),
+        ]
+
+    @property
+    def observation(self) -> np.ndarray:
+        return np.copy(self.state_[:, :, -1])
+
+    @property
+    def state(self) -> np.ndarray:
+        return jnp.array(self.state_, dtype=jnp.float32)
+
+    def reset(self) -> None:
+        self.env.reset()
+
+        self.n_steps = 0
+
+        self.env.env.ale.getScreenGrayscale(self.screen_buffer[0])
+        self.screen_buffer[1].fill(0)
+
+        self.state_ = np.zeros((self.state_height, self.state_width, self.n_stacked_frames), dtype=np.uint8)
+        self.state_[:, :, -1] = self.resize()
+
+    def step(self, action: jnp.int8) -> Tuple[float, bool]:
+        reward = 0
+
+        for idx_frame in range(self.n_skipped_frames):
+            _, reward_, terminal, _, _ = self.env.step(action)
+
+            reward += reward_
+
+            if idx_frame >= self.n_skipped_frames - 2:
+                t = idx_frame - (self.n_skipped_frames - 2)
+                self.env.env.ale.getScreenGrayscale(self.screen_buffer[t])
+
+            if terminal:
+                break
+
+        self.state_ = np.roll(self.state_, -1, axis=-1)
+        self.state_[:, :, -1] = self.pool_and_resize()
+
+        self.n_steps += 1
+
+        return reward, terminal
+
+    def pool_and_resize(self) -> np.ndarray:
+        np.maximum(self.screen_buffer[0], self.screen_buffer[1], out=self.screen_buffer[0])
+
+        return self.resize()
+
+    def resize(self):
+        return np.asarray(
+            cv2.resize(self.screen_buffer[0], (self.state_width, self.state_height), interpolation=cv2.INTER_AREA),
+            dtype=np.uint8,
+        )
+
+    def evaluate_one_simulation(
+        self,
+        q,
+        q_params: Dict,
+        horizon: int,
+        eps_eval: float,
+        exploration_key: jax.random.PRNGKey,
+        video_path: str,
+    ) -> float:
+        ale = self.env.env.ale
+        self.env = RecordVideo(
+            self.env,
+            video_folder=video_path if video_path is not None else ".",
+            name_prefix="",
+            episode_trigger=lambda x: video_path is not None,
+        )
+        self.env.env.ale = ale
+
+        sun_reward = 0
+        terminal = False
+        self.reset()
+
+        while not terminal and self.n_steps < horizon:
+            exploration_key, key = jax.random.split(exploration_key)
+            if jax.random.uniform(key) < eps_eval:
+                action = jax.random.choice(key, jnp.arange(self.n_actions)).astype(jnp.int8)
+            else:
+                action = q.best_action(q_params, self.state)
+
+            reward, terminal = self.step(action)
+
+            sun_reward += reward
+
+        self.env.close()
+        return sun_reward, terminal

evaluate.ipynb

@@ -0,0 +1,77 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "import pickle\n",
+    "from atari import AtariEnv\n",
+    "from networks import QNetwork\n",
+    "\n",
+    "# ------- START TO MODIFY ------- #\n",
+    "ALGO = \"eaudedqn\" # choose between eaudedqn, polyprunedqn, dqn, eaudecql, polyprunecql, and cql.\n",
+    "GAME = \"SpaceInvaders\" # choose between BeamRider, MsPacman, Qbert, Pong, Enduro, SpaceInvaders, Assault, CrazyClimber, Boxing, and VideoPinball.\n",
+    "FEATURE_SIZE = 32 # choose between 32, 512, and 2048.\n",
+    "NETWORK_SEED = 1 # choose between 1, 2, 3, 4, and 5.\n",
+    "EVALUATION_SEED = 0\n",
+    "HORIZON = 27000\n",
+    "EPSILON = 0.01\n",
+    "RECORD_VIDEO = False\n",
+    "# ------- END TO MODIFY ------- #\n",
+    "\n",
+    "params_path = f\"models/{GAME}/{ALGO}/feature_size_{FEATURE_SIZE}_seed_{NETWORK_SEED}\"\n",
+    "\n",
+    "env = AtariEnv(GAME)\n",
+    "\n",
+    "q = QNetwork([32, 64, 64, FEATURE_SIZE], env.n_actions)\n",
+    "\n",
+    "with open(params_path, \"rb\") as handle:\n",
+    "    q_params = pickle.load(handle)\n",
+    "\n",
+    "return_, absorbing = env.evaluate_one_simulation(\n",
+    "    q, q_params, HORIZON, EPSILON, jax.random.PRNGKey(EVALUATION_SEED), params_path + \"_eval\" if RECORD_VIDEO else None\n",
+    ")\n",
+    "print(\"Undiscounted return:\", return_)\n",
+    "print(\"N steps\", env.n_steps, \"; Horizon\", HORIZON, \"; Absorbing\", absorbing)\n",
+    "non_zeros = sum(jax.tree.leaves(jax.tree.map(jnp.count_nonzero, q_params)))\n",
+    "n_params = sum(jax.tree.leaves(jax.tree.map(jnp.size, q_params)))\n",
+    "print(\"Spartity level:\", (1 - jnp.float32(non_zeros) / jnp.float32(n_params)))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "env",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

networks.py

@@ -0,0 +1,38 @@
+from typing import Sequence
+from functools import partial
+import jax
+import jax.numpy as jnp
+import flax.linen as nn
+
+
+class DQNNet(nn.Module):
+    features: Sequence[int]
+    n_actions: int
+
+    @nn.compact
+    def __call__(self, x):
+        initializer = nn.initializers.xavier_uniform()
+        x = nn.relu(
+            nn.Conv(features=self.features[0], kernel_size=(8, 8), strides=(4, 4), kernel_init=initializer)(
+                jnp.array(x, ndmin=4) / 255.0
+            )
+        )
+        x = nn.relu(nn.Conv(features=self.features[1], kernel_size=(4, 4), strides=(2, 2), kernel_init=initializer)(x))
+        x = nn.relu(nn.Conv(features=self.features[2], kernel_size=(3, 3), strides=(1, 1), kernel_init=initializer)(x))
+        x = x.reshape((x.shape[0], -1))
+
+        x = jnp.squeeze(x)
+
+        for idx_layer in range(3, len(self.features)):
+            x = nn.relu((nn.Dense(self.features[idx_layer], kernel_init=initializer)(x)))
+
+        return nn.Dense(self.n_actions, kernel_init=initializer)(x)
+
+
+class QNetwork:
+    def __init__(self, features: Sequence[int], n_actions: int) -> None:
+        self.network = DQNNet(features, n_actions)
+
+    @partial(jax.jit, static_argnames="self")
+    def best_action(self, params, state: jnp.ndarray) -> jnp.int8:
+        return jnp.argmax(self.network.apply(params, state)).astype(jnp.int8)

offline_config.json

@@ -0,0 +1,37 @@
+{
+    "shared_parameters": {
+        "features": [
+            32,
+            64,
+            64,
+            "Feature Size"
+        ],
+        "replay_buffer_capacity": 50000,
+        "batch_size": 32,
+        "update_horizon": 1,
+        "gamma": 0.99,
+        "learning_rate": 5e-05,
+        "architecture_type": "cnn",
+        "target_update_frequency": 2000,
+        "n_buffers_to_load": 5,
+        "n_epochs": 50,
+        "n_fitting_steps": 62500
+    },
+    "eaudecql": {
+        "n_networks": 5,
+        "max_noise": 3.0,
+        "max_speed": 0.01,
+        "reset_optimizer": true,
+        "alpha_cql": 0.1
+    },
+    "polyprunecql": {
+        "sparcity_start_step": 625000,
+        "sparcity_end_step": 2500000,
+        "sparcity_update_freq": 1000,
+        "final_sparsity": 0.95,
+        "alpha_cql": 0.1
+    },
+    "cql": {
+        "alpha_cql": 0.1
+    }
+}

online_config.json

@@ -0,0 +1,37 @@
+{
+    "shared_parameters": {
+        "replay_buffer_capacity": 1000000,
+        "batch_size": 32,
+        "update_horizon": 1,
+        "gamma": 0.99,
+        "learning_rate": 6.25e-05,
+        "horizon": 27000,
+        "n_epochs": 40,
+        "n_training_steps_per_epoch": 250000,
+        "n_initial_samples": 20000,
+        "epsilon_end": 0.01,
+        "epsilon_duration": 250000.0,
+        "target_update_frequency": 8000,
+        "update_to_data": 4.0,
+        "features": [
+            32,
+            64,
+            64,
+            "Feature Size"
+        ],
+        "architecture_type": "cnn"
+    },
+    "eaudedqn": {
+        "n_networks": 5,
+        "max_noise": 3.0,
+        "max_speed": 0.01,
+        "reset_optimizer": true
+    },
+    "polyprunedqn": {
+        "sparcity_start_step": 2000000,
+        "sparcity_end_step": 8000000,
+        "sparcity_update_freq": 4000,
+        "final_sparsity": 0.95
+    },
+    "dqn": {}
+}