Fixed naming

app.py

@@ -4,26 +4,28 @@ import numpy as np
 import torch
 import imageio
 from stable_baselines3 import SAC
-from create_env import create_env
+from custom_env import create_env
 
 # Define the function that runs the model and outputs a video
-def run_model_episode():
-    # 1. Create environment with render_mode="rgb_array" (needed to capture frames)
-    env = create_env(render_mode="rgb_array")
+def run_model_episode(x_start, y_start, x_targ, y_targ, z_targ):
+    # Create environment with user inputs
+    env = create_env(render_mode="rgb_array",
+                     block_xy=(x_start, y_start),
+                     goal_xyz=(x_targ, y_targ, z_targ))
 
-    # 2. Load your trained model
-    checkpoint_path = os.path.join("models", "test", "model.zip")
+    # Load your trained model
+    checkpoint_path = os.path.join("App", "model", "model.zip")
     model = SAC.load(checkpoint_path, env=env, verbose=1)
 
-    # 3. Rollout the episode
+    # Rollout the episode
     frames = []
     obs, info = env.reset()
 
-    for _ in range(200):  # Shorter rollout to avoid giant videos
+    for _ in range(200):  # Shorter rollout
         action, _ = model.predict(obs, deterministic=True)
         obs, reward, done, trunc, info = env.step(action)
 
-        frame = env.render()  # Get current frame as image (rgb_array)
+        frame = env.render()
         frames.append(frame)
 
         if done or trunc:
@@ -31,11 +33,10 @@ def run_model_episode():
 
     env.close()
 
-    # 4. Save the frames into a video
+    # Save frames into a video
     video_path = "run_video.mp4"
     imageio.mimsave(video_path, frames, fps=30)
 
-    # 5. Return path to Gradio to display
     return video_path
 
 # --------------------------------------
@@ -43,13 +44,29 @@ def run_model_episode():
 # --------------------------------------
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 Fetch Robot: Model Demo App")
-    gr.Markdown("Click 'Run Model' to watch the SAC agent interact with the FetchPickAndPlace environment.")
+    gr.Markdown("## Fetch Robot: Model Demo App")
+    gr.Markdown("Enter start and target coordinates, then click 'Run Model' to watch the robot!")
+    gr.Markdown("Coordinates are relative to the center of the table.")
+    gr.Markdown("X and Y coordinates are in meters, Z coordinate is height in meters.") 
+    gr.Markdown("0,0,0 is the center of the table.")
+
+    with gr.Row():
+        x_start = gr.Number(label="Start X", value=0.0)
+        y_start = gr.Number(label="Start Y", value=0.0)
+
+    with gr.Row():
+        x_targ = gr.Number(label="Target X", value=0.1)
+        y_targ = gr.Number(label="Target Y", value=0.1)
+        z_targ = gr.Number(label="Target Z", value=0.1)
 
     run_button = gr.Button("Run Model")
     output_video = gr.Video()
 
-    run_button.click(fn=run_model_episode, inputs=[], outputs=output_video)
+    run_button.click(
+        fn=run_model_episode,
+        inputs=[x_start, y_start, x_targ, y_targ, z_targ],
+        outputs=output_video
+    )
 
 demo.launch(share=True)
 

app_test_5.py

@@ -1,72 +0,0 @@
-import gradio as gr
-import os
-import numpy as np
-import torch
-import imageio
-from stable_baselines3 import SAC
-from custom_env import create_env
-
-# Define the function that runs the model and outputs a video
-def run_model_episode(x_start, y_start, x_targ, y_targ, z_targ):
-    # Create environment with user inputs
-    env = create_env(render_mode="rgb_array",
-                     block_xy=(x_start, y_start),
-                     goal_xyz=(x_targ, y_targ, z_targ))
-
-    # Load your trained model
-    checkpoint_path = os.path.join("App", "model", "model.zip")
-    model = SAC.load(checkpoint_path, env=env, verbose=1)
-
-    # Rollout the episode
-    frames = []
-    obs, info = env.reset()
-
-    for _ in range(200):  # Shorter rollout
-        action, _ = model.predict(obs, deterministic=True)
-        obs, reward, done, trunc, info = env.step(action)
-
-        frame = env.render()
-        frames.append(frame)
-
-        if done or trunc:
-            obs, info = env.reset()
-
-    env.close()
-
-    # Save frames into a video
-    video_path = "run_video.mp4"
-    imageio.mimsave(video_path, frames, fps=30)
-
-    return video_path
-
-# --------------------------------------
-# Build the Gradio App
-# --------------------------------------
-
-with gr.Blocks() as demo:
-    gr.Markdown("## Fetch Robot: Model Demo App")
-    gr.Markdown("Enter start and target coordinates, then click 'Run Model' to watch the robot!")
-    gr.Markdown("Coordinates are relative to the center of the table.")
-    gr.Markdown("X and Y coordinates are in meters, Z coordinate is height in meters.") 
-    gr.Markdown("0,0,0 is the center of the table.")
-
-    with gr.Row():
-        x_start = gr.Number(label="Start X", value=0.0)
-        y_start = gr.Number(label="Start Y", value=0.0)
-
-    with gr.Row():
-        x_targ = gr.Number(label="Target X", value=0.1)
-        y_targ = gr.Number(label="Target Y", value=0.1)
-        z_targ = gr.Number(label="Target Z", value=0.1)
-
-    run_button = gr.Button("Run Model")
-    output_video = gr.Video()
-
-    run_button.click(
-        fn=run_model_episode,
-        inputs=[x_start, y_start, x_targ, y_targ, z_targ],
-        outputs=output_video
-    )
-
-demo.launch(share=True)
-