pid trial

RL Testing/lev_pod_env.py

@@ -10,13 +10,16 @@ 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):
+    def __init__(self, use_gui=False, initial_gap_mm=10.0, max_steps=2000, disturbance_force_std=0.0):
         super(LevPodEnv, self).__init__()
-        
+
         # Store initial gap height parameter
         self.initial_gap_mm = initial_gap_mm
         self.max_episode_steps = max_steps
         self.current_step = 0
+
+        # Stochastic disturbance parameter (standard deviation of random force in Newtons)
+        self.disturbance_force_std = disturbance_force_std
         
         # The following was coded by AI - see [1]
         # --- 1. Define Action & Observation Spaces ---
@@ -95,13 +98,14 @@ class LevPodEnv(gym.Env):
         urdf_path = self._create_modified_urdf()
         
         # Determine start condition
-        if np.random.rand() > 0.5:
-            # Spawn exactly at target
-            spawn_gap_mm = TARGET_GAP * 1000.0
-            # # Add tiny noise
-            # spawn_gap_mm += np.random.uniform(-0.5, 0.5)
-        else:
-            spawn_gap_mm = self.initial_gap_mm
+        # if np.random.rand() > 0.5:
+        #     # Spawn exactly at target
+        #     spawn_gap_mm = TARGET_GAP * 1000.0
+        #     # # Add tiny noise
+        #     # spawn_gap_mm += np.random.uniform(-0.5, 0.5)
+        # else:
+        
+        spawn_gap_mm = self.initial_gap_mm
             
         start_z = -(0.09085 + spawn_gap_mm / 1000.0)
         start_pos = [0, 0, start_z]
@@ -290,7 +294,18 @@ class LevPodEnv(gym.Env):
             flags=p.LINK_FRAME,
             physicsClientId=self.client
         )
-        
+
+        # --- 5b. Apply Stochastic Disturbance Force (if enabled) ---
+        if self.disturbance_force_std > 0:
+            disturbance_force = np.random.normal(0, self.disturbance_force_std)
+            p.applyExternalForce(
+                self.podId, -1,
+                forceObj=[0, 0, disturbance_force],
+                posObj=[0, 0, 0],
+                flags=p.LINK_FRAME,
+                physicsClientId=self.client
+            )
+
         # --- 6. Step Simulation ---
         p.stepSimulation(physicsClientId=self.client)
         self.current_step += 1
@@ -359,6 +374,7 @@ class LevPodEnv(gym.Env):
         if gap_error_mm < 1.0 and abs(np.degrees(roll_angle)) < 2.0 and abs(np.degrees(pitch_angle)) < 2.0:
             reward += 2.0  # Bonus for excellent control
         
+        # Ground truth info dict (from PyBullet physical state, not sensor observations)
         info = {
             'curr_front_L': curr_front_L,
             'curr_front_R': curr_front_R,
@@ -366,7 +382,9 @@ class LevPodEnv(gym.Env):
             'curr_back_R': curr_back_R,
             'gap_front_yoke': avg_gap_front,
             'gap_back_yoke': avg_gap_back,
+            'gap_avg': avg_gap,
             'roll': roll_angle,
+            'pitch': pitch_angle,
             'force_front': force_front,
             'force_back': force_back,
             'torque_front': torque_front,
@@ -375,6 +393,24 @@ class LevPodEnv(gym.Env):
         
         return obs, reward, terminated, truncated, info
 
+    def apply_impulse(self, force_z: float, position: list = None):
+        """
+        Apply a one-time impulse force to the pod.
+
+        Args:
+            force_z: Vertical force in Newtons (positive = upward)
+            position: Local position [x, y, z] to apply force (default: center of mass)
+        """
+        if position is None:
+            position = [0, 0, 0]
+        p.applyExternalForce(
+            self.podId, -1,
+            forceObj=[0, 0, force_z],
+            posObj=position,
+            flags=p.LINK_FRAME,
+            physicsClientId=self.client
+        )
+
     # The following was generated by AI - see [15]
     def _get_obs(self, initial_reset=False):
         """