visualize.py

import time
import numpy as np
import pybullet as p
from environment import SimulationEnv
from enhanced_apf_rrt import Enhanced_APF_RRT
from pso_smoother import PSOSmoother

def draw_tree_in_pybullet(env, planner):
    """
    Physically draws the RRT search tree in the PyBullet 3D environment 
    using the robot's end-effector positions.
    """
    print("Drawing RRT Tree in PyBullet 3D space...")
    eef_idx = 11 # Franka Panda End-Effector
    
    # Dictionary to cache positions so we don't recalculate
    eef_positions = {}
    
    # Calculating End-Effector position for every node in the tree
    for node in planner.tree:
        for j, val in enumerate(node.q):
            p.resetJointState(env.robot_id, j, val)
        eef_positions[node] = p.getLinkState(env.robot_id, eef_idx)[0]
        
    # Drawing the branches
    for node in planner.tree:
        if node.parent:
            p1 = eef_positions[node.parent]
            p2 = eef_positions[node]
            # Drawing faint blue lines for the tree branches
            p.addUserDebugLine(p1, p2, lineColorRGB=[0, 0.4, 0.8], lineWidth=1, lifeTime=0)

def animate_with_trail(env, path, fps=30):
    """Animates the robot and leaves a glowing trail."""
    print("Executing cinematic trajectory...")
    dt = 1.0 / fps
    steps_per_waypoint = 40 # High number for buttery smooth slow-motion
    prev_pos = None
    eef_idx = 11

    for i in range(len(path) - 1):
        q_start = np.array(path[i])
        q_end = np.array(path[i+1])
        
        for step in range(steps_per_waypoint):
            alpha = step / float(steps_per_waypoint)
            q_interp = q_start + alpha * (q_end - q_start)
            
            for j, val in enumerate(q_interp):
                p.resetJointState(env.robot_id, j, val)
            p.stepSimulation()

            # Drawing trail
            curr_pos = p.getLinkState(env.robot_id, eef_idx)[0]
            if prev_pos is not None:
                # Drawing a thick Red/Cyan line for the final optimized path
                p.addUserDebugLine(prev_pos, curr_pos, [1, 0.2, 0.2], 4, 0)
            prev_pos = curr_pos
            
            time.sleep(dt)

def main():
    print("Initializing PyBullet GUI...")
    env = SimulationEnv(use_gui=True)
    env.load_robot()
    
    # Force the Curved Obstacle Course
    # env.spawn_obstacle_course() 
    
    env.spawn_dense_forest()

    # Start and Goal (Forces movement across the central pillar)
    start_q = [0.0, -0.78, 0.0, -2.35, 0.0, 1.57, 0.78]
    goal_q = [1.57, -0.5, 0.0, -2.0, 0.0, 1.57, 0.78]

    print("Planning Enhanced Path...")
    planner = Enhanced_APF_RRT(env.robot_id, env.obstacles, start_q, goal_q)
    
    found = False
    for _ in range(planner.max_iter):
        if planner.expand_tree():
            found = True
            break
            
    if not found:
        print("Failed to find path. Please run again.")
        p.disconnect()
        return

    # Visualizing the Tree in PyBullet
    draw_tree_in_pybullet(env, planner)

    # Extracting and Optimize the Path
    print("Optimizing Path with PSO...")
    raw_path = planner.extract_path(planner.tree[-1])
    short_path = planner.shortcut_path(raw_path, 150)
    
    smoother = PSOSmoother(planner, short_path, iterations=100)
    smoother.w_len = 2.0     # Lower weight to allow for a curve
    smoother.w_smooth = 10.0 # High weight for a fluid motion
    final_path = smoother.optimize()
    
    from plotter import plot_curved_path
    plot_curved_path(planner, final_path)

    # Reseting robot to Start Position before recording
    for j, val in enumerate(start_q):
        p.resetJointState(env.robot_id, j, val)

    # Recording and Animate
    print("Starting Video Recording...")
    log_id = p.startStateLogging(p.STATE_LOGGING_VIDEO_MP4, "ultimate_trajectory.mp4")
    
    # Pause for a second so you can see the tree before the robot moves
    time.sleep(2) 
    
    # Executing the motion
    animate_with_trail(env, final_path)
    
    time.sleep(2)
    p.stopStateLogging(log_id)
    print("Simulation Complete! Video saved as 'ultimate_trajectory.mp4'.")
    
    # Keeping the window open so you can inspect the 3D tree
    input("Press Enter in the terminal to close PyBullet...")
    p.disconnect()

if __name__ == "__main__":
    main()