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 imageio10import time11from stable_baselines3 import SAC12from custom_env import create_env13 14def stream_frames():15    x_start, y_start = 0.0, 0.016    x_targ, y_targ, z_targ = 0.1, 0.1, 0.117 18    env = create_env(render_mode="rgb_array",19                     block_xy=(x_start, y_start),20                     goal_xyz=(x_targ, y_targ, z_targ))21 22    checkpoint_path = os.path.join("App", "model", "model.zip")23    model = SAC.load(checkpoint_path, env=env, verbose=1)24 25    obs, info = env.reset()26 27    while True:28        action, _ = model.predict(obs, deterministic=True)29        obs, reward, done, trunc, info = env.step(action)30 31        frame = env.render()  # Grab RGB frame32        yield frame  # Yield this frame to Gradio33 34        if done or trunc:35            obs, info = env.reset()36 37        time.sleep(0.033)  # ~30 FPS (1/30 seconds)38 39    env.close()40 41# Build Gradio app42with gr.Blocks() as demo:43    gr.Markdown("Fetch Robot: Live Model Demo App")44    frame_output = gr.Image()45    start_button = gr.Button("Start Streaming")46 47    start_button.click(fn=stream_frames, inputs=[], outputs=frame_output)48 49demo.queue()50demo.launch(share=True)