working lunar lander attribution mechanism

app.py

@@ -1,27 +1,65 @@
 import gradio as gr
 from train import TrainingLoop
+from scipy.special import softmax
+import numpy as np
 
 
+train = None
 
-def image_classifier(inp):
-    return {'cat': 0.3, 'dog': 0.7}
+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)}
 
-iface = gr.Interface(fn=image_classifier, inputs="image", outputs="label")
+    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
+    frame = frames[slider_value]
+    attribution = display_softmax(attributions[slider_value])
 
+    return frame, attribution
 
-load_trained = gr.Interface(fn=image_classifier, inputs="image", outputs="label")
+with gr.Blocks() as demo:
+    gr.Markdown("# Introspection in Deep Reinforcement Learning")
 
+    with gr.Tab(label="Attribute"):
+        env_spec = gr.Textbox(label="Environment Specification (e.g.: LunarLander-v2)", lines=1)
+        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=20000, step=1, value=0)
 
+        gr.Interface(fn=get_frame_and_attribution, inputs=slider, live=True, outputs=[gr.Image(), gr.Label()])
 
-attribute = gr.Interface(fn=image_classifier, inputs="image", outputs="label")
 
-
-# iface = gr.TabbedInterface(interface_list=[train_from_scratch, load_trained, attribute], tab_names=["Train from Scratch", "Continue Training", "Attribute"],title="Attribution in Deep Reinforcement Learning")
-iface.launch()
+demo.launch()

train.py

@@ -18,7 +18,9 @@ class TrainingLoop:
             "render_mode": None
         }
 
-        self.env = None
+        self.env = gym.make(
+            **self.defaults
+        )
 
         self.defaults.update(**kwargs)
 
@@ -44,7 +46,7 @@ class TrainingLoop:
 
         score_history = []
 
-        for i in range(1000):
+        for i in range(10000):
             done = False
             score = 0
             obs, _ = self.env.reset()
@@ -177,8 +179,8 @@ class TrainingLoop:
         assert self.agent is not None
         
         baseline_options = {
-            "1": torch.zeros(8),
-            "2": self._collect_running_baseline_average(num_iterations), 
+            0: torch.zeros(8),
+            1: self._collect_running_baseline_average(num_iterations), 
         }
 
         baseline = baseline_options[option]