Create Model.py

Model.py

@@ -0,0 +1,24 @@
+import torch
+import numpy as np
+from diffusion_policy.model.diffusion_policy import DiffusionPolicy
+
+# Load pretrained model
+def load_model(device="cuda" if torch.cuda.is_available() else "cpu"):
+    model_path = "diffusion-policy/pretrained/6dof_robot"  # Replace with actual model path
+    model = DiffusionPolicy.load_from_checkpoint(model_path, map_location=device)
+    model.to(device)
+    model.eval()
+    return model
+
+# Generate trajectory for a 6-DOF robot
+def generate_trajectory(model, start_state, goal_state, num_steps=50, device="cuda"):
+    trajectory = torch.zeros((num_steps, 6), device=device)  # Store trajectory
+    state = torch.tensor(start_state, dtype=torch.float32, device=device)
+
+    with torch.no_grad():
+        for t in range(num_steps):
+            action = model.predict(state)  # Predict next action
+            state = state + action  # Simulate movement
+            trajectory[t] = state
+
+    return trajectory.cpu().numpy()  # Convert to NumPy for plotting