Add application file

app.py

@@ -0,0 +1,135 @@
+import os
+if os.getenv("SPACE_ID") is not None:
+    os.environ["SDL_VIDEODRIVER"] = "dummy"
+    os.environ["SDL_AUDIODRIVER"] = "dummy"
+
+import torch
+import torch.nn as nn
+import numpy as np
+import gymnasium as gym
+import imageio
+import gradio as gr
+from huggingface_hub import hf_hub_download
+
+# 1. Policy Network Architecture
+def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
+    nn.init.orthogonal_(layer.weight, std)
+    nn.init.constant_(layer.bias, bias_const)
+    return layer
+
+# Reconstruct pure Sequential actor
+def get_actor_network(state_dim=8, action_dim=4):
+    actor = nn.Sequential(
+        layer_init(nn.Linear(state_dim, 64)),
+        nn.Tanh(),
+        layer_init(nn.Linear(64, 64)),
+        nn.Tanh(),
+        layer_init(nn.Linear(64, action_dim), std=0.01),
+    )
+    return actor
+
+# 2. Inference & Rendering
+def simulate_agent(stage_selection):
+    weight_mapping = {
+        "Stage 1: Baseline": "1_baseline.pth",
+        "Stage 2: Surrogate Hacking": "2_surrogate_hacking_attention.pth",
+        "Stage 3: Temporal Paradox ": "3_temporal_paradox_variance.pth",
+        "Stage 4: Target Decoupling": "4_target_decoupling_final.pth"
+    }
+    filename = weight_mapping.get(stage_selection)
+    repo_id = "ben-dlwlrma/Representation-Over-Routing" 
+    
+    # Download weights from HF Hub
+    try:
+        weights_path = hf_hub_download(repo_id=repo_id, filename=filename)
+    except Exception as e:
+        return None, f"Weight download failed. Error: {str(e)}"
+
+    # Initialize env
+    env = gym.make("LunarLander-v2", render_mode="rgb_array")
+    
+    # Initialize model on CPU
+    device = torch.device("cpu")
+    actor = get_actor_network(state_dim=8, action_dim=4).to(device)
+    
+    # Load weights
+    try:
+        actor.load_state_dict(torch.load(weights_path, map_location=device, weights_only=True))
+        actor.eval()
+    except Exception as e:
+        env.close()
+        return None, f"Architecture mismatch. Error: {str(e)}"
+
+    state, _ = env.reset(seed=32)
+    done = False
+    frames = []
+    total_reward = 0.0
+    step_count = 0
+
+    while not done and step_count < 600:
+        try:
+            frame = env.render()
+            if frame is not None:
+                frames.append(frame)
+        except Exception as e:
+            env.close()
+            return None, f"Render failed: {str(e)}"
+            
+        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
+        with torch.no_grad():
+            action_logits = actor(state_tensor)
+            action = torch.argmax(action_logits, dim=1).item()
+            
+        state, reward, terminated, truncated, _ = env.step(action)
+        total_reward += reward
+        step_count += 1
+        done = terminated or truncated
+
+    env.close()
+
+    # Export to MP4
+    video_filename = "eval_output.mp4"
+    fps = 30
+    try:
+        imageio.mimsave(video_filename, frames, fps=fps, codec='libx264', pixelformat='yuv420p')
+    except Exception as e:
+        return None, f"Video encoding failed: {str(e)}"
+
+    logs = (f"Status: Inference complete\n"
+            f"Stage: {stage_selection}\n"
+            f"Total Reward: {total_reward:.2f}\n"
+            f"Steps: {step_count}")
+    
+    return video_filename, logs
+
+# 3. Gradio Web UI
+with gr.Blocks(title="Representation over Routing", theme=gr.themes.Base()) as demo:
+    gr.Markdown("## Representation over Routing")
+    gr.Markdown("Multi-timescale RL evaluation environment. Select an ablation stage to visualize policy behavior.")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            model_dropdown = gr.Dropdown(
+                choices=[
+                    "Stage 1: Baseline",
+                    "Stage 2: Surrogate Hacking",
+                    "Stage 3: Temporal Paradox ",
+                    "Stage 4: Target Decoupling"
+                ],
+                value="Stage 4: Target Decoupling",
+                label="Model Stage"
+            )
+            run_button = gr.Button("Run Inference", variant="primary")
+            
+        with gr.Column(scale=2):
+            video_output = gr.Video(label="Environment Render")
+            text_output = gr.Textbox(label="Execution Logs", lines=4)
+
+    run_button.click(
+        fn=simulate_agent,
+        inputs=[model_dropdown],
+        outputs=[video_output, text_output]
+    )
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,8 @@
+torch>=2.0.0
+numpy
+gymnasium[box2d]
+imageio
+imageio-ffmpeg
+huggingface_hub
+gradio
+spaces