perhaps optimal

This commit is contained in:
2026-02-21 23:53:53 -06:00
parent 799ee6c956
commit 725f227943
19 changed files with 1925 additions and 951 deletions

View File

@@ -4,13 +4,15 @@ import pybullet as p
import pybullet_data
import numpy as np
import os
from datetime import datetime
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, max_steps=2000, disturbance_force_std=0.0):
def __init__(self, use_gui=False, initial_gap_mm=10.0, max_steps=2000, disturbance_force_std=0.0,
record_video=False, record_telemetry=False, record_dir="recordings"):
super(LevPodEnv, self).__init__()
# Store initial gap height parameter
@@ -20,6 +22,17 @@ class LevPodEnv(gym.Env):
# Stochastic disturbance parameter (standard deviation of random force in Newtons)
self.disturbance_force_std = disturbance_force_std
# Recording parameters
self.record_video = record_video
self.record_telemetry = record_telemetry
self.record_dir = record_dir
self._frame_skip = 4 # Capture every 4th step → 60fps video from 240Hz sim
self._video_width = 640
self._video_height = 480
self._frames = []
self._telemetry = {}
self._recording_active = False
# The following was coded by AI - see [1]
# --- 1. Define Action & Observation Spaces ---
@@ -143,7 +156,8 @@ class LevPodEnv(gym.Env):
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"
# +Y = Left, -Y = Right (consistent with yoke labeling)
label = "Center_Left" if local_pos[1] > 0 else "Center_Right"
else: # Front/back sensors (Y ≈ 0)
label = "Front" if local_pos[0] > 0 else "Back"
self.sensor_labels.append(label)
@@ -156,7 +170,49 @@ class LevPodEnv(gym.Env):
self.prev_sensor_gaps = None
obs = self._get_obs(initial_reset=True)
self.current_step = 0
# --- Recording setup ---
# Finalize any previous episode's recording before starting new one
if self._recording_active:
self._finalize_recording()
if self.record_video or self.record_telemetry:
self._recording_active = True
os.makedirs(self.record_dir, exist_ok=True)
if self.record_video:
self._frames = []
# Pod body center ≈ start_z, yoke tops at ≈ start_z + 0.091
# Track bottom at Z=0. Focus camera on the pod body, looking from the side
# so both the pod and the track bottom (with gap between) are visible.
pod_center_z = start_z + 0.045 # Approximate visual center of pod
self._view_matrix = p.computeViewMatrixFromYawPitchRoll(
cameraTargetPosition=[0, 0, pod_center_z],
distance=0.55,
yaw=50, # Front-right angle
pitch=-5, # Nearly horizontal to see gap from the side
roll=0,
upAxisIndex=2,
physicsClientId=self.client
)
self._proj_matrix = p.computeProjectionMatrixFOV(
fov=35, # Narrow FOV for less distortion at close range
aspect=self._video_width / self._video_height,
nearVal=0.01,
farVal=2.0,
physicsClientId=self.client
)
if self.record_telemetry:
self._telemetry = {
'time': [], 'gap_FL': [], 'gap_FR': [], 'gap_BL': [], 'gap_BR': [],
'gap_front_avg': [], 'gap_back_avg': [], 'gap_avg': [],
'roll_deg': [], 'pitch_deg': [],
'curr_FL': [], 'curr_FR': [], 'curr_BL': [], 'curr_BR': [],
'force_front': [], 'force_back': [],
'torque_front': [], 'torque_back': [],
}
return obs, {}
# The following was generated by AI - see [14]
@@ -295,7 +351,7 @@ class LevPodEnv(gym.Env):
physicsClientId=self.client
)
# --- 5b. Apply Stochastic Disturbance Force (if enabled) ---
# --- 5b. Apply Stochastic Disturbance Force and Torques (if enabled) ---
if self.disturbance_force_std > 0:
disturbance_force = np.random.normal(0, self.disturbance_force_std)
p.applyExternalForce(
@@ -305,6 +361,17 @@ class LevPodEnv(gym.Env):
flags=p.LINK_FRAME,
physicsClientId=self.client
)
# Roll and pitch disturbance torques, scaled from heave force (torque ~ force * moment_arm)
# Moment arm ~ 0.15 m so e.g. 1 N force -> 0.15 N·m torque std
disturbance_torque_std = self.disturbance_force_std * 0.15
roll_torque = np.random.normal(0, disturbance_torque_std)
pitch_torque = np.random.normal(0, disturbance_torque_std)
p.applyExternalTorque(
self.podId, -1,
torqueObj=[roll_torque, pitch_torque, 0],
flags=p.LINK_FRAME,
physicsClientId=self.client
)
# --- 6. Step Simulation ---
p.stepSimulation(physicsClientId=self.client)
@@ -390,7 +457,43 @@ class LevPodEnv(gym.Env):
'torque_front': torque_front,
'torque_back': torque_back
}
# --- Recording ---
if self.record_video and self.current_step % self._frame_skip == 0:
_, _, rgb, _, _ = p.getCameraImage(
width=self._video_width,
height=self._video_height,
viewMatrix=self._view_matrix,
projectionMatrix=self._proj_matrix,
physicsClientId=self.client
)
self._frames.append(np.array(rgb, dtype=np.uint8).reshape(
self._video_height, self._video_width, 4)[:, :, :3]) # RGBA → RGB
if self.record_telemetry:
t = self._telemetry
t['time'].append(self.current_step * self.dt)
t['gap_FL'].append(front_left_gap * 1000)
t['gap_FR'].append(front_right_gap * 1000)
t['gap_BL'].append(back_left_gap * 1000)
t['gap_BR'].append(back_right_gap * 1000)
t['gap_front_avg'].append(avg_gap_front * 1000)
t['gap_back_avg'].append(avg_gap_back * 1000)
t['gap_avg'].append(avg_gap * 1000)
t['roll_deg'].append(np.degrees(roll_angle))
t['pitch_deg'].append(np.degrees(pitch_angle))
t['curr_FL'].append(curr_front_L)
t['curr_FR'].append(curr_front_R)
t['curr_BL'].append(curr_back_L)
t['curr_BR'].append(curr_back_R)
t['force_front'].append(force_front)
t['force_back'].append(force_back)
t['torque_front'].append(torque_front)
t['torque_back'].append(torque_back)
if terminated or truncated:
self._finalize_recording()
return obs, reward, terminated, truncated, info
def apply_impulse(self, force_z: float, position: list = None):
@@ -411,6 +514,20 @@ class LevPodEnv(gym.Env):
physicsClientId=self.client
)
def apply_torque_impulse(self, torque_nm: list):
"""
Apply a one-time impulse torque to the pod (body frame).
Args:
torque_nm: [Tx, Ty, Tz] in N·m (LINK_FRAME: X=roll, Y=pitch, Z=yaw)
"""
p.applyExternalTorque(
self.podId, -1,
torqueObj=torque_nm,
flags=p.LINK_FRAME,
physicsClientId=self.client
)
# The following was generated by AI - see [15]
def _get_obs(self, initial_reset=False):
"""
@@ -511,12 +628,103 @@ class LevPodEnv(gym.Env):
return temp_urdf_path
def _finalize_recording(self):
"""Save recorded video and/or telemetry plot to disk."""
if not self._recording_active:
return
self._recording_active = False
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
# --- Save video ---
if self.record_video and len(self._frames) > 0:
import imageio.v3 as iio
video_path = os.path.join(self.record_dir, f"sim_{timestamp}.mp4")
fps = int(round(1.0 / (self.dt * self._frame_skip))) # 60fps
iio.imwrite(video_path, self._frames, fps=fps, codec="h264")
print(f"Video saved: {video_path} ({len(self._frames)} frames, {fps}fps)")
self._frames = []
# --- Save telemetry plot ---
if self.record_telemetry and len(self._telemetry.get('time', [])) > 0:
self._save_telemetry_plot(timestamp)
def _save_telemetry_plot(self, timestamp):
"""Generate and save a 4-panel telemetry figure."""
import matplotlib.pyplot as plt
t = {k: np.array(v) for k, v in self._telemetry.items()}
time = t['time']
target_mm = TARGET_GAP * 1000
weight = 5.8 * 9.81 # Pod weight in N
fig, axes = plt.subplots(4, 1, figsize=(12, 14), sharex=True)
fig.suptitle(f'Simulation Telemetry | gap₀={self.initial_gap_mm}mm target={target_mm:.1f}mm',
fontsize=14, fontweight='bold')
# --- Panel 1: Gap heights ---
ax = axes[0]
ax.plot(time, t['gap_FL'], lw=1, alpha=0.6, label='FL')
ax.plot(time, t['gap_FR'], lw=1, alpha=0.6, label='FR')
ax.plot(time, t['gap_BL'], lw=1, alpha=0.6, label='BL')
ax.plot(time, t['gap_BR'], lw=1, alpha=0.6, label='BR')
ax.plot(time, t['gap_avg'], lw=2, color='black', label='Average')
ax.axhline(y=target_mm, color='orange', ls='--', lw=2, label=f'Target ({target_mm:.1f}mm)')
ax.set_ylabel('Gap Height (mm)')
ax.legend(loc='best', ncol=3, fontsize=9)
ax.grid(True, alpha=0.3)
final_err = abs(t['gap_avg'][-1] - target_mm)
ax.text(0.98, 0.02, f'Final error: {final_err:.3f}mm',
transform=ax.transAxes, ha='right', va='bottom', fontsize=10,
bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
# --- Panel 2: Roll & Pitch ---
ax = axes[1]
ax.plot(time, t['roll_deg'], lw=1.5, label='Roll')
ax.plot(time, t['pitch_deg'], lw=1.5, label='Pitch')
ax.axhline(y=0, color='gray', ls='--', lw=1)
ax.set_ylabel('Angle (degrees)')
ax.legend(loc='best', fontsize=9)
ax.grid(True, alpha=0.3)
# --- Panel 3: Coil currents ---
ax = axes[2]
ax.plot(time, t['curr_FL'], lw=1, alpha=0.8, label='FL')
ax.plot(time, t['curr_FR'], lw=1, alpha=0.8, label='FR')
ax.plot(time, t['curr_BL'], lw=1, alpha=0.8, label='BL')
ax.plot(time, t['curr_BR'], lw=1, alpha=0.8, label='BR')
total = np.abs(t['curr_FL']) + np.abs(t['curr_FR']) + np.abs(t['curr_BL']) + np.abs(t['curr_BR'])
ax.plot(time, total, lw=2, color='black', ls='--', label='Total |I|')
ax.set_ylabel('Current (A)')
ax.legend(loc='best', ncol=3, fontsize=9)
ax.grid(True, alpha=0.3)
# --- Panel 4: Forces ---
ax = axes[3]
ax.plot(time, t['force_front'], lw=1.5, label='Front yoke')
ax.plot(time, t['force_back'], lw=1.5, label='Back yoke')
f_total = t['force_front'] + t['force_back']
ax.plot(time, f_total, lw=2, color='black', label='Total')
ax.axhline(y=weight, color='red', ls='--', lw=1.5, label=f'Weight ({weight:.1f}N)')
ax.set_ylabel('Force (N)')
ax.set_xlabel('Time (s)')
ax.legend(loc='best', ncol=2, fontsize=9)
ax.grid(True, alpha=0.3)
plt.tight_layout()
plot_path = os.path.join(self.record_dir, f"sim_{timestamp}_telemetry.png")
fig.savefig(plot_path, dpi=200, bbox_inches='tight')
plt.close(fig)
print(f"Telemetry plot saved: {plot_path}")
self._telemetry = {}
def close(self):
self._finalize_recording()
try:
p.disconnect(physicsClientId=self.client)
except p.error:
pass # Already disconnected
def render(self):
"""Rendering is handled by PyBullet GUI mode"""
pass