perhaps optimal
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user