Mujoco_WASM/web_platform/physics/agent.ts

import { MainThreadPhysicsAdapter } from '../src/simulation/PhysicsAdapter';
import { importBrowserFiles } from '../src/project/importer';
import type { ProjectManifest } from '../src/project/types';
import { composePickPlaceScene, PICK_PLACE_ROBOT } from '../src/mobile/agent/PickPlaceScene';
import { DeterministicBaseline } from '../src/mobile/agent/DeterministicBaseline';
import { AgentTaskController } from '../src/mobile/agent/AgentTaskController';
import { MockDecisionProvider } from '../src/mobile/agent/MockDecisionProvider';
import { DecisionClient } from '../src/mobile/agent/DecisionClient';
import { ModelBindings } from '../src/simulation/ModelBindings';
import { PICK_PLACE as P } from '../src/mobile/agent/PickPlaceScene';
type Fault =
  | 'empty-grasp'
  | 'empty-once'
  | 'alignment'
  | 'slip'
  | 'jam'
  | 'navigation'
  | 'unreachable'
  | 'overspeed';

const adapter = new MainThreadPhysicsAdapter();
let source: ProjectManifest | undefined;
let active: DeterministicBaseline | undefined;
let task: AgentTaskController | undefined;
let generation = 0;
const api = {
  async boot(url: string) {
    const response = await fetch(url);
    if (!response.ok) throw new Error(`A 资产 HTTP ${response.status}`);
    source = await importBrowserFiles([new File([await response.blob()], 'lekiwi-v1.zip')]);
  },
  async baseline(seed: number, governed = false) {
    const epoch = ++generation;
    task?.dispose();
    task = undefined;
    active?.dispose();
    active = undefined;
    if (!source) throw new Error('尚未导入 A 资产');
    await adapter.load(source, 'URDF/LeKiwi.urdf', {
      robotProfileId: 'lekiwi-v1',
      configureRobotRuntime: false,
      mobileRobot: governed ? PICK_PLACE_ROBOT : undefined,
      urdfMode: 'mjcf',
      baseMode: 'floating',
      map: { kind: 'none' },
      sceneComposer: (xml) => composePickPlaceScene(xml, seed, governed),
    });
    adapter.releaseRetired();
    const session = adapter.session!;
    let baseline!: DeterministicBaseline;
    let ownerAtStart = 'manual';
    if (governed) {
      session.startAgent((write) => {
        baseline = new DeterministicBaseline(session, seed, write);
        let cancelled = false;
        return {
          beforeStep: () => {
            if (cancelled) return false;
            baseline.beforeStep();
            return baseline.status === 'running';
          },
          afterStep: (dt) => {
            baseline.afterStep(dt);
            if (baseline.status !== 'running') session.stopAgent(baseline.status);
          },
          hold: () => baseline.hold(),
          cancel: () => {
            cancelled = true;
          },
          dispose: () => baseline.dispose(),
        };
      });
      session.setPaused(false);
      ownerAtStart = session.snapshot().controlOwner!;
    } else
      baseline = new DeterministicBaseline(session, seed, (id, value) =>
        session.setActuator(id, value),
      );
    active = baseline;
    const opt = session.model.opt,
      dt = Number(opt.timestep);
    opt.delete();
    let steps = 0;
    const began = performance.now();
    while (baseline.status === 'running') {
      if (epoch !== generation) throw new Error('基线已取消');
      if (!governed) baseline.beforeStep();
      if (baseline.status !== 'running') break;
      session.singleStep();
      if (++steps % 20 === 0) session.module.mj_forward(session.model, session.data);
      if (!governed) baseline.afterStep(dt);
      if (steps % 200 === 0) await new Promise<void>((resolve) => setTimeout(resolve, 0));
    }
    const result = {
      seed,
      ownerAtStart,
      status: baseline.status,
      error: baseline.error,
      steps,
      simTime: Number(session.data.time),
      wallMs: performance.now() - began,
      transported: baseline.evaluator.transported,
      succeeded: baseline.evaluator.succeeded,
      maxJointSpeed: baseline.maxJointSpeed,
      final: baseline.sample,
      trace: baseline.trace,
    };
    baseline.dispose();
    active = undefined;
    return result;
  },
  async task(
    seed: number,
    connection?: { address: string; token: string },
    options: { fault?: Fault; goal?: number[] } = {},
  ) {
    if (connection && options.fault) throw new Error('故障夹具禁止真实 API');
    const epoch = ++generation;
    task?.dispose();
    active?.dispose();
    if (!source) throw new Error('尚未导入 A 资产');
    await adapter.load(source, 'URDF/LeKiwi.urdf', {
      robotProfileId: 'lekiwi-v1',
      configureRobotRuntime: false,
      mobileRobot: PICK_PLACE_ROBOT,
      urdfMode: 'mjcf',
      baseMode: 'floating',
      map: { kind: 'none' },
      sceneComposer: (xml) => composePickPlaceScene(xml, seed, true, options.goal),
    });
    adapter.releaseRetired();
    const session = adapter.session!;
    const provider = connection
      ? new DecisionClient(connection.address, connection.token)
      : new MockDecisionProvider();
    let controller!: AgentTaskController;
    let skills!: DeterministicBaseline;
    session.startAgent((write) => {
      const b = new ModelBindings(session.model),
        r = PICK_PLACE_ROBOT;
      const grip = b.scalarActuator(
        r.gripperActuators[0].name,
        r.gripperJoint,
        'position',
      ).actuatorId;
      const jam = b.scalarActuator(r.armActuators[1], r.armJoints[1].name, 'position').actuatorId;
      const wheels = new Set(
        r.baseActuators.map(
          (name, i) => b.scalarActuator(name, r.baseJoints[i], 'velocity').actuatorId,
        ),
      );
      let jamTarget = 0,
        slipAt: number | undefined,
        slipOpening = 0;
      // Test-only actuator faults: no live qpos writes, attachments or object teleportation.
      const disturbedWrite = (id: number, value: number) => {
        if (options.fault === 'jam' && id === jam) value = jamTarget;
        if (options.fault === 'overspeed' && id === jam) value = r.armJoints[1].max;
        if (options.fault === 'navigation' && wheels.has(id)) value = 0;
        if (
          id === grip &&
          (options.fault === 'empty-grasp' ||
            (options.fault === 'empty-once' && controller.failures.length === 0))
        )
          value = r.gripperOpen;
        if (id === grip && options.fault === 'slip') {
          if (
            slipAt === undefined &&
            skills.phase === 'carry' &&
            skills.evaluator.transported > 0.12
          ) {
            slipAt = skills.sample.time;
            slipOpening = skills.sample.opening;
          }
          if (slipAt !== undefined) {
            const opening = Math.min(
              1,
              slipOpening + P.gripperOpeningRate * (skills.sample.time - slipAt),
            );
            value = r.gripperClosed + opening * (r.gripperOpen - r.gripperClosed);
          }
        }
        write(id, value);
      };
      skills = new DeterministicBaseline(session, seed, disturbedWrite, true, options.goal);
      jamTarget = skills.sample.joints[1];
      if (options.fault === 'unreachable' || options.fault === 'alignment') {
        const solve = skills.ik.solve.bind(skills.ik);
        // Actual scratch-data IK, not fake success/failure results or live qpos writes.
        skills.ik.solve = (target) =>
          solve(
            options.fault === 'unreachable' ? [2, 2, 2] : [target[0] + 0.02, target[1], target[2]],
          );
      }
      controller = new AgentTaskController(
        skills,
        provider,
        '将方块抓起，使用底盘搬运到指定支撑台并释放、撤离。',
      );
      return controller;
    });
    task = controller;
    active = skills;
    session.setPaused(false);
    const ownerAtStart = session.snapshot().controlOwner;
    const began = performance.now();
    let frames = 0;
    while (!controller.terminal) {
      if (epoch !== generation) throw new Error('任务已取消');
      session.singleStep();
      if (++frames % 20 === 0 && controller.state === 'executing')
        session.module.mj_forward(session.model, session.data);
      if (frames % 200 === 0) await new Promise<void>((resolve) => setTimeout(resolve, 0));
    }
    const result = {
      ...controller.snapshot(),
      fault: options.fault ?? null,
      ownerAtStart,
      seed,
      simTime: Number(session.data.time),
      wallMs: performance.now() - began,
      maxJointSpeed: skills.maxJointSpeed,
      results: controller.results,
      failures: controller.failures,
      trace: skills.trace,
      receipts: provider instanceof DecisionClient ? provider.receipts : [],
    };
    session.stopAgent('任务夹具已结束');
    task = undefined;
    active = undefined;
    return result;
  },
  async lifecycle(action: 'pause-task' | 'pause-global' | 'stop' | 'force' | 'reload') {
    task?.dispose();
    active?.dispose();
    if (!source) throw new Error('尚未导入 A 资产');
    const load = () =>
      adapter.load(source!, 'URDF/LeKiwi.urdf', {
        robotProfileId: 'lekiwi-v1',
        configureRobotRuntime: false,
        mobileRobot: PICK_PLACE_ROBOT,
        urdfMode: 'mjcf',
        baseMode: 'floating',
        map: { kind: 'none' },
        sceneComposer: (xml) => composePickPlaceScene(xml, 0, true),
      });
    await load();
    adapter.releaseRetired();
    const session = adapter.session!,
      provider = new MockDecisionProvider();
    let release!: () => void;
    const delayed = new Promise<void>((resolve) => {
      release = resolve;
    });
    const plan = provider.plan.bind(provider);
    // Deliberately ignore abort AFTER creating the reply: prove the local epoch gate,
    // not just HTTP cancellation, prevents a broken provider from moving real WASM.
    provider.plan = async (...args) => {
      const reply = await plan(...args);
      await delayed;
      return reply;
    };
    let controller!: AgentTaskController, skills!: DeterministicBaseline;
    session.startAgent((write) => {
      skills = new DeterministicBaseline(session, 0, write, true);
      controller = new AgentTaskController(skills, provider, '抓取并搬运方块');
      return controller;
    });
    task = controller;
    active = skills;
    session.setPaused(false);
    session.singleStep();
    await new Promise<void>((resolve) => setTimeout(resolve, 0));
    if (controller.state !== 'planning') throw new Error('生命周期夹具未进入等待规划');
    if (action === 'pause-task') controller.pause();
    if (action === 'pause-global') session.setPaused(true);
    if (action === 'stop') session.stopAgent('夹具停止');
    if (action === 'force') session.setExternalForce(1, [1, 0, 0]);
    if (action === 'reload') {
      await load();
      adapter.releaseRetired();
    }
    const current = adapter.session!;
    const before = {
      time: Number(current.data.time),
      ctrl: Array.from(current.data.ctrl),
      qpos: Array.from(current.data.qpos),
    };
    release();
    await new Promise<void>((resolve) => setTimeout(resolve, 20));
    for (let i = 0; i < 10; i++) current.advance(performance.now() + i * 20);
    const after = {
      time: Number(current.data.time),
      ctrl: Array.from(current.data.ctrl),
      qpos: Array.from(current.data.qpos),
    };
    const result = {
      action,
      source: 'delayed-mock-plus-real-WASM',
      state: controller.state,
      owner: current.controlOwner,
      jevCalls: controller.jevCalls,
      before,
      after,
    };
    controller.cancel('生命周期验收结束');
    return result;
  },
  async resources() {
    const rows: { heapBytes: number; resources: { isDeleted(): boolean }[] }[] = [];
    for (let i = 0; i < 8; i++) {
      await api.lifecycle('stop');
      const current = adapter.session!;
      // Test-only reflection retains handles, not their typed-array views, to verify
      // native deletions rather than inferring a leak from linear-memory capacity.
      const ik = active!.ik as unknown as {
        scratch: { isDeleted(): boolean };
        positionJacobian: { isDeleted(): boolean };
        rotationJacobian: { isDeleted(): boolean };
      };
      const sampler = active! as unknown as { sampler: { force: { isDeleted(): boolean } } };
      rows.push({
        heapBytes: (current.data.qpos as Float64Array).buffer.byteLength,
        resources: [
          current.model,
          current.data,
          ik.scratch,
          ik.positionJacobian,
          ik.rotationJacobian,
          sampler.sampler.force,
        ],
      });
      task?.dispose();
      active?.dispose();
    }
    api.dispose();
    return rows.map((row) => ({
      heapBytes: row.heapBytes,
      deleted: row.resources.map((resource) => resource.isDeleted()),
    }));
  },
  dispose() {
    generation++;
    task?.dispose();
    task = undefined;
    active?.dispose();
    active = undefined;
    adapter.dispose();
  },
};
declare global {
  interface Window {
    lekiwiAgent: typeof api;
  }
}
window.lekiwiAgent = api;