| 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() | |