| import gradio as gr | |
| from train import TrainingLoop | |
| from scipy.special import softmax | |
| import numpy as np | |
| train = None | |
| frames, attributions = None, None | |
| lunar_lander_spec_conversion = { | |
| 0: "X-coordinate", | |
| 1: "Y-coordinate", | |
| 2: "Linear velocity in the X-axis", | |
| 3: "Linear velocity in the Y-axis", | |
| 4: "Angle", | |
| 5: "Angular velocity", | |
| 6: "Left leg touched the floor", | |
| 7: "Right leg touched the floor" | |
| } | |
| def create_training_loop(env_spec): | |
| global train | |
| train = TrainingLoop(env_spec=env_spec) | |
| train.create_agent() | |
| return train.env.spec | |
| def display_softmax(inputs): | |
| inputs = np.array(inputs) | |
| probabilities = softmax(inputs) | |
| softmax_dict = {name: float(prob) for name, prob in zip(lunar_lander_spec_conversion.values(), probabilities)} | |
| return softmax_dict | |
| def generate_output(num_iterations, option): | |
| global frames, attributions | |
| frames, attributions = train.explain_trained(num_iterations=num_iterations, option=option) | |
| slider.maximum = len(frames) | |
| def get_frame_and_attribution(slider_value): | |
| global frames, attributions | |
| slider_value = min(slider_value, len(frames) - 1) | |
| frame = frames[slider_value] | |
| print(f"{frame.shape=}") | |
| attribution = display_softmax(attributions[slider_value]) | |
| return frame, attribution | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Introspection in Deep Reinforcement Learning") | |
| gr.Markdown(r""" | |
| \#\# How this space works: | |
| This space was created for trying to apply [Integrated Gradients](https://captum.ai/docs/extension/integrated_gradients\#:~:text=Integrated%20gradients%20is%20a%20simple,and%20feature%20or%20rule%20extraction.) \ | |
| into Deep Reinforcement Learning Scenarions. It uses PyTorch's captum library for interpretability, and Gymnasium for the emulator of the continuous lunar lander. | |
| \#\#\# Training algorithm: [DDPG](https://arxiv.org/abs/1509.02971) | |
| This agent was trained with Deep Deterministic Policy Gradients, and outputs an average reward of 260.8 per episode (successful) | |
| \#\#\# Using this space: | |
| - First, select the environment (futurely there will be more environments available) | |
| - Then, select if you want the baseline (see IG paper for more detail) to be \ | |
| a torch `tensor` of zeroes, or a running average of the initial frames of a few episodes (selected on the right) \ | |
| - Click attribute and wait a few seconds (usually 20-25s) for the attributions to be computed with the trained agent over 10 episodes | |
| - Finally, use the slider to get a key frame that tells the attributions of the agent. They're under a Softmax to fit the component's requirements for a probability distribution. | |
| """) | |
| with gr.Tab(label="Attribute"): | |
| env_spec = gr.Dropdown(choices=["LunarLander-v2"],type="value",multiselect=False, label="Environment Specification (e.g.: LunarLander-v2)") | |
| env = gr.Interface(title="Create the Environment", allow_flagging="never", inputs=env_spec, fn=create_training_loop, outputs=gr.JSON()) | |
| with gr.Row(): | |
| option = gr.Dropdown(choices=["Torch Tensor of 0's", "Running Average"], type="index") | |
| baselines = gr.Slider(label="Number of Baseline Iterations", interactive=True, minimum=0, maximum=100, value=10, step=5, info="Baseline inputs to collect for the average", render=True) | |
| gr.Button("ATTRIBUTE").click(fn=generate_output, inputs=[baselines, option]) | |
| slider = gr.Slider(label="Key Frame", minimum=0, maximum=1000, step=1, value=0) | |
| gr.Interface(fn=get_frame_and_attribution, inputs=slider, live=True, outputs=[gr.Image(label="Timestep"),gr.Label(label="Attributions")]) | |
| gr.Markdown(r"""\#\# Local Usage and Packages \ | |
| `pip install torch gymnasium 'gymnasium[box2d]'` \ | |
| You might need to install Box2D Separately, which requires a swig package to compile code from Python into C/C++, which is the language that Box2d was built in: \ | |
| `brew install swig` \ | |
| `pip install box2d \n \#\# Average Score: 164.38 (significant improvement from discrete action spaces) \ | |
| For each step, the reward: \ | |
| - is increased/decreased the closer/further the lander is to the landing pad. \ | |
| - is increased/decreased the slower/faster the lander is moving.\ | |
| - is decreased the more the lander is tilted (angle not horizontal). \ | |
| - is increased by 10 points for each leg that is in contact with the ground. \ | |
| - is decreased by 0.03 points each frame a side engine is firing.\ | |
| - is decreased by 0.3 points each frame the main engine is firing. \ | |
| The episode receives an additional reward of -100 or +100 points for crashing or landing safely respectively. An episode is considered a solution if it scores at least 200 points.\*\* \ | |
| \#\# `train()` and `load_trained()` \ | |
| `load_trained()` function loads a pre-trained model that ran through 1000 episodes of training, while `train()` does training from scratch. You can edit which one of the functions is running from the bottom of the main.py file. If you set render_mode=False, the program will train a lot faster.)\n demo.launch()""") | |
| demo.launch() |