import gymnasium as gym from gymnasium import spaces import pybullet as p import pybullet_data import numpy as np import os from mag_lev_coil import MagLevCoil from maglev_predictor import MaglevPredictor TARGET_GAP = 16.491741 / 1000 # target gap height for 5.8 kg pod in meters class LevPodEnv(gym.Env): def __init__(self, use_gui=False, initial_gap_mm=10.0): super(LevPodEnv, self).__init__() # Store initial gap height parameter self.initial_gap_mm = initial_gap_mm # --- 1. Define Action & Observation Spaces --- # Action: 4 PWM duty cycles between -1 and 1 (4 independent coils) # [front_left, front_right, back_left, back_right] corresponding to +Y and -Y ends of each U-yoke self.action_space = spaces.Box(low=-1, high=1, shape=(4,), dtype=np.float32) # Observation: 4 noisy sensor gap heights + 4 velocities # [sensor_center_right, sensor_center_left, sensor_front, sensor_back, # vel_center_right, vel_center_left, vel_front, vel_back] self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(8,), dtype=np.float32) # --- 2. Setup Physics & Actuators --- self.dt = 1./240. # PyBullet default timestep # Four coils: one for each end of the two U-shaped yokes # Front yoke: left (+Y) and right (-Y) ends # Back yoke: left (+Y) and right (-Y) ends self.coil_front_L = MagLevCoil(1.1, 0.0025, 12, 10.2) # From your specs self.coil_front_R = MagLevCoil(1.1, 0.0025, 12, 10.2) self.coil_back_L = MagLevCoil(1.1, 0.0025, 12, 10.2) self.coil_back_R = MagLevCoil(1.1, 0.0025, 12, 10.2) # Sensor noise parameters self.sensor_noise_std = 0.0001 # 0.1mm standard deviation # Maglev force/torque predictor self.predictor = MaglevPredictor() # Connect to PyBullet (DIRECT is faster for training, GUI for debugging) self.client = p.connect(p.GUI if use_gui else p.DIRECT) p.setAdditionalSearchPath(pybullet_data.getDataPath()) # Store references self.trackId = None self.podId = None self.collision_local_positions = [] self.yoke_indices = [] # For force application self.yoke_labels = [] self.sensor_indices = [] # For gap height measurement self.sensor_labels = [] # For velocity calculation self.prev_sensor_gaps = None def reset(self, seed=None, options=None): # Reset PyBullet simulation p.resetSimulation(physicsClientId=self.client) p.setGravity(0, 0, -9.81, physicsClientId=self.client) p.setTimeStep(self.dt, physicsClientId=self.client) # Create the maglev track (inverted system - track above, pod hangs below) # Track bottom surface at Z=0 track_collision = p.createCollisionShape( shapeType=p.GEOM_BOX, halfExtents=[1.0, 0.2, 0.010], physicsClientId=self.client ) track_visual = p.createVisualShape( shapeType=p.GEOM_BOX, halfExtents=[1.0, 0.2, 0.010], rgbaColor=[0.3, 0.3, 0.3, 0.8], physicsClientId=self.client ) self.trackId = p.createMultiBody( baseMass=0, # Static baseCollisionShapeIndex=track_collision, baseVisualShapeIndex=track_visual, basePosition=[0, 0, 0.010], # Track center at Z=10mm, bottom at Z=0 physicsClientId=self.client ) p.changeDynamics(self.trackId, -1, lateralFriction=0.3, restitution=0.1, physicsClientId=self.client) # Load pod URDF - hangs below track # Create modified URDF with gap-dependent bolt positions urdf_path = self._create_modified_urdf() # Calculate starting position based on desired initial gap # Pod center Z = -(90.85mm + initial_gap_mm) / 1000 start_z = -(0.09085 + self.initial_gap_mm / 1000.0) start_pos = [0, 0, start_z] start_orientation = p.getQuaternionFromEuler([0, 0, 0]) self.podId = p.loadURDF(urdf_path, start_pos, start_orientation, physicsClientId=self.client) # Parse collision shapes to identify yokes and sensors collision_shapes = p.getCollisionShapeData(self.podId, -1, physicsClientId=self.client) self.collision_local_positions = [] self.yoke_indices = [] self.yoke_labels = [] self.sensor_indices = [] self.sensor_labels = [] # Expected heights for detection expected_yoke_sensor_z = 0.08585 # Yokes and sensors always at this height expected_bolt_z = 0.08585 + self.initial_gap_mm / 1000.0 # Bolts at gap-dependent height for i, shape in enumerate(collision_shapes): shape_type = shape[2] local_pos = shape[5] self.collision_local_positions.append(local_pos) # Check if at sensor/yoke height (Z ≈ 0.08585m) - NOT bolts if abs(local_pos[2] - expected_yoke_sensor_z) < 0.001: if shape_type == p.GEOM_BOX: # Yokes are BOX type at the four corners (size 0.0254) self.yoke_indices.append(i) x_pos = "Front" if local_pos[0] > 0 else "Back" y_pos = "Left" if local_pos[1] > 0 else "Right" self.yoke_labels.append(f"{x_pos}_{y_pos}") elif shape_type == p.GEOM_CYLINDER or shape_type == p.GEOM_MESH: # Sensors: distinguish by position pattern if abs(local_pos[0]) < 0.06 or abs(local_pos[1]) < 0.02: self.sensor_indices.append(i) if abs(local_pos[0]) < 0.001: # Center sensors (X ≈ 0) label = "Center_Right" if local_pos[1] > 0 else "Center_Left" else: # Front/back sensors (Y ≈ 0) label = "Front" if local_pos[0] > 0 else "Back" self.sensor_labels.append(label) # Reset coil states self.coil_front_L.current = 0 self.coil_front_R.current = 0 self.coil_back_L.current = 0 self.coil_back_R.current = 0 # Reset velocity tracking self.prev_sensor_gaps = None return self._get_obs(), {} def step(self, action): # --- 1. Update Coil Currents from PWM Actions --- pwm_front_L = action[0] # Front yoke, +Y end pwm_front_R = action[1] # Front yoke, -Y end pwm_back_L = action[2] # Back yoke, +Y end pwm_back_R = action[3] # Back yoke, -Y end curr_front_L = self.coil_front_L.update(pwm_front_L, self.dt) curr_front_R = self.coil_front_R.update(pwm_front_R, self.dt) curr_back_L = self.coil_back_L.update(pwm_back_L, self.dt) curr_back_R = self.coil_back_R.update(pwm_back_R, self.dt) # --- 2. Get Current Pod State --- pos, orn = p.getBasePositionAndOrientation(self.podId, physicsClientId=self.client) lin_vel, ang_vel = p.getBaseVelocity(self.podId, physicsClientId=self.client) # Convert quaternion to rotation matrix rot_matrix = np.array(p.getMatrixFromQuaternion(orn)).reshape(3, 3) # --- 3. Calculate Gap Heights at Yoke Positions (for force prediction) --- # Calculate world positions of the 4 yokes (ends of U-yokes) yoke_gap_heights_dict = {} # Store by label for easy access for i, yoke_idx in enumerate(self.yoke_indices): local_pos = self.collision_local_positions[yoke_idx] local_vec = np.array(local_pos) world_offset = rot_matrix @ local_vec world_pos = np.array(pos) + world_offset # Top surface of yoke box (add half-height = 5mm) yoke_top_z = world_pos[2] + 0.005 # Gap height: track bottom (Z=0) to yoke top (negative Z) gap_height = -yoke_top_z # Convert to positive gap in meters yoke_gap_heights_dict[self.yoke_labels[i]] = gap_height # Average gap heights for each U-shaped yoke (average left and right ends) # Front yoke: average of Front_Left and Front_Right # Back yoke: average of Back_Left and Back_Right avg_gap_front = (yoke_gap_heights_dict.get('Front_Left', 0.010) + yoke_gap_heights_dict.get('Front_Right', 0.010)) / 2 avg_gap_back = (yoke_gap_heights_dict.get('Back_Left', 0.010) + yoke_gap_heights_dict.get('Back_Right', 0.010)) / 2 # Calculate roll angle from yoke end positions # Roll = arctan((right_z - left_z) / y_distance) # Average front and back measurements front_left_gap = yoke_gap_heights_dict.get('Front_Left', 0.010) front_right_gap = yoke_gap_heights_dict.get('Front_Right', 0.010) back_left_gap = yoke_gap_heights_dict.get('Back_Left', 0.010) back_right_gap = yoke_gap_heights_dict.get('Back_Right', 0.010) # Rigid body distances (hypotenuses - don't change with rotation) y_distance = 0.1016 # 2 * 0.0508m (left to right distance) x_distance = 0.2518 # 2 * 0.1259m (front to back distance) # Roll angle (rotation about X-axis) - using arcsin since y_distance is the hypotenuse # When right side has larger gap (higher), roll is negative (right-side-up convention) roll_angle_front = np.arcsin(-(front_right_gap - front_left_gap) / y_distance) roll_angle_back = np.arcsin(-(back_right_gap - back_left_gap) / y_distance) roll_angle = (roll_angle_front + roll_angle_back) / 2 # Pitch angle (rotation about Y-axis) - using arcsin since x_distance is the hypotenuse # When back has larger gap (higher), pitch is positive (nose-down convention) pitch_angle_left = np.arcsin((back_left_gap - front_left_gap) / x_distance) pitch_angle_right = np.arcsin((back_right_gap - front_right_gap) / x_distance) pitch_angle = (pitch_angle_left + pitch_angle_right) / 2 # --- 4. Predict Forces and Torques using Maglev Predictor --- # Currents can be positive or negative (as seen in CSV data) # Positive weakens permanent magnet field, negative strengthens it # Gap heights in mm gap_front_mm = avg_gap_front * 1000 gap_back_mm = avg_gap_back * 1000 # Roll angle in degrees roll_deg = np.degrees(roll_angle) # Predict force and torque for each U-shaped yoke # Front yoke (at X = +0.1259m) force_front, torque_front = self.predictor.predict( curr_front_L, curr_front_R, roll_deg, gap_front_mm ) # Back yoke (at X = -0.1259m) force_back, torque_back = self.predictor.predict( curr_back_L, curr_back_R, roll_deg, gap_back_mm ) # --- 5. Apply Forces and Torques to Pod --- # Forces are applied at Y=0 (center of U-yoke) at each X position # This is where the actual magnetic force acts on the U-shaped yoke # Apply force at front yoke center (X=+0.1259, Y=0) front_yoke_center = [0.1259, 0, 0.08585] # From pod.xml yoke positions p.applyExternalForce( self.podId, -1, forceObj=[0, 0, force_front], posObj=front_yoke_center, flags=p.LINK_FRAME, physicsClientId=self.client ) # Apply force at back yoke center (X=-0.1259, Y=0) back_yoke_center = [-0.1259, 0, 0.08585] p.applyExternalForce( self.podId, -1, forceObj=[0, 0, force_back], posObj=back_yoke_center, flags=p.LINK_FRAME, physicsClientId=self.client ) # Apply roll torques # Each yoke produces its own torque about X axis torque_front_Nm = torque_front / 1000 # Convert from mN·m to N·m torque_back_Nm = torque_back / 1000 # Apply torques at respective yoke positions p.applyExternalTorque( self.podId, -1, torqueObj=[torque_front_Nm, 0, 0], flags=p.LINK_FRAME, physicsClientId=self.client ) p.applyExternalTorque( self.podId, -1, torqueObj=[torque_back_Nm, 0, 0], flags=p.LINK_FRAME, physicsClientId=self.client ) # --- 6. Step Simulation --- p.stepSimulation(physicsClientId=self.client) # --- 7. Get New Observation --- obs = self._get_obs() # --- 8. Calculate Reward --- # Goal: Hover at target gap (10mm = 0.010m), minimize roll, minimize power target_gap = TARGET_GAP # 10mm in meters avg_gap = (avg_gap_front + avg_gap_back) / 2 gap_error = abs(avg_gap - target_gap) roll_error = abs(roll_angle) pitch_error = abs(pitch_angle) z_vel_penalty = abs(lin_vel[2]) # Power dissipation (all 4 coils) power = (curr_front_L**2 * self.coil_front_L.R + curr_front_R**2 * self.coil_front_R.R + curr_back_L**2 * self.coil_back_L.R + curr_back_R**2 * self.coil_back_R.R) # Reward function reward = 1.0 reward -= gap_error * 100 # Penalize gap error heavily reward -= roll_error * 50 # Penalize roll (about X-axis) reward -= pitch_error * 50 # Penalize pitch (about Y-axis) reward -= z_vel_penalty * 10 # Penalize vertical motion reward -= power * 0.01 # Small penalty for power usage # --- 9. Check Termination --- terminated = False truncated = False # Check if pod crashed or flew away if avg_gap < 0.002 or avg_gap > 0.030: # Outside 2-30mm range terminated = True reward -= 100 # Check if pod rolled or pitched too much if abs(roll_angle) > np.radians(10): # More than 10 degrees roll terminated = True reward -= 50 if abs(pitch_angle) > np.radians(10): # More than 10 degrees pitch terminated = True reward -= 50 info = { 'curr_front_L': curr_front_L, 'curr_front_R': curr_front_R, 'curr_back_L': curr_back_L, 'curr_back_R': curr_back_R, 'gap_front_yoke': avg_gap_front, 'gap_back_yoke': avg_gap_back, 'roll': roll_angle, 'force_front': force_front, 'force_back': force_back, 'torque_front': torque_front, 'torque_back': torque_back } return obs, reward, terminated, truncated, info def _get_obs(self): """ Returns observation: [gaps(4), velocities(4)] Uses noisy sensor readings + computed velocities for microcontroller-friendly deployment """ pos, orn = p.getBasePositionAndOrientation(self.podId, physicsClientId=self.client) # Convert quaternion to rotation matrix rot_matrix = np.array(p.getMatrixFromQuaternion(orn)).reshape(3, 3) # Calculate sensor gap heights with noise sensor_gap_heights = {} for i, sensor_idx in enumerate(self.sensor_indices): local_pos = self.collision_local_positions[sensor_idx] local_vec = np.array(local_pos) world_offset = rot_matrix @ local_vec world_pos = np.array(pos) + world_offset # Top surface of sensor (add half-height = 5mm) sensor_top_z = world_pos[2] + 0.005 # Gap height: track bottom (Z=0) to sensor top gap_height = -sensor_top_z # Add measurement noise noisy_gap = gap_height + np.random.normal(0, self.sensor_noise_std) sensor_gap_heights[self.sensor_labels[i]] = noisy_gap # Pack sensor measurements in consistent order # [center_right, center_left, front, back] gaps = np.array([ sensor_gap_heights.get('Center_Right', 0.010), sensor_gap_heights.get('Center_Left', 0.010), sensor_gap_heights.get('Front', 0.010), sensor_gap_heights.get('Back', 0.010) ], dtype=np.float32) # Compute velocities (d_gap/dt) if self.prev_sensor_gaps is None: # First observation - no velocity information yet velocities = np.zeros(4, dtype=np.float32) else: # Compute velocity as finite difference velocities = (gaps - self.prev_sensor_gaps) / self.dt # Store for next step self.prev_sensor_gaps = gaps.copy() # Concatenate: [gaps, velocities] obs = np.concatenate([gaps, velocities]) return obs def _create_modified_urdf(self): """ Create a modified URDF with bolt positions adjusted based on initial gap height. Bolts are at Z = 0.08585 + gap_mm/1000 (relative to pod origin). Yokes and sensors remain at Z = 0.08585 (relative to pod origin). """ import tempfile # Calculate bolt Z position bolt_z = 0.08585 + self.initial_gap_mm / 1000.0 # Read the original URDF template urdf_template_path = os.path.join(os.path.dirname(__file__), "pod.xml") with open(urdf_template_path, 'r') as f: urdf_content = f.read() # Replace the bolt Z positions (originally at 0.09585) # There are 4 bolts at different X,Y positions but same Z urdf_modified = urdf_content.replace( 'xyz="0.285 0.03 0.09585"', f'xyz="0.285 0.03 {bolt_z:.6f}"' ).replace( 'xyz="0.285 -0.03 0.09585"', f'xyz="0.285 -0.03 {bolt_z:.6f}"' ).replace( 'xyz="-0.285 0.03 0.09585"', f'xyz="-0.285 0.03 {bolt_z:.6f}"' ).replace( 'xyz="-0.285 -0.03 0.09585"', f'xyz="-0.285 -0.03 {bolt_z:.6f}"' ) # Write to a temporary file with tempfile.NamedTemporaryFile(mode='w', suffix='.urdf', delete=False) as f: f.write(urdf_modified) temp_urdf_path = f.name return temp_urdf_path def close(self): p.disconnect(physicsClientId=self.client) def render(self): """Rendering is handled by PyBullet GUI mode""" pass