First commit

.gitignore

@@ -0,0 +1,8 @@
+datasets
+videos
+figures
+*.log
+*.png
+mujoco210
+.DS_Store
+__pycache__/

app.py

@@ -0,0 +1,168 @@
+import glob
+import os
+
+import gymnasium as gym
+import numpy as np
+from gymnasium.wrappers import RecordVideo
+from moviepy.video.compositing.concatenate import concatenate_videoclips
+from moviepy.video.io.VideoFileClip import VideoFileClip
+from sympy import latex
+
+from interpretable import InterpretablePolicyExtractor
+from utils import generate_dataset_from_expert, rollouts
+import matplotlib.pyplot as plt
+
+import torch
+
+import gradio as gr
+import sys
+
+intro = """
+# Making RL Policy Interpretable with Kolmogorov-Arnold Network 🧠 ➙ 🔢
+
+Waris Radji<sup>1</sup>, Corentin Léger<sup>2</sup>, Hector Kohler<sup>1</sup>  
+<small><sup>1</sup>[Inria, team Scool](https://team.inria.fr/scool/)    <sup>2</sup>[Inria, team Flowers](https://flowers.inria.fr/)</small>
+
+
+
+In this demo, we showcase a method to make a trained Reinforcement Learning (RL) policy interpretable using the Kolmogorov-Arnold Network (KAN). The process involves transferring the knowledge from a pre-trained RL policy to a KAN. We achieve this by training the KAN to map actions from observations obtained from trajectories of the pre-trained policy.
+
+## Procedure
+
+- Train the KAN using observations from trajectories generated by a pre-trained RL policy, the KAN learns to map observations to corresponding actions.
+- Apply symbolic regression algorithms to the KAN's learned mapping.
+- Extract an interpretable policy expressed in symbolic form.
+
+For more information about KAN you can read the [paper](https://arxiv.org/abs/2404.19756), and check the [PyTorch official information](https://github.com/KindXiaoming/pykan).
+To follow the progress of KAN in RL you can check the repo [kanrl](https://github.com/riiswa/kanrl).
+"""
+
+envs = ["CartPole-v1", "MountainCar-v0", "Acrobot-v1", "Pendulum-v1", "MountainCarContinuous-v0", "LunarLander-v2", "Swimmer-v4", "Hopper-v4"]
+
+
+class Logger:
+    def __init__(self, filename):
+        self.terminal = sys.stdout
+        self.log = open(filename, "w")
+
+    def write(self, message):
+        self.terminal.write(message)
+        self.log.write(message)
+
+    def flush(self):
+        self.terminal.flush()
+        self.log.flush()
+
+    def isatty(self):
+        return False
+
+
+sys.stdout = Logger("output.log")
+sys.stderr = Logger("output.log")
+
+
+def read_logs():
+    sys.stdout.flush()
+    with open("output.log", "r") as f:
+        return f.read()
+
+
+if __name__ == "__main__":
+    torch.set_default_dtype(torch.float32)
+    dataset_path = None
+    ipe = None
+    env_name = None
+
+    def load_video_and_dataset(_env_name):
+        global dataset_path
+        global env_name
+        env_name = _env_name
+
+        dataset_path, video_path = generate_dataset_from_expert("ppo", _env_name, 15, 3)
+        return video_path, gr.Button("Compute the symbolic policy!", interactive=True)
+
+
+    def parse_integer_list(input_str):
+        if not input_str or input_str.isspace():
+            return None
+
+        elements = input_str.split(',')
+
+        try:
+            int_list = tuple([int(elem.strip()) for elem in elements])
+            return int_list
+        except ValueError:
+            return False
+
+    def extract_interpretable_policy(env_name, kan_widths):
+        global ipe
+
+        widths = parse_integer_list(kan_widths)
+        if kan_widths is False:
+            gr.Warning(f"Please enter widths {kan_widths} in the right format... The current run is executed with no hidden layer.")
+            widths = None
+
+        ipe = InterpretablePolicyExtractor(env_name, widths)
+        ipe.train_from_dataset(dataset_path, steps=50)
+
+        ipe.policy.prune()
+        ipe.policy.plot(mask=True, scale=5)
+
+        fig = plt.gcf()
+        fig.canvas.draw()
+        return np.array(fig.canvas.renderer.buffer_rgba())
+
+    def symbolic_policy():
+        global ipe
+        global env_name
+        lib = ['x', 'x^2', 'x^3', 'x^4', 'exp', 'log', 'sqrt', 'tanh', 'sin', 'abs']
+        ipe.policy.auto_symbolic(lib=lib)
+        env = gym.make(env_name, render_mode="rgb_array")
+        env = RecordVideo(env, video_folder="videos", episode_trigger=lambda x: True, name_prefix=f"kan-{env_name}")
+
+        rollouts(env, ipe.forward, 2)
+
+        video_path = os.path.join("videos", f"kan-{env_name}.mp4")
+        video_files = glob.glob(os.path.join("videos", f"kan-{env_name}-episode*.mp4"))
+        clips = [VideoFileClip(file) for file in video_files]
+        final_clip = concatenate_videoclips(clips)
+        final_clip.write_videofile(video_path, codec="libx264", fps=24)
+
+        symbolic_formula = f"### The symbolic formula of the policy is:"
+        formulas = ipe.policy.symbolic_formula()[0]
+        for i, formula in enumerate(formulas):
+            symbolic_formula += "\n$$ a_" + str(i) + "=" + latex(formula) + "$$"
+        if ipe._action_is_discrete:
+            symbolic_formula += "\n" + r"$$ a = \underset{i}{\mathrm{argmax}} \ a_i.$$"
+
+        return video_path, symbolic_formula
+
+
+    css = """
+#formula {overflow-x: auto!important};
+    """
+
+    with gr.Blocks(theme='gradio/monochrome', css=css) as app:
+        gr.Markdown(intro)
+
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### Pretrained policy loading (PPO from [rl-baselines3-zoo](https://github.com/DLR-RM/rl-baselines3-zoo))")
+                choice = gr.Dropdown(envs, label="Environment name")
+                expert_video = gr.Video(label="Expert policy video", interactive=False, autoplay=True)
+                kan_widths = gr.Textbox(value="2", label="Widths of the hidden layers of the KAN, separated by commas (e.g. `3,3`). Leave empty if there are no hidden layers.")
+                button = gr.Button("Compute the symbolic policy!", interactive=False)
+            with gr.Column():
+                gr.Markdown("### Symbolic policy extraction")
+                kan_architecture = gr.Image(interactive=False, label="KAN architecture")
+                sym_video = gr.Video(label="Symbolic policy video", interactive=False, autoplay=True)
+        sym_formula = gr.Markdown(elem_id="formula")
+        with gr.Accordion("See logs"):
+            logs = gr.Textbox(label="Logs", interactive=False)
+        choice.input(load_video_and_dataset, inputs=[choice], outputs=[expert_video, button])
+        button.click(extract_interpretable_policy, inputs=[choice, kan_widths], outputs=[kan_architecture]).then(
+            symbolic_policy, inputs=[], outputs=[sym_video, sym_formula]
+        )
+        app.load(read_logs, None, logs, every=1)
+
+    app.launch()

example.py

@@ -0,0 +1,26 @@
+import gym
+from gym.wrappers import RecordVideo
+from matplotlib import pyplot as plt
+
+from interpretable.interpretable import InterpretablePolicyExtractor
+from interpretable.utils import generate_dataset_from_expert, rollouts
+
+if __name__ == "__main__":
+    env_name = "CartPole-v1"
+    dataset_path = generate_dataset_from_expert("ppo", env_name, force=True)
+    ipe = InterpretablePolicyExtractor(env_name)
+    results = ipe.train_from_dataset(dataset_path)
+    ipe.policy.prune()
+    ipe.policy.plot(mask=True)
+    plt.savefig("kan-policy.png")
+
+    env = gym.make(env_name, render_mode="rgb_array")
+    env = RecordVideo(env, video_folder="videos", episode_trigger=lambda x: True, name_prefix=f"kan-{env_name}")
+
+    ipe.policy.auto_symbolic()
+
+    ipe.policy.plot(mask=True)
+    plt.savefig("sym-policy.png")
+    print(ipe.policy.symbolic_formula())
+
+    rollouts(env, ipe.forward, 2)

interpretable.py

@@ -0,0 +1,47 @@
+import torch
+from typing import Dict, Tuple, Optional, Callable, Union
+import gymnasium as gym
+from kan import KAN
+import numpy as np
+
+
+def extract_dim(space: gym.Space):
+    if isinstance(space, gym.spaces.Box) and len(space.shape) == 1:
+        return space.shape[0], False
+    elif isinstance(space, gym.spaces.Discrete):
+        return space.n, True
+    else:
+        raise NotImplementedError(f"There is no support for space {space}.")
+
+
+class InterpretablePolicyExtractor:
+    lib = ['x', 'x^2', 'x^3', 'x^4', 'exp', 'log', 'sqrt', 'tanh', 'sin', 'abs']
+
+    def __init__(self, env_name: str, hidden_widths: Optional[Tuple[int]]=None):
+        self.env = gym.make(env_name)
+        if hidden_widths is None:
+            hidden_widths = []
+        observation_dim, self._observation_is_discrete = extract_dim(self.env.observation_space)
+        action_dim, self._action_is_discrete = extract_dim(self.env.action_space)
+        self.policy = KAN(width=[observation_dim, *hidden_widths, action_dim])
+        self.loss_fn = torch.nn.MSELoss() if not self._action_is_discrete else torch.nn.CrossEntropyLoss()
+
+    def train_from_dataset(self, dataset: Union[Dict[str, torch.Tensor], str], steps: int = 20):
+        if isinstance(dataset, str):
+            dataset = torch.load(dataset)
+        if dataset["train_label"].ndim == 1 and not self._action_is_discrete:
+            dataset["train_label"] = dataset["train_label"][:, None]
+        if dataset["train_label"].ndim == 1 and not self._action_is_discrete:
+            dataset["test_label"] = dataset["test_label"][:, None]
+        return self.policy.train(dataset, opt="LBFGS", steps=steps, loss_fn=self.loss_fn)
+
+    def forward(self, observation):
+        observation = torch.from_numpy(observation)
+        action = self.policy(observation.unsqueeze(0))
+        if self._action_is_discrete:
+            return action.argmax(axis=-1).squeeze().item()
+        else:
+            return action.squeeze(0).detach().numpy()
+
+    def train_from_policy(self, policy: Callable[[np.ndarray], Union[np.ndarray, int, float]], steps: int):
+        raise NotImplementedError()  # TODO

requirements.txt

@@ -0,0 +1,17 @@
+gymnasium[box2d,mujoco]~=0.29.1
+torch~=2.2.2
+numpy~=1.26.4
+pykan~=0.0.2
+tqdm~=4.66.2
+scikit-learn~=1.4.2
+matplotlib~=3.8.4
+moviepy~=1.0.3
+huggingface_hub
+gradio
+huggingface_sb3
+stable_baselines3
+rl_zoo3
+gym
+shimmy>=0.2.1
+mujoco-py
+cpython<3

utils.py

@@ -0,0 +1,178 @@
+import glob
+import os
+import pickle
+
+import torch
+import numpy as np
+import gymnasium as gym
+from huggingface_hub.utils import EntryNotFoundError
+from huggingface_sb3 import load_from_hub
+from moviepy.video.compositing.concatenate import concatenate_videoclips
+from moviepy.video.io.VideoFileClip import VideoFileClip
+from rl_zoo3 import ALGOS
+from gymnasium.wrappers import RecordVideo
+from stable_baselines3.common.running_mean_std import RunningMeanStd
+
+import os
+import tarfile
+import urllib.request
+
+
+def install_mujoco():
+    mujoco_url = "https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz"
+    mujoco_file = "mujoco210-linux-x86_64.tar.gz"
+    mujoco_dir = "mujoco210"
+
+    # Check if the directory already exists
+    if not os.path.exists("mujoco210"):
+        # Download Mujoco if not exists
+        print("Downloading Mujoco...")
+        urllib.request.urlretrieve(mujoco_url, mujoco_file)
+
+        # Extract Mujoco
+        print("Extracting Mujoco...")
+        with tarfile.open(mujoco_file, "r:gz") as tar:
+            tar.extractall()
+
+        # Clean up the downloaded tar file
+        os.remove(mujoco_file)
+
+        print("Mujoco installed successfully!")
+    else:
+        print("Mujoco already installed.")
+
+    # Set environment variable MUJOCO_PY_MUJOCO_PATH
+    os.environ["MUJOCO_PY_MUJOCO_PATH"] = os.path.abspath(mujoco_dir)
+
+    ld_library_path = os.environ.get("LD_LIBRARY_PATH", "")
+    mujoco_bin_path = os.path.join(os.path.abspath(mujoco_dir), "bin")
+    if mujoco_bin_path not in ld_library_path:
+        os.environ["LD_LIBRARY_PATH"] = ld_library_path + ":" + mujoco_bin_path
+
+
+
+class NormalizeObservation(gym.Wrapper):
+    def __init__(self, env: gym.Env, clip_obs: float, obs_rms: RunningMeanStd, epsilon: float):
+        gym.Wrapper.__init__(self, env)
+        self.clip_obs = clip_obs
+        self.obs_rms = obs_rms
+        self.epsilon = epsilon
+
+    def step(self, action):
+        observation, reward, terminated, truncated, info = self.env.step(action)
+        observation = self.normalize(np.array([observation]))[0]
+        return observation, reward, terminated, truncated, info
+
+    def reset(self, **kwargs):
+        observation, info = self.env.reset(**kwargs)
+        return self.normalize(np.array([observation]))[0], info
+
+    def normalize(self, obs):
+        return np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon), -self.clip_obs, self.clip_obs)
+
+
+class CreateDataset(gym.Wrapper):
+    def __init__(self, env: gym.Env):
+        gym.Wrapper.__init__(self, env)
+        self.observations = []
+        self.actions = []
+        self.last_observation = None
+
+    def step(self, action):
+        self.observations.append(self.last_observation)
+        self.actions.append(action)
+        observation, reward, terminated, truncated, info = self.env.step(action)
+        self.last_observation = observation
+        return observation, reward, terminated, truncated, info
+
+    def reset(self, **kwargs):
+        observation, info = self.env.reset(**kwargs)
+        self.last_observation = observation
+        return observation, info
+
+    def get_dataset(self):
+        if isinstance(self.env.action_space, gym.spaces.Box) and self.env.action_space.shape != (1,):
+            actions = np.vstack(self.actions)
+        else:
+            actions = np.hstack(self.actions)
+        return np.vstack(self.observations), actions
+
+
+def rollouts(env, policy, num_episodes=1):
+    for episode in range(num_episodes):
+        done = False
+        observation, _ = env.reset()
+        while not done:
+            action = policy(observation)
+            observation, reward, terminated, truncated, _ = env.step(action)
+            done = terminated or truncated
+    env.close()
+
+
+def generate_dataset_from_expert(algo, env_name, num_train_episodes=5, num_test_episodes=2, force=False):
+    if env_name.startswith("Swimmer") or env_name.startswith("Hopper"):
+        install_mujoco()
+    if env_name == "Swimmer-v4":
+        env_name = "Swimmer-v3"
+    elif env_name == "Hopper-v4":
+        env_name = "Hopper-v3"
+    dataset_path = os.path.join("datasets", f"{algo}-{env_name}.pt")
+    video_path = os.path.join("videos", f"{algo}-{env_name}.mp4")
+    if os.path.exists(dataset_path) and os.path.exists(video_path) and not force:
+        return dataset_path, video_path
+    repo_id = f"sb3/{algo}-{env_name}"
+    policy_file = f"{algo}-{env_name}.zip"
+
+    expert_path = load_from_hub(repo_id, policy_file)
+    try:
+        vec_normalize_path = load_from_hub(repo_id, "vec_normalize.pkl")
+        with open(vec_normalize_path, "rb") as f:
+            vec_normalize = pickle.load(f)
+            if vec_normalize.norm_obs:
+                vec_normalize_params = {"clip_obs": vec_normalize.clip_obs, "obs_rms": vec_normalize.obs_rms, "epsilon": vec_normalize.epsilon}
+            else:
+                vec_normalize_params = None
+    except EntryNotFoundError:
+        vec_normalize_params = None
+
+    expert = ALGOS[algo].load(expert_path)
+    train_env = gym.make(env_name)
+    train_env = CreateDataset(train_env)
+    if vec_normalize_params is not None:
+        train_env = NormalizeObservation(train_env, **vec_normalize_params)
+    test_env = gym.make(env_name, render_mode="rgb_array")
+    test_env = CreateDataset(test_env)
+    if vec_normalize_params is not None:
+        test_env = NormalizeObservation(test_env, **vec_normalize_params)
+    test_env = RecordVideo(test_env, video_folder="videos", episode_trigger=lambda x: True, name_prefix=f"{algo}-{env_name}")
+
+    def policy(obs):
+        return expert.predict(obs, deterministic=True)[0]
+
+    os.makedirs("videos", exist_ok=True)
+    rollouts(train_env, policy, num_train_episodes)
+    rollouts(test_env, policy, num_test_episodes)
+
+    train_observations, train_actions = train_env.get_dataset()
+    test_observations, test_actions = test_env.get_dataset()
+
+    dataset = {
+        "train_input": torch.from_numpy(train_observations),
+        "test_input": torch.from_numpy(test_observations),
+        "train_label": torch.from_numpy(train_actions),
+        "test_label": torch.from_numpy(test_actions)
+    }
+
+    os.makedirs("datasets", exist_ok=True)
+    torch.save(dataset, dataset_path)
+
+    video_files = glob.glob(os.path.join("videos", f"{algo}-{env_name}-episode*.mp4"))
+    clips = [VideoFileClip(file) for file in video_files]
+    final_clip = concatenate_videoclips(clips)
+    final_clip.write_videofile(video_path, codec="libx264", fps=24)
+
+    return dataset_path, video_path
+
+
+if __name__ == "__main__":
+    generate_dataset_from_expert("ppo", "CartPole-v1", force=True)