import pybullet as p import pybullet_data import time import numpy as np from PIL import Image import os import math # Global variables to store simulation state obstacleId = None planeId = None def setup_pybullet_environment(): """Setup PyBullet environment with ground plane and obstacle""" global obstacleId, planeId # Connect to PyBullet physics server p.connect(p.GUI) # Set additional search path for URDF files p.setAdditionalSearchPath(pybullet_data.getDataPath()) # Set gravity p.setGravity(0, 0, -9.81) # Load ground plane planeId = p.loadURDF("plane.urdf") # Create obstacle - Box: width=0.5m, depth=0.1m, height=0.05m at x=0.75m obstacle_collision_shape = p.createCollisionShape( p.GEOM_BOX, halfExtents=[0.25, 0.05, 0.025] # Half extents for box ) obstacle_visual_shape = p.createVisualShape( p.GEOM_BOX, halfExtents=[0.25, 0.05, 0.025], rgbaColor=[1, 0, 0, 1] # Red color ) obstacleId = p.createMultiBody( baseMass=0, # Static obstacle baseCollisionShapeIndex=obstacle_collision_shape, baseVisualShapeIndex=obstacle_visual_shape, basePosition=[0.75, 0, 0.025] # Centered at x=0.75m, base on ground ) # Set camera view p.resetDebugVisualizerCamera( cameraDistance=2.0, cameraYaw=0, cameraPitch=-30, cameraTargetPosition=[0.5, 0, 0] ) return obstacleId, planeId def create_vehicle(vehicle_specs): """Create vehicle (robot or drone) based on specifications""" vehicle_type = vehicle_specs.get("vehicle_type", "robot") if vehicle_type == "robot": return create_robot(vehicle_specs) elif vehicle_type == "drone": return create_drone(vehicle_specs) else: raise ValueError(f"Unknown vehicle type: {vehicle_type}") def create_robot(robot_specs): """Create robot based on specifications with corrected wheel orientation and joint axis.""" # Extract specifications with defaults wheel_type = robot_specs.get("wheel_type", "large_smooth") body_clearance_cm = robot_specs.get("body_clearance_cm", 7) # approach_sensor_enabled = robot_specs.get("approach_sensor_enabled", True) # Not used directly in physics main_material = robot_specs.get("main_material", "light_plastic") # Wheel parameters wheel_params = { "small_high_grip": {"radius": 0.06, "friction": 1.5, "width": 0.03}, "large_smooth": {"radius": 0.07, "friction": 0.8, "width": 0.04}, "tracked_base": {"radius": 0.065, "friction": 2.0, "width": 0.05} # Simplified as wide wheels } wheel_config = wheel_params.get(wheel_type, wheel_params["large_smooth"]) wheel_radius = wheel_config["radius"] wheel_friction = wheel_config["friction"] wheel_width = wheel_config["width"] # This is the thickness of the wheel # Body parameters body_length = 0.25 # Along X body_width = 0.20 # Along Y body_height = 0.04 # Along Z obstacle_height = 0.05 min_clearance = obstacle_height + 0.01 body_clearance = max(body_clearance_cm / 100.0, min_clearance) # This is vertical clearance # Calculate Z position for the body's center # The wheels should touch the ground (Z=0), so: # wheel_center_z = wheel_radius (center of wheel above ground) # body_bottom_z = wheel_center_z + body_clearance # body_center_z = body_bottom_z + body_height/2 body_center_z_pos = wheel_radius + body_clearance + (body_height / 2.0) # Material properties material_mass = { "light_plastic": 2.0, "sturdy_metal_alloy": 3.0 } body_mass = material_mass.get(main_material, 2.0) wheel_mass = 0.3 # Create collision shapes body_collision_shape = p.createCollisionShape( p.GEOM_BOX, halfExtents=[body_length/2, body_width/2, body_height/2] ) # For GEOM_CYLINDER, the 'height' is along the cylinder's local Z-axis. # 'radius' is in its local XY plane. # We want wheels whose flat sides are perpendicular to the Y-axis of the robot. # So, the cylinder's length (PyBullet's 'height' param) should be wheel_width. wheel_collision_shape = p.createCollisionShape( p.GEOM_CYLINDER, radius=wheel_radius, height=wheel_width # This is the thickness of the wheel disk ) # Create visual shapes body_visual_shape = p.createVisualShape( p.GEOM_BOX, halfExtents=[body_length/2, body_width/2, body_height/2], rgbaColor=[0, 0, 1, 1] ) # For visual shape, 'length' corresponds to cylinder's Z-axis length wheel_visual_shape = p.createVisualShape( p.GEOM_CYLINDER, radius=wheel_radius, length=wheel_width, # Thickness of the wheel disk rgbaColor=[0.2, 0.2, 0.2, 1] ) # Link properties for two wheels link_masses = [wheel_mass, wheel_mass] link_collision_shape_indices = [wheel_collision_shape, wheel_collision_shape] link_visual_shape_indices = [wheel_visual_shape, wheel_visual_shape] # Position of wheel links relative to the robot's base (body) center # Wheels should be at ground level (Z=0), so their centers are at Z=wheel_radius # Relative to body center: wheel_z = wheel_radius - body_center_z_pos wheel_link_z_offset = wheel_radius - body_center_z_pos # This should be negative link_positions = [ [0, body_width/2 + wheel_width/2, wheel_link_z_offset], # Right wheel [0, -(body_width/2 + wheel_width/2), wheel_link_z_offset] # Left wheel ] # Orientation of wheel links: # Use identity orientation (no rotation) and set joint axis to X directly wheel_link_orientation_quat = p.getQuaternionFromEuler([0, 0, 0]) # No rotation link_orientations = [wheel_link_orientation_quat, wheel_link_orientation_quat] link_inertial_frame_positions = [[0, 0, 0], [0, 0, 0]] # Relative to link frame link_inertial_frame_orientations = [[0,0,0,1], [0,0,0,1]] # Identity quaternion link_parent_indices = [0, 0] # Both wheels attached to the base (index 0 for links) link_joint_types = [p.JOINT_REVOLUTE, p.JOINT_REVOLUTE] # Joint Axis: Use X-axis directly for forward motion link_joint_axis = [[1, 0, 0], [1, 0, 0]] robotId = p.createMultiBody( baseMass=body_mass, baseCollisionShapeIndex=body_collision_shape, baseVisualShapeIndex=body_visual_shape, basePosition=[0, 0, body_center_z_pos], # Initial position of the base baseOrientation=[0,0,0,1], linkMasses=link_masses, linkCollisionShapeIndices=link_collision_shape_indices, linkVisualShapeIndices=link_visual_shape_indices, linkPositions=link_positions, linkOrientations=link_orientations, linkInertialFramePositions=link_inertial_frame_positions, linkInertialFrameOrientations=link_inertial_frame_orientations, linkParentIndices=link_parent_indices, linkJointTypes=link_joint_types, linkJointAxis=link_joint_axis ) # Set dynamics properties for body p.changeDynamics(robotId, -1, lateralFriction=0.8, spinningFriction=0.1, rollingFriction=0.05, # Rolling friction for the body itself if it contacts linearDamping=0.1, angularDamping=0.3) # Set dynamics properties for wheels # Joint indices for createMultiBody start from 0 for the first link. wheel_joint_indices = [0, 1] for joint_idx in wheel_joint_indices: p.changeDynamics(robotId, joint_idx, lateralFriction=wheel_friction, # Friction against sideways slip spinningFriction=0.05, rollingFriction=0.001, # Low rolling friction for the wheel itself linearDamping=0.05, angularDamping=0.1) # Enable motor for velocity control p.setJointMotorControl2(robotId, joint_idx, p.VELOCITY_CONTROL, force=0) print(f"Created robot: body_z_pos={body_center_z_pos:.3f}m, wheel_radius={wheel_radius:.3f}m, actual_clearance_under_body={(body_center_z_pos - body_height/2 - wheel_radius):.3f}m") return robotId, wheel_joint_indices, "robot" def create_drone(drone_specs): """Create drone based on specifications""" # Extract specifications with defaults propeller_size = drone_specs.get("propeller_size", "medium") flight_height_cm = drone_specs.get("flight_height_cm", 20) stability_mode = drone_specs.get("stability_mode", "auto_hover") main_material = drone_specs.get("main_material", "light_carbon_fiber") # Propeller parameters propeller_params = { "small_agile": {"radius": 0.05, "thrust_coeff": 1.2, "mass": 0.1}, "medium": {"radius": 0.08, "thrust_coeff": 1.5, "mass": 0.15}, "large_stable": {"radius": 0.12, "thrust_coeff": 2.0, "mass": 0.2} } prop_config = propeller_params.get(propeller_size, propeller_params["medium"]) prop_radius = prop_config["radius"] thrust_coeff = prop_config["thrust_coeff"] prop_mass = prop_config["mass"] # Body parameters body_length = 0.20 body_width = 0.20 body_height = 0.05 flight_height = max(flight_height_cm / 100.0, 0.15) # Minimum 15cm flight height # Material properties material_mass = { "light_carbon_fiber": 0.8, "sturdy_aluminum": 1.2 } body_mass = material_mass.get(main_material, 0.8) # Calculate starting position - above obstacle body_z_pos = flight_height prop_offset = body_length / 2 + prop_radius / 2 # Create collision shapes body_collision_shape = p.createCollisionShape( p.GEOM_BOX, halfExtents=[body_length/2, body_width/2, body_height/2] ) prop_collision_shape = p.createCollisionShape( p.GEOM_CYLINDER, radius=prop_radius, height=0.01 # Very thin propellers ) # Create visual shapes body_visual_shape = p.createVisualShape( p.GEOM_BOX, halfExtents=[body_length/2, body_width/2, body_height/2], rgbaColor=[0, 1, 0, 1] # Green body for drone ) prop_visual_shape = p.createVisualShape( p.GEOM_CYLINDER, radius=prop_radius, length=0.01, rgbaColor=[0.1, 0.1, 0.1, 0.8] # Semi-transparent dark propellers ) # Create the drone with body and 4 propellers link_masses = [prop_mass] * 4 # Four propellers link_collision_shape_indices = [prop_collision_shape] * 4 link_visual_shape_indices = [prop_visual_shape] * 4 link_positions = [ [prop_offset, prop_offset, 0.03], # Front right [-prop_offset, prop_offset, 0.03], # Front left [-prop_offset, -prop_offset, 0.03], # Rear left [prop_offset, -prop_offset, 0.03] # Rear right ] link_orientations = [[0, 0, 0, 1]] * 4 link_inertial_frame_positions = [[0, 0, 0]] * 4 link_inertial_frame_orientations = [[0, 0, 0, 1]] * 4 link_parent_indices = [0, 0, 0, 0] # All propellers connected to base link_joint_types = [p.JOINT_REVOLUTE] * 4 # Revolute joints for propellers link_joint_axis = [[0, 0, 1]] * 4 # Rotate around Z-axis (vertical) droneId = p.createMultiBody( baseMass=body_mass, baseCollisionShapeIndex=body_collision_shape, baseVisualShapeIndex=body_visual_shape, basePosition=[0, 0, body_z_pos], linkMasses=link_masses, linkCollisionShapeIndices=link_collision_shape_indices, linkVisualShapeIndices=link_visual_shape_indices, linkPositions=link_positions, linkOrientations=link_orientations, linkInertialFramePositions=link_inertial_frame_positions, linkInertialFrameOrientations=link_inertial_frame_orientations, linkParentIndices=link_parent_indices, linkJointTypes=link_joint_types, linkJointAxis=link_joint_axis ) # Set dynamics properties for body p.changeDynamics(droneId, -1, lateralFriction=0.1, spinningFriction=0.1, rollingFriction=0.1, linearDamping=0.3, angularDamping=0.5) # Set dynamics properties for propellers for prop_idx in range(4): p.changeDynamics(droneId, prop_idx, lateralFriction=0.1, spinningFriction=0.05, rollingFriction=0.01, linearDamping=0.2, angularDamping=0.3) # Store thrust coefficient for flight control drone_props = { "thrust_coeff": thrust_coeff, "target_height": flight_height, "stability_mode": stability_mode } propeller_joint_indices = [0, 1, 2, 3] # Joint indices for the four propellers print(f"Created drone: flight_height={flight_height:.3f}m, prop_radius={prop_radius:.3f}m, thrust_coeff={thrust_coeff}") return droneId, propeller_joint_indices, "drone", drone_props def run_simulation_step(vehicleId, joint_indices, control_params, vehicle_type="robot", vehicle_props=None): """Run one simulation step with vehicle control""" if vehicle_type == "robot": run_robot_simulation_step(vehicleId, joint_indices, control_params) elif vehicle_type == "drone": run_drone_simulation_step(vehicleId, joint_indices, control_params, vehicle_props) # Step simulation p.stepSimulation() time.sleep(1./240.) # 240 Hz simulation def run_robot_simulation_step(robotId, wheel_joint_indices, control_params): """Run robot simulation step with wheel control""" if wheel_joint_indices: # Robot has wheel joints # Set target velocity for forward motion - TRY POSITIVE direction target_velocity = 5.0 # rad/s - positive for forward motion max_force = 50.0 # Nm - much more torque for climbing # Apply velocity control to both wheels for forward motion for joint_idx in wheel_joint_indices: p.setJointMotorControl2( robotId, joint_idx, p.VELOCITY_CONTROL, targetVelocity=target_velocity, force=max_force ) else: # Fallback: apply direct force force_magnitude = 5.0 # Newtons p.applyExternalForce(robotId, -1, [force_magnitude, 0, 0], [0, 0, 0], p.WORLD_FRAME) def run_drone_simulation_step(droneId, propeller_joint_indices, control_params, drone_props): """Run drone simulation step with flight control""" # Get current drone state drone_pos, drone_orn = p.getBasePositionAndOrientation(droneId) drone_vel, drone_ang_vel = p.getBaseVelocity(droneId) target_height = drone_props.get("target_height", 0.2) thrust_coeff = drone_props.get("thrust_coeff", 1.5) stability_mode = drone_props.get("stability_mode", "auto_hover") # Calculate thrust needed for hovering drone_mass = p.getDynamicsInfo(droneId, -1)[0] gravity_force = drone_mass * 9.81 base_thrust_per_prop = gravity_force / 4 # Four propellers # Height control (PID-like) height_error = target_height - drone_pos[2] height_thrust_correction = height_error * 10.0 # Higher proportional gain # Forward motion - apply body force instead of individual propeller forces forward_force = 3.0 # Newtons - direct forward force # Apply main thrust for hovering total_thrust = (base_thrust_per_prop + height_thrust_correction) * thrust_coeff # Apply upward thrust at drone center if total_thrust > 0: p.applyExternalForce( droneId, -1, [0, 0, total_thrust * 4], # Total upward force [0, 0, 0], # At center of mass p.WORLD_FRAME ) # Apply forward force directly to drone body p.applyExternalForce( droneId, -1, [forward_force, 0, 0], # Forward force [0, 0, 0], # At center of mass p.WORLD_FRAME ) # Add slight damping to prevent oscillations linear_damping = -0.1 p.applyExternalForce( droneId, -1, [drone_vel[0] * linear_damping, drone_vel[1] * linear_damping, 0], [0, 0, 0], p.WORLD_FRAME ) # Spin propellers for visual effect if propeller_joint_indices: for prop_idx in propeller_joint_indices: p.setJointMotorControl2( droneId, prop_idx, p.VELOCITY_CONTROL, targetVelocity=20.0, # Fast spinning for visual effect force=0.1 ) def get_simulation_feedback(vehicleId, obstacleId, start_time, current_sim_time, vehicle_type="robot"): """Get feedback from current simulation state""" # Get vehicle position and orientation vehicle_pos, vehicle_orn = p.getBasePositionAndOrientation(vehicleId) # Check if vehicle is upright/stable euler_angles = p.getEulerFromQuaternion(vehicle_orn) roll, pitch, yaw = euler_angles if vehicle_type == "robot": # Robot is considered upright if roll and pitch are small is_stable = abs(roll) < 0.5 and abs(pitch) < 0.5 elif vehicle_type == "drone": # Drone is considered stable if not completely inverted and at reasonable height is_stable = abs(roll) < 1.0 and abs(pitch) < 1.0 and vehicle_pos[2] > 0.05 else: is_stable = True # Check for contacts with obstacle contact_points = p.getContactPoints(vehicleId, obstacleId) obstacle_contacts_exist = len(contact_points) > 0 feedback = { "robot_position": list(vehicle_pos), # Keep "robot_position" for compatibility "robot_orientation_quaternion": list(vehicle_orn), "obstacle_contacts_exist": obstacle_contacts_exist, "is_robot_upright": is_stable, # Keep "is_robot_upright" for compatibility "current_sim_time_sec": current_sim_time, "vehicle_type": vehicle_type } return feedback def reset_simulation(): """Reset and disconnect from PyBullet simulation""" p.resetSimulation() p.disconnect() def capture_frame(): """Capture current frame from simulation""" # Get camera image width, height, rgb_img, depth_img, seg_img = p.getCameraImage( width=640, height=480, viewMatrix=p.computeViewMatrixFromYawPitchRoll( cameraTargetPosition=[0.5, 0, 0], distance=2.0, yaw=0, pitch=-30, roll=0, upAxisIndex=2 ), projectionMatrix=p.computeProjectionMatrixFOV( fov=60, aspect=640/480, nearVal=0.1, farVal=100.0 ) ) # Convert to PIL Image rgb_array = np.array(rgb_img).reshape(height, width, 4)[:, :, :3] # Remove alpha channel image = Image.fromarray(rgb_array, 'RGB') return image def get_obstacle_info(): """Get information about the obstacle""" return { "width_m": 0.5, "depth_m": 0.1, "height_m": 0.05, "position_x_m": 0.75, "position_y_m": 0.0, "position_z_m": 0.025, "material": "static_red_box", "success_threshold_x_m": 0.8 }