Mujoco_WASM/web_platform/src/mobile/agent/DeterministicBaseline.ts

import type { SimulationSession } from '../../simulation/SimulationSession';
import { ModelBindings } from '../../simulation/ModelBindings';
import { LeKiwiIK } from './LeKiwiIK';
import { PhysicalSampler, PickPlaceEvaluator, type PhysicalSample } from './PhysicalEvidence';
import { PICK_PLACE as P, PICK_PLACE_ROBOT as R, seedPositions } from './PickPlaceScene';

const phases = [
  'open',
  'pregrasp',
  'descend',
  'close',
  'verify',
  'carry',
  'stop-base',
  'place',
  'release',
  'retreat',
  'settle',
] as const;
export type BaselinePhase = (typeof phases)[number];
const durations = [2, 6, 3, 3, 4, 22, 1, 4, 3, 4, 1];
const clamp = (v: number, lo: number, hi: number) => Math.max(lo, Math.min(hi, v));

/** No network, RL reward, object attachment or independent physics loop.
 * The caller exclusively owns the existing simulation clock and validates writes.
 * Ungated mode is the frozen Stage-5 reference. Gated mode waits at each verified
 * skill boundary; only a current synchronous owner callback may authorize/recover. */
export class DeterministicBaseline {
  readonly evaluator: PickPlaceEvaluator;
  readonly trace: (PhysicalSample & {
    phase: BaselinePhase;
    secure: boolean;
    transported: number;
  })[] = [];
  readonly ik: LeKiwiIK;
  private readonly sampler: PhysicalSampler;
  private readonly arm;
  private readonly gripper;
  private readonly wheels;
  private readonly positions;
  private graspSource: number[];
  private readonly baseStart: number[];
  private readonly servo: number[];
  private readonly velocities = Array<number>(5).fill(0);
  private readonly baseVelocity = [0, 0, 0];
  private readonly translation: readonly number[];
  private opening: number;
  private phaseIndex = 0;
  private phaseStart: number;
  private readonly start: number;
  private nextControl: number;
  private lostGraspTime = 0;
  private readonly armStallTime = Array<number>(5).fill(0);
  private disposed = false;
  private completedPhase = false;
  waiting = false;
  sample: PhysicalSample;
  status: 'running' | 'succeeded' | 'failed' = 'running';
  error = '';
  maxJointSpeed = 0;
  constructor(
    private readonly session: SimulationSession,
    seed: number,
    private readonly writeServo: (actuator: number, value: number) => void,
    private readonly gated = false,
    goal?: readonly number[],
  ) {
    const b = new ModelBindings(session.model);
    this.positions = seedPositions(seed, goal);
    this.translation = goal
      ? this.positions.goal.slice(0, 2).map((v, i) => v - this.positions.object[i])
      : P.baseTranslation;
    this.graspSource = [...this.positions.object];
    this.waiting = gated;
    this.arm = R.armJoints.map((spec, i) =>
      b.scalarActuator(R.armActuators[i], spec.name, 'position'),
    );
    this.gripper = b.scalarActuator(R.gripperActuators[0].name, R.gripperJoint, 'position');
    this.wheels = R.baseActuators.map((name, i) =>
      b.scalarActuator(name, R.baseJoints[i], 'velocity'),
    );
    this.sampler = new PhysicalSampler(session);
    let ik: LeKiwiIK | undefined;
    try {
      this.ik = ik = new LeKiwiIK(session);
      this.sample = this.sampler.sample();
    } catch (error) {
      ik?.dispose();
      this.sampler.dispose();
      throw error;
    }
    this.baseStart = [...this.sample.base];
    this.servo = [...this.sample.joints];
    this.opening = this.sample.opening;
    this.start = this.phaseStart = this.nextControl = this.sample.time;
    this.evaluator = new PickPlaceEvaluator(this.positions.object, this.positions.goal);
  }
  get phase(): BaselinePhase {
    return phases[this.phaseIndex];
  }
  get goal(): readonly number[] {
    return this.positions.goal;
  }
  get candidate(): BaselinePhase {
    return phases[this.phaseIndex + (this.completedPhase ? 1 : 0)] ?? 'settle';
  }
  /** Re-check facts even if both models recommend continuation. No model can grant this gate. */
  authorize(skill: BaselinePhase): void {
    if (!this.gated || !this.waiting || this.status !== 'running' || skill !== this.candidate)
      throw new Error('非法技能切换');
    const s = this.sample;
    if (skill === 'pregrasp' && s.baseSpeed > 0.01) throw new Error('base_not_stopped');
    if (skill === 'descend' && s.tcp[2] < s.object[2] + 0.04) throw new Error('tcp_not_above');
    if (skill === 'close' && Math.hypot(...s.tcp.map((v, i) => v - s.object[i])) > 0.012)
      throw new Error('alignment_timeout');
    if (skill === 'verify' && s.fingerForces.some((v) => v < P.fingerForceMin))
      throw new Error('empty_grasp');
    if (['carry', 'stop-base'].includes(skill) && !this.evaluator.secure)
      throw new Error('grasp_unverified');
    if (skill === 'place' && (this.evaluator.transported < P.transportMin || s.baseSpeed > 0.01))
      throw new Error('transport_unverified');
    if (skill === 'release' && !s.onGoalSupport) throw new Error('placement_unsupported');
    if (
      skill === 'retreat' &&
      (s.opening < 0.65 || s.fingerForces.some((v) => v >= P.fingerForceMin))
    )
      throw new Error('release_blocked');
    if (this.completedPhase) this.phaseIndex++;
    this.completedPhase = false;
    this.waiting = false;
    this.phaseStart = this.nextControl = s.time;
  }
  canRecover(): boolean {
    const s = this.sample;
    return (
      this.gated &&
      this.status === 'failed' &&
      this.phaseIndex <= phases.indexOf('verify') &&
      ['alignment_timeout', 'empty_grasp', 'grasp_unverified'].includes(this.error) &&
      s.supported &&
      s.objectSpeed < P.placementSpeedMax &&
      s.baseSpeed < 0.01 &&
      this.evaluator.transported === 0 &&
      Math.hypot(...s.object.map((v, i) => v - this.positions.object[i])) < 0.012
    );
  }
  /** Only a supported, stationary object still at the source may be realigned/regrasped. */
  recover(): void {
    if (!this.canRecover()) throw new Error('unsafe_recovery');
    this.hold();
    // An empty/blocked attempt may not have reached its requested gripper position.
    // Re-seed targets from measured pose so a released fault cannot cause a servo jump.
    this.opening = this.sample.opening;
    this.servo.splice(0, this.servo.length, ...this.sample.joints);
    this.armStallTime.fill(0);
    this.graspSource = [...this.sample.object];
    this.phaseIndex = 0;
    this.phaseStart = this.nextControl = this.sample.time;
    this.completedPhase = false;
    this.waiting = true;
    this.status = 'running';
    this.error = '';
    // Do not reset the episode deadline or fabricate/reset the evaluator's transport history.
  }
  private target(): number[] | undefined {
    const source = this.graspSource,
      goal = this.positions.goal;
    switch (this.phase) {
      case 'pregrasp':
        return [source[0], source[1], P.pregraspHeight];
      case 'descend':
      case 'close':
        return [...source];
      case 'verify':
        return [source[0], source[1], P.carryHeight];
      case 'place':
      case 'release':
        return [...goal];
      case 'retreat':
      case 'settle':
        return [goal[0], goal[1], P.retreatHeight];
      default:
        return undefined;
    }
  }
  beforeStep(): void {
    if (
      this.disposed ||
      this.waiting ||
      this.status !== 'running' ||
      this.sample.time + 1e-9 < this.nextControl
    )
      return;
    this.nextControl = this.sample.time + P.controlDt;
    try {
      const target = this.target();
      if (target) {
        const solution = this.ik.solve(target);
        if (solution.residual > 0.004) throw new Error('unreachable');
        this.arm.forEach((joint, i) => {
          const desired = clamp(5 * (solution.targets[i] - this.servo[i]), -P.armSpeed, P.armSpeed);
          this.velocities[i] += clamp(
            desired - this.velocities[i],
            -P.armAcceleration * P.controlDt,
            P.armAcceleration * P.controlDt,
          );
          const old = this.servo[i],
            spec = R.armJoints[i];
          const next = clamp(
            clamp(
              old + this.velocities[i] * P.controlDt,
              this.sample.joints[i] - P.armTrackingError,
              this.sample.joints[i] + P.armTrackingError,
            ),
            old - P.armSpeed * P.controlDt,
            old + P.armSpeed * P.controlDt,
          );
          this.servo[i] = clamp(next, spec.min, spec.max);
          this.velocities[i] = (this.servo[i] - old) / P.controlDt;
        });
      }
      this.arm.forEach((joint, i) => this.writeServo(joint.actuatorId, this.servo[i]));
      const close = ['close', 'verify', 'carry', 'stop-base', 'place'].includes(this.phase);
      this.opening += clamp(
        (close ? 0 : 1) - this.opening,
        -P.gripperOpeningRate * P.controlDt,
        P.gripperOpeningRate * P.controlDt,
      );
      this.writeServo(
        this.gripper.actuatorId,
        R.gripperClosed + this.opening * (R.gripperOpen - R.gripperClosed),
      );
      const desired = [0, 0, 0];
      if (this.phase === 'carry') {
        const dx = this.baseStart[0] + this.translation[0] - this.sample.base[0];
        const dy = this.baseStart[1] + this.translation[1] - this.sample.base[1];
        const scale = Math.min(1, P.baseSpeed / Math.max(1e-12, Math.hypot(dx, dy)));
        const c = Math.cos(this.sample.yaw),
          s = Math.sin(this.sample.yaw);
        desired[0] = scale * (c * dx + s * dy);
        desired[1] = scale * (-s * dx + c * dy);
        desired[2] = clamp(-1.5 * this.sample.yaw, -P.yawSpeed, P.yawSpeed);
      }
      for (let i = 0; i < 3; i++) {
        const bound = (i === 2 ? P.yawAcceleration : P.baseAcceleration) * P.controlDt;
        this.baseVelocity[i] += clamp(desired[i] - this.baseVelocity[i], -bound, bound);
      }
      const wheelSpeeds = R.baseMix.map((row) =>
        row.reduce((sum, v, i) => sum + v * this.baseVelocity[i], 0),
      );
      const wheelScale = R.wheelLimit / Math.max(R.wheelLimit, ...wheelSpeeds.map(Math.abs));
      this.wheels.forEach((wheel, i) =>
        this.writeServo(wheel.actuatorId, wheelSpeeds[i] * wheelScale),
      );
      this.trace.push({
        ...this.sample,
        phase: this.phase,
        secure: this.evaluator.secure,
        transported: this.evaluator.transported,
      });
    } catch (error) {
      this.fail(error instanceof Error ? error.message : String(error));
    }
  }
  afterStep(dt: number): void {
    if (this.disposed || this.status !== 'running') return;
    try {
      const previous = this.sample;
      this.sample = this.sampler.sample();
      this.evaluator.update(this.sample, dt);
      // Only arm joints: gripper contact intentionally blocks a closing servo.
      // Frozen model waits do not call afterStep and cannot accumulate stall time.
      for (let i = 0; i < this.arm.length; i++) {
        const speed = Math.abs(this.sample.joints[i] - previous.joints[i]) / dt;
        const blocked =
          Math.abs(this.servo[i] - this.sample.joints[i]) >= P.armStallError &&
          speed < P.armStallSpeed;
        this.armStallTime[i] = blocked ? this.armStallTime[i] + dt : 0;
        if (this.armStallTime[i] >= P.armStallDuration) throw new Error('joint_stall');
      }
      this.maxJointSpeed = Math.max(this.maxJointSpeed, this.sample.peakJointSpeed);
      if (this.sample.peakJointSpeed > P.jointSpeedStop) throw new Error('joint_velocity');
      if (this.sample.base[2] < 0.03 || this.sample.base[2] > 0.1 || this.sample.tilt > 0.35)
        throw new Error('base_unstable');
      if (this.sample.time - this.start > P.episodeTimeout) throw new Error('episode_timeout');
      if (this.phase === 'carry' || this.phase === 'stop-base') {
        this.lostGraspTime = this.evaluator.secure ? 0 : this.lostGraspTime + dt;
        if (this.lostGraspTime > 0.2) throw new Error('lost_grasp');
      }
      let ready = this.sample.time - this.phaseStart >= durations[this.phaseIndex];
      if (this.phase === 'carry') {
        const distance = Math.hypot(
          ...this.sample.base
            .slice(0, 2)
            .map((v, i) => v - this.baseStart[i] - this.translation[i]),
        );
        if (distance < 0.007) ready = true;
        else if (ready) throw new Error('navigation_timeout');
      }
      if (!ready) return;
      const target = this.target();
      if (target && Math.hypot(...target.map((v, i) => v - this.sample.tcp[i])) > 0.012)
        throw new Error('alignment_timeout');
      if (this.phase === 'close' && this.sample.fingerForces.some((f) => f < P.fingerForceMin))
        throw new Error('empty_grasp');
      if (this.phase === 'verify' && !this.evaluator.secure) throw new Error('grasp_unverified');
      if (this.phase === 'carry' && this.evaluator.transported < P.transportMin)
        throw new Error('transport_unverified');
      if (this.phase === 'place' && !this.sample.onGoalSupport)
        throw new Error('placement_unsupported');
      if (
        this.phase === 'release' &&
        (this.sample.opening < 0.65 || this.sample.fingerForces.some((f) => f >= P.fingerForceMin))
      )
        throw new Error('release_blocked');
      if (this.phase === 'settle') {
        if (!this.evaluator.succeeded) throw new Error('placement_unverified');
        this.status = 'succeeded';
        this.hold();
      } else if (this.gated) {
        this.completedPhase = true;
        this.waiting = true;
        this.hold();
      } else {
        this.phaseIndex++;
        this.phaseStart = this.sample.time;
      }
    } catch (error) {
      this.fail(error instanceof Error ? error.message : String(error));
    }
  }
  /** Called only through the session's current owner write capability. */
  hold(): void {
    this.baseVelocity.fill(0);
    this.velocities.fill(0);
    this.wheels.forEach((wheel) => this.writeServo(wheel.actuatorId, 0));
    // Keep arm/gripper's LAST safe servo targets; never zero loaded position servos.
  }
  private fail(reason: string): void {
    this.status = 'failed';
    this.error = reason;
    this.hold();
  }
  dispose(): void {
    if (this.disposed) return;
    this.disposed = true;
    this.sampler.dispose();
    this.ik.dispose();
  }
}