From b4739362a5f6dc21758c62c8297c8edf950c8bd2 Mon Sep 17 00:00:00 2001 From: Brunsmeier <2970937094@qq.com> Date: Tue, 7 Jul 2026 17:08:04 +0800 Subject: [PATCH] Refactor IMU velocity estimator to improve gravity compensation and integrate world-frame velocity calculations --- record_episode.py | 147 +++++++++++++++++++++++++++++----------------- 1 file changed, 92 insertions(+), 55 deletions(-) diff --git a/record_episode.py b/record_episode.py index d47e9c9..d51deb6 100644 --- a/record_episode.py +++ b/record_episode.py @@ -239,27 +239,23 @@ class D405Camera: # ============================================================ -# Simple IMU velocity estimator +# Simple IMU velocity estimator Updated # ============================================================ -class ProviderIMUVelocityEstimator: +class ProviderWorldIMUVelocityEstimator: """ - Provider-style roll/pitch gravity compensation + velocity integration. + Provider-style gravity/support compensation + world-frame velocity integration. - This follows verify.py: - gravity_body_g = [ - -sin(pitch), - cos(pitch) * sin(roll), - cos(pitch) * cos(roll), - ] * gravity_sign + Pipeline: + 1. raw acc_body_g - gravity_body_g -> linear_acc_body_g + 2. estimate residual body-frame acceleration bias + 3. linear_acc_body_ms2 -> linear_acc_world_ms2 using R_body_to_world + 4. integrate velocity in world frame - Then: - linear_acc_body = raw_acc_body - gravity_body - - Important: - This is still an estimated IMU velocity, not ground truth. - It removes the static gravity/support acceleration approximately, - but velocity integration will still drift. + Output: + velocity_ms: world-frame velocity, unit m/s + linear_acc_ms2: world-frame linear acceleration, unit m/s^2 + acc_bias_ms2: body-frame residual acceleration bias, unit m/s^2 """ def __init__( @@ -267,21 +263,20 @@ class ProviderIMUVelocityEstimator: gravity_sign=1.0, bias_seconds=1.0, acc_deadband=0.10, - vel_decay=0.995, ): self.gravity_sign = gravity_sign self.bias_seconds = bias_seconds self.acc_deadband = acc_deadband - self.vel_decay = vel_decay self.start_t_ms = None self.last_t_ms = None self.bias_ready = False self.bias_samples = [] - self.acc_bias = np.zeros(3, dtype=float) + self.acc_bias_body_ms2 = np.zeros(3, dtype=float) - self.v = np.zeros(3, dtype=float) + self.v_world = np.zeros(3, dtype=float) + self.last_acc_world_ms2 = np.zeros(3, dtype=float) def reset(self): self.start_t_ms = None @@ -289,9 +284,10 @@ class ProviderIMUVelocityEstimator: self.bias_ready = False self.bias_samples = [] - self.acc_bias[:] = 0.0 + self.acc_bias_body_ms2[:] = 0.0 - self.v[:] = 0.0 + self.v_world[:] = 0.0 + self.last_acc_world_ms2[:] = 0.0 def gravity_body_g(self, roll_deg, pitch_deg): roll = math.radians(roll_deg) @@ -303,6 +299,22 @@ class ProviderIMUVelocityEstimator: return np.array([g_x, g_y, g_z], dtype=float) * self.gravity_sign + def rot_body_to_world(self, roll_deg, pitch_deg, yaw_deg): + roll = math.radians(roll_deg) + pitch = math.radians(pitch_deg) + yaw = math.radians(yaw_deg) + + cr, sr = math.cos(roll), math.sin(roll) + cp, sp = math.cos(pitch), math.sin(pitch) + cy, sy = math.cos(yaw), math.sin(yaw) + + # R_body_to_world = Rz(yaw) @ Ry(pitch) @ Rx(roll) + return np.array([ + [cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr], + [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr], + [-sp, cp * sr, cp * cr], + ], dtype=float) + def update(self, packet): t_ms = packet["t_ms"] @@ -313,61 +325,80 @@ class ProviderIMUVelocityEstimator: roll_deg = packet["angle_deg"][0] pitch_deg = packet["angle_deg"][1] + yaw_deg = packet["angle_deg"][2] + # ---------------- gravity/support removal in body frame ---------------- gravity_g = self.gravity_body_g( roll_deg=roll_deg, pitch_deg=pitch_deg, ) - linear_acc_g_raw = acc_body_g - gravity_g - linear_acc_ms2_raw = linear_acc_g_raw * 9.81 + linear_acc_body_g_raw = acc_body_g - gravity_g + linear_acc_body_ms2_raw = linear_acc_body_g_raw * 9.80665 elapsed_s = ((t_ms - self.start_t_ms) & 0xFFFFFFFF) / 1000.0 - # First N seconds are used as residual bias estimation. - # During this period velocity stays zero. + # ---------------- bias estimation in body frame ---------------- + # During this period, velocity stays zero. if not self.bias_ready: - self.bias_samples.append(linear_acc_ms2_raw) + self.bias_samples.append(linear_acc_body_ms2_raw.copy()) if elapsed_s >= self.bias_seconds: if self.bias_samples: - self.acc_bias = np.mean( + self.acc_bias_body_ms2 = np.mean( np.array(self.bias_samples, dtype=float), axis=0, ) self.bias_ready = True - print(f"[IMU] Provider acc bias estimated: {self.acc_bias}") + print( + "[IMU] Provider body-frame acc bias estimated: " + f"{self.acc_bias_body_ms2}" + ) self.last_t_ms = t_ms - - linear_acc_ms2 = linear_acc_ms2_raw - self.acc_bias + self.last_acc_world_ms2[:] = 0.0 return { "gravity_g": gravity_g.copy(), - "linear_acc_ms2": linear_acc_ms2.copy(), - "velocity_ms": self.v.copy(), - "acc_bias_ms2": self.acc_bias.copy(), + "linear_acc_ms2": self.last_acc_world_ms2.copy(), + "velocity_ms": self.v_world.copy(), + "acc_bias_ms2": self.acc_bias_body_ms2.copy(), "bias_ready": self.bias_ready, } + # ---------------- remove body-frame residual bias ---------------- + linear_acc_body_ms2 = linear_acc_body_ms2_raw - self.acc_bias_body_ms2 + + # ---------------- deadband before rotation ---------------- + # Norm is unchanged by pure rotation, so before/after rotation both okay. + if np.linalg.norm(linear_acc_body_ms2) < self.acc_deadband: + linear_acc_body_ms2[:] = 0.0 + + # ---------------- body frame -> world frame ---------------- + R_bw = self.rot_body_to_world( + roll_deg=roll_deg, + pitch_deg=pitch_deg, + yaw_deg=yaw_deg, + ) + + linear_acc_world_ms2 = R_bw @ linear_acc_body_ms2 + + # ---------------- dt ---------------- delta_ms = (t_ms - self.last_t_ms) & 0xFFFFFFFF self.last_t_ms = t_ms - linear_acc_ms2 = linear_acc_ms2_raw - self.acc_bias - - if np.linalg.norm(linear_acc_ms2) < self.acc_deadband: - linear_acc_ms2[:] = 0.0 - if 0 < delta_ms <= 200: dt = delta_ms / 1000.0 - self.v = self.v * self.vel_decay + linear_acc_ms2 * dt + self.v_world += linear_acc_world_ms2 * dt + + self.last_acc_world_ms2 = linear_acc_world_ms2.copy() return { "gravity_g": gravity_g.copy(), - "linear_acc_ms2": linear_acc_ms2.copy(), - "velocity_ms": self.v.copy(), - "acc_bias_ms2": self.acc_bias.copy(), + "linear_acc_ms2": self.last_acc_world_ms2.copy(), + "velocity_ms": self.v_world.copy(), + "acc_bias_ms2": self.acc_bias_body_ms2.copy(), "bias_ready": self.bias_ready, } @@ -407,6 +438,7 @@ def infer_scissor_state(packet): return "invalid_both_active" + # ============================================================ # Episode writer # ============================================================ @@ -647,12 +679,17 @@ def run(args): active_value=args.record_active_value, debounce_s=args.debounce, ) - velocity_estimator = ProviderIMUVelocityEstimator( - gravity_sign=args.gravity_sign, - bias_seconds=args.bias_seconds, - acc_deadband=args.acc_deadband, - vel_decay=args.velocity_decay, -) +# velocity_estimator = ProviderIMUVelocityEstimator( +# gravity_sign=args.gravity_sign, +# bias_seconds=args.bias_seconds, +# acc_deadband=args.acc_deadband, +# vel_decay=args.velocity_decay, +# ) + velocity_estimator = ProviderWorldIMUVelocityEstimator( + gravity_sign=args.gravity_sign, + bias_seconds=args.bias_seconds, + acc_deadband=args.acc_deadband, + ) writer = EpisodeWriter( root_dir=args.dataset, @@ -888,12 +925,12 @@ def parse_args(): help="JPEG quality for RGB image saving.", ) - parser.add_argument( - "--velocity-decay", - type=float, - default=0.995, - help="Decay factor for simple raw IMU velocity integration.", - ) + # parser.add_argument( + # "--velocity-decay", + # type=float, + # default=0.995, + # help="Decay factor for simple raw IMU velocity integration.", + # ) parser.add_argument( "--flush-every",