Mujoco_WASM/web_platform/physics/mobile.ts

import { RobotManager } from '../src/mobile/RobotManager';
import { MOBILE_ROBOTS } from '../src/mobile/RobotDescriptor';
import type { MobilePolicyMetadata } from '../src/mobile/ONNXPolicyRunner';
import { MuJoCoViewer } from '../src/viewer/MuJoCoViewer';
import { TaskDragController } from '../src/mobile/TaskDragController';
import { Vector3 } from 'three';

let viewer: MuJoCoViewer | undefined;
let drag: TaskDragController | undefined;
const manager = new RobotManager({
  attach(session) {
    viewer?.attach(session);
  },
});
function snapshot() {
  const env = manager.env!;
  const observation = env.observe();
  return {
    robotId: env.config.id,
    state: Array.from(env.state),
    observation: Array.from(observation),
    ctrl: Array.from(env.session.data.ctrl, Number),
    qpos: Array.from(env.session.data.qpos, Number),
    time: Number(env.session.data.time),
    reward: env.kernel.result.reward,
    stage: env.kernel.result.info.stage,
    revision: env.revision,
    heap: env.wasmHeapBytes,
    version: env.session.module.mj_versionString(),
    workspace: manager.adapter.workspace!.root,
  };
}
const api = {
  async boot(url: string, id: string) {
    const config = MOBILE_ROBOTS.find((c) => c.id === id)!;
    const response = await fetch(url);
    if (!response.ok) throw new Error(`asset HTTP ${response.status}`);
    const previous = manager.adapter.session;
    await manager.loadZip(new File([await response.blob()], 'robot.zip'), config);
    return {
      ...snapshot(),
      retiredModelDeleted: previous?.model.isDeleted() ?? true,
      retiredDataDeleted: previous?.data.isDeleted() ?? true,
    };
  },
  async invalidCandidate(url: string) {
    const config = structuredClone(MOBILE_ROBOTS[1]);
    const wrong = {
      ...config,
      armActuators: ['missing_actuator', ...config.armActuators.slice(1)],
    };
    await manager.loadZip(new File([await (await fetch(url)).blob()], 'robot.zip'), wrong);
  },
  step(n: number, action?: number[], manual = false) {
    const env = manager.env!;
    if (action) env.action.set(action);
    const started = performance.now();
    for (let i = 0; i < n; i++) env.step(env.action, manual);
    return { ...snapshot(), wallMs: performance.now() - started };
  },
  reset() {
    manager.reset();
    return snapshot();
  },
  move(entity: 'object' | 'goal', position: number[]) {
    manager.move(entity, position);
    return snapshot();
  },
  state: snapshot,
  workspaces() {
    const module = manager.adapter.session!.module as unknown as {
      FS: { readdir(p: string): string[] };
    };
    return module.FS.readdir('/workspace').filter((s) => s !== '.' && s !== '..');
  },
  async download() {
    const bytes = await manager.exportTrainingBundle();
    const href = URL.createObjectURL(new Blob([bytes.slice().buffer]));
    const link = document.createElement('a');
    link.href = href;
    link.download = 'training.zip';
    link.click();
    setTimeout(() => URL.revokeObjectURL(href), 1000);
  },
  async policy(modelUrl: string, metadata: MobilePolicyMetadata) {
    await manager.loadPolicy(new Uint8Array(await (await fetch(modelUrl)).arrayBuffer()), metadata);
    manager.setMode('policy');
    manager.setRunning(true);
  },
  tick(now: number) {
    return {
      ...manager.advance(now),
      metrics: manager.policy?.metrics,
      time: Number(manager.env!.session.data.time),
    };
  },
  show() {
    const host = document.createElement('div');
    host.style.cssText = 'width:1000px;height:650px;position:relative';
    document.body.append(host);
    const noop = () => {};
    viewer = new MuJoCoViewer(host, {
      advance: (now) => manager.advance(now),
      onSelection: noop,
      onFrame: noop,
      onError: (e) => {
        throw e;
      },
      onMapEditorSelect: noop,
      onMapEditorPreviewSelect: noop,
      onParametricMapSelect: noop,
      onParametricMapTransform: noop,
      onMapEditorTransform: noop,
    });
    viewer.attach(manager.adapter.session);
    drag = new TaskDragController(viewer, manager);
  },
  screen(entity: 'object' | 'goal') {
    const s = manager.env!.state,
      offset = entity === 'object' ? 37 : 47;
    const p = new Vector3(s[offset], s[offset + 1], s[offset + 2]).project(viewer!.camera);
    const r = viewer!.renderer.domElement.getBoundingClientRect();
    return { x: r.left + (p.x + 1) * 0.5 * r.width, y: r.top + (1 - p.y) * 0.5 * r.height };
  },
  async dispose() {
    drag?.dispose();
    await manager.dispose();
    viewer?.dispose();
  },
};
declare global {
  interface Window {
    mobilePhysics: typeof api;
  }
}
window.mobilePhysics = api;