streamlitapp.py

import streamlit as st
import torch
import numpy as np
import gymnasium as gym
import time
import os

from DQN import DQN, FrameStack, DEVICE
from PPO import ActorCritic

DQN_MODEL_PATH = "dqn_spaceinvaders.pth"
PPO_MODEL_PATH = "ppo_spaceinvaders.pth"
ENV_NAME = "ALE/SpaceInvaders-v5"
STATE_SHAPE = (4, 84, 84)
FRAME_DELAY_SECONDS = 0.03

st.title("🎮 Space Invaders - DQN vs PPO Demo")
st.markdown("Created By: Rhichard Koh")
st.markdown("""
This app lets you run a trained agent (DQN or PPO) to play Space Invaders in real-time.

- **DQN**: Uses value-based greedy policy  
- **PPO**: Uses policy gradient and actor-critic structure  

Due to Streamlit's rendering requirements, the OpenAI Gym environment is set to 'rgb_array'
mode instead of 'human'. As a result, some dynamic visual elements like enemy missiles 
may not be visible in the rendered frames, since they are often rendered in intermediate or 
rapidly updating frames that are missed in this display mode.
""")

@st.cache_resource
def load_env():
    env = gym.make(ENV_NAME, render_mode="rgb_array")
    return env, env.action_space.n

@st.cache_resource
def load_model(agent_type):
    _, n_actions = load_env()

    if agent_type == "DQN":
        model = DQN(STATE_SHAPE, n_actions).to(DEVICE)
        path = DQN_MODEL_PATH
    else:
        model = ActorCritic(STATE_SHAPE, n_actions).to(DEVICE)
        path = PPO_MODEL_PATH

    if os.path.exists(path):
        state_dict = torch.load(path, map_location=DEVICE)
        model.load_state_dict(state_dict)
        model.eval()
        return model

    st.error(f"Model not found: {path}")
    st.stop()

def to_state_tensor(state):
    return torch.as_tensor(state, dtype=torch.float32, device=DEVICE).unsqueeze(0)

def select_action_dqn(model, state):
    with torch.inference_mode():
        state_tensor = to_state_tensor(state)
        return model(state_tensor).argmax(dim=1).item()

def select_action_ppo(model, state):
    with torch.inference_mode():
        state_tensor = to_state_tensor(state)
        logits, _ = model(state_tensor)
        dist = torch.distributions.Categorical(logits=logits)
        return dist.sample().item()

def run_episode(env, model, agent_type, frame_placeholder=None):
    frame_stack = FrameStack(STATE_SHAPE[0])
    state, _ = env.reset()
    state = frame_stack.reset(state)
    total_reward = 0
    done = False

    while not done:
        if agent_type == "DQN":
            action = select_action_dqn(model, state)
        else:
            action = select_action_ppo(model, state)

        obs, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        state = frame_stack.step(obs)
        total_reward += reward

        if frame_placeholder is not None:
            frame = env.render()
            frame_placeholder.image(
                frame,
                caption=f"Reward: {total_reward:.0f}",
                channels="RGB",
                use_container_width=True,
            )
            time.sleep(FRAME_DELAY_SECONDS)

    return total_reward

def benchmark_agents(env, num_games):
    dqn_model = load_model("DQN")
    ppo_model = load_model("PPO")
    progress_bar = st.progress(0)

    dqn_scores = []
    ppo_scores = []

    for game_idx in range(num_games):
        dqn_scores.append(run_episode(env, dqn_model, "DQN"))
        progress_bar.progress((2 * game_idx + 1) / (2 * num_games))

        ppo_scores.append(run_episode(env, ppo_model, "PPO"))
        progress_bar.progress((2 * game_idx + 2) / (2 * num_games))

    progress_bar.empty()
    return np.mean(dqn_scores), np.mean(ppo_scores)

env, _ = load_env()
agent_type = st.selectbox("Select Agent", ["DQN", "PPO"])

if st.button("▶️ Run Agent"):
    model = load_model(agent_type)
    stframe = st.empty()
    total_reward = run_episode(env, model, agent_type, frame_placeholder=stframe)
    st.success(f"🎉 Game Over! Total Reward: {total_reward:.0f}")

if st.button("📊 Run 100 Games & Compare DQN vs PPO"):
    st.info("Running 100 games for each agent... Please wait ⏳")
    avg_dqn, avg_ppo = benchmark_agents(env, num_games=100)

    st.subheader("🏁 Results after 100 games")
    st.write(f"**DQN Average Reward:** {avg_dqn:.2f}")
    st.write(f"**PPO Average Reward:** {avg_ppo:.2f}")

    if avg_dqn > avg_ppo:
        st.success("✅ **DQN performed better on average!**")
    elif avg_ppo > avg_dqn:
        st.success("✅ **PPO performed better on average!**")
    else:
        st.warning("⚖️ Both agents performed equally!")