RL-Project/Fetch-Reinforcement_learning_Project

1# <-- this must come first, before any mujoco / gym imports2import os3os.environ["MUJOCO_GL"] = "osmesa"4 5 6import gradio as gr7import numpy as np8import torch9import imageio10from stable_baselines3 import SAC11from custom_env import create_env12 13# Define the function that runs the model and outputs a video14def run_model_episode():15    # 1. Create environment with render_mode="rgb_array" (needed to capture frames)16    # e.g. user inputs:17    # Relative to center of table18    x_start, y_start = 0.0, 0.019    x_targ, y_targ, z_targ = 0.1, 0.1, 0.120 21    env = create_env(render_mode="rgb_array",22                    block_xy=(x_start, y_start),23                    goal_xyz=(x_targ, y_targ, z_targ))24 25    # 2. Load your trained model26    checkpoint_path = os.path.join("model", "model.zip")27    model = SAC.load(checkpoint_path, env=env, verbose=1)28 29    # 3. Rollout the episode30    frames = []31    obs, info = env.reset()32 33    for _ in range(200):  # Shorter rollout to avoid giant videos34        action, _ = model.predict(obs, deterministic=True)35        obs, reward, done, trunc, info = env.step(action)36 37        frame = env.render()  # Get current frame as image (rgb_array)38        frames.append(frame)39 40        if done or trunc:41            obs, info = env.reset()42 43    env.close()44 45    # TODO This will probably need to save into a unique directory 46    # so it doesnt override when multiple people are running the app47 48    # 4. Save the frames into a video49    video_path = "run_video_2.mp4"50    imageio.mimsave(video_path, frames, fps=30)51 52    # 5. Return path to Gradio to display53    return video_path54 55# --------------------------------------56# Build the Gradio App57# --------------------------------------58 59with gr.Blocks() as demo:60    gr.Markdown("Fetch Robot: Model Demo App")61    gr.Markdown("Click 'Run Model' to watch the SAC agent interact with the FetchPickAndPlace environment.")62 63    run_button = gr.Button("Run Model")64    output_video = gr.Video()65 66    run_button.click(fn=run_model_episode, inputs=[], outputs=output_video)67 68demo.launch(share=True)