Update app with improved UI and markdown formatting

.gitignore

@@ -1 +1,53 @@
-.DS_Store
+.DS_Store
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual environments
+.venv/
+venv/
+ENV/
+env/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# Model files (if large)
+*.pth
+*.pt
+*.h5
+*.pkl
+*.joblib
+
+# Logs
+*.log
+logs/
+
+# Temporary files
+tmp/
+temp/

.python-version

@@ -0,0 +1 @@
+3.10

README.md

@@ -54,4 +54,4 @@ For each step, the reward:
 
 ## `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.
+`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.

app.py

@@ -3,14 +3,14 @@ from train import TrainingLoop
 from scipy.special import softmax
 import numpy as np
 
-
+# Global variables for training and data storage
 train = None
-
 frames, attributions = None, None
 
-lunar_lander_spec_conversion = {
+# Lunar Lander environment state feature mapping
+LUNAR_LANDER_FEATURES = {
     0: "X-coordinate",
-    1: "Y-coordinate",
+    1: "Y-coordinate", 
     2: "Linear velocity in the X-axis",
     3: "Linear velocity in the Y-axis",
     4: "Angle",
@@ -20,76 +20,206 @@ lunar_lander_spec_conversion = {
 }
 
 def create_training_loop(env_spec):
+    """Initialize the training loop with the specified environment."""
     global train
     train = TrainingLoop(env_spec=env_spec)
     train.create_agent()
-    
     return train.env.spec
 
 def display_softmax(inputs):
+    """Convert raw attribution values to softmax probabilities for visualization."""
     inputs = np.array(inputs)
     probabilities = softmax(inputs)
-
-    softmax_dict = {name: float(prob) for name, prob in zip(lunar_lander_spec_conversion.values(), probabilities)}
-
+    
+    softmax_dict = {
+        name: float(prob) 
+        for name, prob in zip(LUNAR_LANDER_FEATURES.values(), probabilities)
+    }
     return softmax_dict
 
 def generate_output(num_iterations, option):
+    """Generate attribution explanations for the trained agent."""
     global frames, attributions
-    frames, attributions = train.explain_trained(num_iterations=num_iterations, option=option)
+    frames, attributions = train.explain_trained(
+        num_iterations=num_iterations, 
+        option=option
+    )
     slider.maximum = len(frames)
 
 def get_frame_and_attribution(slider_value):
+    """Get frame and attribution data for the selected timestep."""
     global frames, attributions
     slider_value = min(slider_value, len(frames) - 1)
     frame = frames[slider_value]
-
-    print(f"{frame.shape=}")
-
+    
+    print(f"Frame shape: {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.Blocks(
+    title="Deep RL Explainability",
+    theme=gr.themes.Soft(),
+    css="""
+    .gradio-container {
+        max-width: 1200px !important;
+    }
+    .tab-nav {
+        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
+    }
+    """
+) as demo:
+    
+    # Header section
+    gr.Markdown("""
+    # 🚀 Deep Reinforcement Learning Explainability
+    
+    **Exploring AI decision-making through Integrated Gradients in RL environments**
+    
+    ---
     """)
-    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())
-
+    
+    # Introduction section
+    gr.Markdown("""
+    ## 📖 How This Works
+    
+    This application demonstrates the application of **[Integrated Gradients](https://captum.ai/docs/extension/integrated_gradients)** 
+    to Deep Reinforcement Learning scenarios. We use PyTorch's Captum library for interpretability 
+    and Gymnasium for the continuous Lunar Lander environment.
+    
+    ### 🧠 Training Algorithm: [DDPG](https://arxiv.org/abs/1509.02971)
+    
+    The agent is trained using **Deep Deterministic Policy Gradients** and achieves an average reward 
+    of **260.8** per episode (successful landings).
+    
+    ### 🎯 How to Use This Space
+    
+    1. **Select Environment**: Choose the Lunar Lander environment
+    2. **Choose Baseline**: Select between zero tensor or running average baseline
+    3. **Generate Attributions**: Click "ATTRIBUTE" and wait ~20-25 seconds
+    4. **Explore Results**: Use the slider to examine attributions at different timesteps
+    
+    The attributions are normalized using Softmax to provide interpretable probability distributions.
+    """)
+    
+    # Main interface tab
+    with gr.Tab("🔍 Attribution Analysis", elem_id="attribution-tab"):
+        
+        # Environment setup
+        gr.Markdown("### 🌙 Environment Setup")
+        env_spec = gr.Dropdown(
+            choices=["LunarLander-v2"],
+            type="value",
+            multiselect=False,
+            label="Environment Specification",
+            value="LunarLander-v2",
+            info="Select the RL environment to analyze"
+        )
+        
+        env_interface = gr.Interface(
+            title="Initialize Environment",
+            allow_flagging="never",
+            inputs=env_spec,
+            fn=create_training_loop,
+            outputs=gr.JSON(label="Environment Spec"),
+            description="Click to initialize the training environment"
+        )
+        
+        # Attribution controls
+        gr.Markdown("### ⚙️ Attribution Configuration")
+        
         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()""")
+            with gr.Column(scale=1):
+                option = gr.Dropdown(
+                    choices=["Torch Tensor of 0's", "Running Average"],
+                    type="index",
+                    label="Baseline Method",
+                    info="Choose the baseline for Integrated Gradients"
+                )
+            
+            with gr.Column(scale=1):
+                baselines = gr.Slider(
+                    label="Number of Baseline Iterations",
+                    interactive=True,
+                    minimum=0,
+                    maximum=100,
+                    value=10,
+                    step=5,
+                    info="Number of baseline inputs to collect for averaging"
+                )
+        
+        # Generate button
+        generate_btn = gr.Button(
+            "🚀 GENERATE ATTRIBUTIONS",
+            variant="primary",
+            size="lg"
+        )
+        generate_btn.click(
+            fn=generate_output,
+            inputs=[baselines, option],
+            outputs=[]
+        )
+        
+        # Results section
+        gr.Markdown("### 📊 Results Visualization")
+        
+        slider = gr.Slider(
+            label="🎬 Key Frame Selector",
+            minimum=0,
+            maximum=1000,
+            step=1,
+            value=0,
+            info="Navigate through different timesteps to see attributions"
+        )
+        
+        results_interface = gr.Interface(
+            fn=get_frame_and_attribution,
+            inputs=slider,
+            live=True,
+            outputs=[
+                gr.Image(label="🎮 Environment State", type="numpy"),
+                gr.Label(label="📈 Feature Attributions", num_top_classes=8)
+            ],
+            title="Real-time Attribution Analysis"
+        )
+    
+    gr.Markdown("""
+    ---
+    
+    ## 🛠️ Local Usage & Installation
+    
+    ### Required Packages
+    ```bash
+    pip install torch gymnasium 'gymnasium[box2d]'
+    ```
+    
+    ### Box2D Installation (macOS)
+    ```bash
+    brew install swig
+    pip install box2d
+    ```
+    
+    ## 🎯 Lunar Lander Environment Details
+    
+    ### Reward Structure
+    - **Position**: Increased/decreased based on distance to landing pad
+    - **Velocity**: Increased/decreased based on speed (slower is better)
+    - **Angle**: Decreased when lander is tilted (horizontal is ideal)
+    - **Landing**: +10 points for each leg touching ground
+    - **Fuel**: -0.03 points per frame for side engine, -0.3 for main engine
+    - **Episode End**: -100 for crash, +100 for safe landing
+    
+    **Success Threshold**: 200+ points per episode
+    
+    ### Training Functions
+    - `load_trained()`: Loads pre-trained model (1000 episodes)
+    - `train()`: Trains from scratch
+    - Set `render_mode=False` for faster training
+    
+    ---
+    
+    *Built with ❤️ using Gradio, PyTorch, and Captum*
+    """)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()

pyproject.toml

@@ -0,0 +1,23 @@
+[project]
+name = "deep-rl-explainability"
+version = "0.1.0"
+description = "Add your description here"
+requires-python = ">=3.10"
+dependencies = [
+    "ale-py==0.8.1",
+    "autorom==0.4.2",
+    "autorom-accept-rom-license==0.6.1",
+    "captum==0.6.0",
+    "gradio>=5.44.1",
+    "gymnasium[box2d]==0.29.1",
+    "huggingface-hub>=0.34.4",
+    "imageio==2.31.5",
+    "imageio-ffmpeg==0.4.9",
+    "matplotlib==3.8.0",
+    "matplotlib-inline==0.1.6",
+    "moviepy==1.0.3",
+    "mujoco==2.3.7",
+    "numpy==1.26.0",
+    "scipy>=1.15.3",
+    "torch==2.1.0",
+]

requirements_hf.txt

@@ -0,0 +1,13 @@
+torch==2.1.0
+gymnasium==0.29.1
+gymnasium[box2d]
+gradio==5.44.1
+captum==0.6.0
+numpy==1.26.0
+scipy
+matplotlib==3.8.0
+moviepy==1.0.3
+imageio==2.31.5
+imageio-ffmpeg==0.4.9
+box2d-py==2.3.5
+swig==4.*