Files
guadaloop_lev_control/RL Testing/lev_pod_env.py

446 lines
19 KiB
Python
Raw Normal View History

2025-12-10 15:50:20 -06:00
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