app.py

from models.TD3.TD3 import TD3

import torch
import numpy as np
from sim import SIM_ENV
import yaml
import gradio as gr
import os
from pathlib import Path
import matplotlib
import random

matplotlib.use('Agg')

def generate_random_points(num_scenarios=2):
    """Generate random robot poses and goals"""
    robot_poses = []
    robot_goals = []
    
    for _ in range(num_scenarios):
        # Random pose: x, y, orientation, velocity
        pose = [
            [random.uniform(1, 9)],  # x position
            [random.uniform(1, 9)],  # y position
            [random.uniform(0, 3.14)],  # orientation
            [0]  # initial velocity
        ]
        
        # Random goal: x, y, orientation
        goal = [
            [random.uniform(1, 9)],  # x position
            [random.uniform(1, 9)],  # y position
            [0]  # orientation
        ]
        
        robot_poses.append(pose)
        robot_goals.append(goal)
    
    return robot_poses, robot_goals

def get_predefined_scenarios():
    """Return predefined robot poses and goals"""
    robot_poses = [
        [[3], [4], [0], [0]],
        [[8], [1], [1], [0]],
        [[2], [6], [1], [0]],
        [[7], [1], [0], [0]],
        [[7], [6.5], [2], [0]],
        [[9], [9], [3], [0]],
        [[2], [9], [1], [0]],
        [[3], [6], [3], [0]],
        [[1], [7], [0], [0]],
        [[5], [7], [3], [0]]
    ]
    
    robot_goals = [
        [[8], [8], [0]],
        [[2], [9], [0]],
        [[7], [1], [0]],
        [[7.2], [9], [0]],
        [[1], [1], [0]],
        [[5], [1], [0]],
        [[7], [4], [0]],
        [[9], [4], [0]],
        [[1], [9], [0]],
        [[5], [1], [0]]
    ]
    
    return robot_poses, robot_goals

def run_simulation():
    """Run the simulation and return the path to the generated GIF"""
    action_dim = 2  # number of actions produced by the model
    max_action = 1  # maximum absolute value of output actions
    state_dim = 25  # number of input values in the neural network (vector length of state input)
    device = torch.device(
        "cuda" if torch.cuda.is_available() else "cpu"
    )  # using cuda if it is available, cpu otherwise
    epoch = 0  # epoch number
    max_steps = 300  # maximum number of steps in single episode

    model = TD3(
        state_dim=state_dim,
        action_dim=action_dim,
        max_action=max_action,
        device=device,
        load_model=True,
    )  # instantiate a model

    sim = SIM_ENV(world_file="eval_world.yaml", save_ani=True)  # instantiate environment
    
    # Generate random evaluation points instead of loading from YAML
    # num_scenarios = random.randint(2, 5)  # Random number of scenarios between 2 and 5
    
    # Get predefined scenarios and randomly select one
    all_poses, all_goals = get_predefined_scenarios()
    scenario_index = random.randint(0, len(all_poses) - 1)
    
    robot_poses = [all_poses[scenario_index]]
    robot_goals = [all_goals[scenario_index]]
    
    print(f"Selected scenario {scenario_index+1} of {len(all_poses)}")
    total_reward = 0.0
    total_steps = 0
    col = 0
    goals = 0
    for idx in range(len(robot_poses)):
        count = 0
        latest_scan, distance, cos, sin, collision, goal, a, reward = sim.reset(
            robot_state=robot_poses[idx],
            robot_goal=robot_goals[idx],
            random_obstacles=False,
        )
        done = False
        while not done and count < max_steps:
            state, terminal = model.prepare_state(
                latest_scan, distance, cos, sin, collision, goal, a
            )
            action = model.get_action(np.array(state), False)
            a_in = [(action[0] + 1) / 4, action[1]]
            latest_scan, distance, cos, sin, collision, goal, a, reward = sim.step(
                lin_velocity=a_in[0], ang_velocity=a_in[1]
            )
            total_reward += reward
            total_steps += 1
            count += 1
            if collision:
                col += 1
            if goal:
                goals += 1
            done = collision or goal
    avg_step_reward = total_reward / total_steps
    avg_reward = total_reward / len(robot_poses)
    avg_col = col / len(robot_poses)
    avg_goal = goals / len(robot_poses)
    print(f"Total Reward: {total_reward}")
    print(f"Average Reward: {avg_reward}")
    print(f"Average Step Reward: {avg_step_reward}")
    print(f"Average Collision rate: {avg_col}")
    print(f"Average Goal rate: {avg_goal}")
    print("..............................................")
    model.writer.add_scalar("test/total_reward", total_reward, epoch)
    model.writer.add_scalar("test/avg_reward", avg_reward, epoch)
    model.writer.add_scalar("test/avg_step_reward", avg_step_reward, epoch)
    model.writer.add_scalar("test/avg_col", avg_col, epoch)
    model.writer.add_scalar("test/avg_goal", avg_goal, epoch)
    
    sim.env.end(ending_time=3)
    
    # Find the latest generated GIF file in the animation folder
    animation_dir = Path("animation")
    if animation_dir.exists():
        gif_files = list(animation_dir.glob("*.gif"))
        if gif_files:
            # Sort by creation time (newest first)
            latest_gif = max(gif_files, key=lambda x: x.stat().st_ctime)
            return str(latest_gif), {
                "Total Reward": f"{total_reward:.2f}",
                "Average Reward": f"{avg_reward:.2f}",
                "Average Step Reward": f"{avg_step_reward:.2f}",
                "Collision Rate": f"{avg_col:.2f}",
                "Goal Rate": f"{avg_goal:.2f}"
            }
    
    return None, {"Error": "No GIF file was generated"}

def get_default_data():
    """Return default animation and statistics data"""
    # Find any existing GIF in the animation folder
    animation_dir = Path("animation")
    default_gif = None
    
    if animation_dir.exists():
        gif_files = list(animation_dir.glob("*.gif"))
        if gif_files:
            # Get the most recent GIF
            default_gif = str(max(gif_files, key=lambda x: x.stat().st_ctime))
    
    # Default statistics
    default_stats = {
        "Total Reward": "99.12",
        "Average Reward": "99.12",
        "Average Step Reward": "1.40",
        "Collision Rate": "0.00",
        "Goal Rate": "1.00"
    }
    
    return default_gif, default_stats

def main(args=None):
    """Main function with Gradio interface"""
    # Get default data for initial display
    default_gif, default_stats = get_default_data()
    
    with gr.Blocks(title="Robot Navigation Simulation") as demo:
        with gr.Row():
            with gr.Column():
                run_button = gr.Button("Run Simulation", variant="primary")
                
                with gr.Row():
                    with gr.Column():
                        output_image = gr.Image(
                            type="filepath", 
                            label="Simulation Animation",
                            value=default_gif  # Set default value
                        )
                    
                    with gr.Column():
                        output_stats = gr.JSON(
                            label="Simulation Statistics",
                            value=default_stats  # Set default value
                        )
        
        run_button.click(
            fn=run_simulation,
            outputs=[output_image, output_stats]
        )
    
    demo.launch(share=False)


if __name__ == "__main__":
    main()