Mujoco_WASM/wasm/web_platform/src/simulation/SimulationSession.ts

import type {MainModule, MjData, MjModel, MjvPerturb, MjvScene} from '@mujoco/mujoco';
import {meshIdFromSceneDataId} from './geometry';

export interface ActuatorInfo {id: number; name: string; value: number; min: number; max: number; limited: boolean;}
export interface JointInfo {id: number; name: string; type: number; value: number; min: number; max: number; limited: boolean; editable: boolean; bodyId: number; axis: [number, number, number];}
export interface BodyInfo {id: number; name: string;}
export interface SimulationSnapshot {time: number; qpos: number[]; qvel: number[]; ctrl: number[]; actuators: ActuatorInfo[]; joints: JointInfo[]; bodies: BodyInfo[]; warnings: string[]; model: {nbody: number; njnt: number; ngeom: number; nu: number; nq: number; nv: number};}
export interface FrameResult {steps: number; stepMs: number; overBudget: boolean;}

export class SimulationSession {
  readonly model: MjModel;
  readonly data: MjData;
  readonly perturb: MjvPerturb;
  paused = true;
  speed = 1;
  readonly frameBudgetMs = 8;
  readonly maxCatchUpSteps = 100;
  private accumulator = 0;
  private lastNow?: number;
  private forceBody = -1;
  private force: [number, number, number] = [0, 0, 0];
  private disposed = false;

  constructor(readonly module: MainModule, modelPath: string, readonly warnings: string[] = []) {
    let model: MjModel | undefined; let data: MjData | undefined; let perturb: MjvPerturb | undefined;
    try {
      model = module.MjModel.mj_loadXML(modelPath) ?? undefined;
      if (!model) throw new Error(`MuJoCo 无法编译模型：${modelPath}`);
      data = new module.MjData(model);
      perturb = new module.MjvPerturb(); module.mjv_defaultPerturb(perturb);
      this.model = model; this.data = data; this.perturb = perturb;
      module.mj_forward(model, data);
    } catch (error) { perturb?.delete(); data?.delete(); model?.delete(); throw error; }
  }

  setPaused(paused: boolean): void {this.paused = paused; this.accumulator = 0; this.lastNow = undefined;}
  setSpeed(speed: number): void {this.speed = Math.min(4, Math.max(0.1, speed));}
  reset(): void {this.setPaused(true);this.module.mj_resetData(this.model,this.data);this.module.mj_forward(this.model,this.data);this.clearExternalForce();}
  singleStep(): void {this.applyForce(); this.module.mj_step(this.model, this.data);}

  advance(now: number): FrameResult {
    if (this.lastNow === undefined) {this.lastNow = now; return {steps: 0, stepMs: 0, overBudget: false};}
    const elapsed = Math.min(0.1, Math.max(0, (now - this.lastNow) / 1000)); this.lastNow = now;
    if (this.paused) return {steps: 0, stepMs: 0, overBudget: false};
    this.accumulator += elapsed * this.speed;
    const dt = Number(this.model.opt.timestep) || 0.002; const started = performance.now(); let steps = 0;
    while (this.accumulator >= dt && steps < this.maxCatchUpSteps && performance.now() - started < this.frameBudgetMs) {
      this.applyForce(); this.module.mj_step(this.model, this.data); this.accumulator -= dt; steps++;
    }
    const overBudget = this.accumulator >= dt;
    if (steps >= this.maxCatchUpSteps) this.accumulator = Math.min(this.accumulator, dt);
    return {steps, stepMs: performance.now() - started, overBudget};
  }

  setActuator(id: number, value: number): void {
    if (id < 0 || id >= this.model.nu) return;
    const actuator = this.model.actuator(id);
    try {
      const limited = Boolean(actuator.ctrllimited);
      const min = limited ? Number(actuator.ctrlrange[0]) : -1;
      const max = limited ? Number(actuator.ctrlrange[1]) : 1;
      const address = Number(this.model.actuator_ctrladr[id] ?? id);
      this.data.ctrl[address] = Math.min(max, Math.max(min, value));
    } finally {
      actuator.delete();
    }
  }

  setJointPosition(id: number, value: number): boolean {
    if (id < 0 || id >= this.model.njnt) return false;
    const joint = this.model.jnt(id);
    try {
      const type = Number(joint.type);
      if (type !== 2 && type !== 3) return false;
      const limited = Boolean(joint.limited);
      const min = limited ? Number(joint.range[0]) : -Math.PI;
      const max = limited ? Number(joint.range[1]) : Math.PI;
      this.setPaused(true);
      this.data.qpos[Number(joint.qposadr)] = Math.min(max, Math.max(min, value));
      this.module.mj_forward(this.model, this.data);
      return true;
    } finally {
      joint.delete();
    }
  }

  setExternalForce(bodyId: number, force: [number, number, number]): void {this.forceBody = bodyId > 0 && bodyId < this.model.nbody ? bodyId : -1; this.force = force;}
  clearExternalForce(): void {this.forceBody = -1; this.force = [0, 0, 0]; this.data.xfrc_applied.fill(0); this.perturb.active = 0;}
  initializePerturb(scene: MjvScene, bodyId: number): void {this.perturb.select = bodyId; this.module.mjv_initPerturb(this.model, this.data, scene, this.perturb);}
  applyPerturbForce(): void {if (this.forceBody > 0) this.module.mjv_applyPerturbForce(this.model, this.data, this.perturb);}

  private applyForce(): void {
    this.data.xfrc_applied.fill(0); if (this.forceBody < 1) return;
    this.applyPerturbForce(); const offset = this.forceBody * 6;
    this.data.xfrc_applied[offset] += this.force[0]; this.data.xfrc_applied[offset + 1] += this.force[1]; this.data.xfrc_applied[offset + 2] += this.force[2];
  }

  /** 用有限几何的包围球估算视图中心与范围，忽略地面等无限平面。 */
  geometryBounds():{center:[number,number,number];extent:number} {
    const lower=[Number.POSITIVE_INFINITY,Number.POSITIVE_INFINITY,Number.POSITIVE_INFINITY];const upper=[Number.NEGATIVE_INFINITY,Number.NEGATIVE_INFINITY,Number.NEGATIVE_INFINITY];
    for(let geom=0;geom<this.model.ngeom;geom+=1){if(Number(this.model.geom_type[geom])===this.module.mjtGeom.mjGEOM_PLANE.value)continue;const radius=Math.max(0,Number(this.model.geom_rbound[geom]));for(let axis=0;axis<3;axis+=1){const value=Number(this.data.geom_xpos[geom*3+axis]);lower[axis]=Math.min(lower[axis],value-radius);upper[axis]=Math.max(upper[axis],value+radius);}}
    if(!lower.every(Number.isFinite)||!upper.every(Number.isFinite))return {center:[0,0,0],extent:2};
    return {center:[(lower[0]+upper[0])/2,(lower[1]+upper[1])/2,(lower[2]+upper[2])/2],extent:Math.max(.5,upper[0]-lower[0],upper[1]-lower[1],upper[2]-lower[2])};
  }

  /** 返回当前姿态全部模型几何（不含无限平面）的最低世界坐标。 */
  minimumGeometryZ():number {
    const option=new this.module.MjvOption();const camera=new this.module.MjvCamera();const scene=new this.module.MjvScene(this.model,32768);
    this.module.mjv_defaultOption(option);this.module.mjv_defaultCamera(camera);
    let minimum=Number.POSITIVE_INFINITY;
    try{
      this.module.mjv_updateScene(this.model,this.data,option,this.perturb,camera,this.module.mjtCatBit.mjCAT_ALL.value,scene);
      const geoms=scene.geoms;
      try{for(let index=0;index<geoms.size();index+=1){const geom=geoms.get(index);if(!geom)continue;try{
        if(geom.type===this.module.mjtGeom.mjGEOM_PLANE.value)continue;
        const z=this.geomMinimumZ(geom);if(Number.isFinite(z))minimum=Math.min(minimum,z);
      }finally{geom.delete();}}}finally{geoms.delete();}
    }finally{scene.delete();camera.delete();option.delete();}
    return Number.isFinite(minimum)?minimum:0;
  }

  /** 平移所有世界根 body，使当前姿态的最低点位于 z=0。 */
  alignLowestPointToGround():number {
    const offset=-this.minimumGeometryZ();
    if(Math.abs(offset)<1e-9)return 0;
    for(let body=1;body<this.model.nbody;body+=1)if(Number(this.model.body_parentid[body])===0)this.model.body_pos[body*3+2]+=offset;
    this.module.mj_forward(this.model,this.data);
    return offset;
  }

  private geomMinimumZ(geom:import('@mujoco/mujoco').MjvGeom):number {
    const m=this.module,type=geom.type,s=geom.size,r0=geom.mat[6],r1=geom.mat[7],r2=geom.mat[8],center=geom.pos[2];
    if(type===m.mjtGeom.mjGEOM_SPHERE.value)return center-s[0];
    if(type===m.mjtGeom.mjGEOM_BOX.value)return center-(Math.abs(r0)*s[0]+Math.abs(r1)*s[1]+Math.abs(r2)*s[2]);
    if(type===m.mjtGeom.mjGEOM_ELLIPSOID.value)return center-Math.hypot(r0*s[0],r1*s[1],r2*s[2]);
    if(type===m.mjtGeom.mjGEOM_CYLINDER.value)return center-(Math.hypot(r0,r1)*s[0]+Math.abs(r2)*s[2]);
    if(type===m.mjtGeom.mjGEOM_CAPSULE.value)return center-(s[0]+Math.abs(r2)*s[2]);
    if(type===m.mjtGeom.mjGEOM_MESH.value&&geom.dataid>=0){
      const id=meshIdFromSceneDataId(geom.dataid),first=Number(this.model.mesh_vertadr[id]),count=Number(this.model.mesh_vertnum[id]);let minimum=Number.POSITIVE_INFINITY;
      for(let vertex=0;vertex<count;vertex+=1){const offset=(first+vertex)*3;minimum=Math.min(minimum,center+r0*this.model.mesh_vert[offset]+r1*this.model.mesh_vert[offset+1]+r2*this.model.mesh_vert[offset+2]);}
      return minimum;
    }
    const radius=geom.size[0]||0;return center-radius;
  }

  snapshot(): SimulationSnapshot {
    const actuators = Array.from({length: this.model.nu}, (_, id): ActuatorInfo => {
      const actuator = this.model.actuator(id);
      try {
        const limited = Boolean(actuator.ctrllimited);
        const address = Number(this.model.actuator_ctrladr[id] ?? id);
        return {id,name:actuator.name || `actuator_${id}`,value:Number(this.data.ctrl[address]),min:limited?Number(actuator.ctrlrange[0]):-1,max:limited?Number(actuator.ctrlrange[1]):1,limited};
      } finally {
        actuator.delete();
      }
    });
    const joints = Array.from({length: this.model.njnt}, (_, id): JointInfo => {
      const joint = this.model.jnt(id);
      try {
        const type=Number(joint.type); const limited=Boolean(joint.limited);
        return {id,name:joint.name || `joint_${id}`,type,value:Number(this.data.qpos[Number(joint.qposadr)]),min:limited?Number(joint.range[0]):-Math.PI,max:limited?Number(joint.range[1]):Math.PI,limited,editable:type===2||type===3,bodyId:Number(joint.bodyid),axis:[Number(joint.axis[0]),Number(joint.axis[1]),Number(joint.axis[2])]};
      } finally {
        joint.delete();
      }
    });
    const bodies = Array.from({length: this.model.nbody}, (_,id): BodyInfo => {
      const body = this.model.body(id);
      try { return {id,name:body.name || `body_${id}`}; }
      finally { body.delete(); }
    });
    return {time:Number(this.data.time),qpos:Array.from(this.data.qpos),qvel:Array.from(this.data.qvel),ctrl:Array.from(this.data.ctrl),actuators,joints,bodies,warnings:this.warnings,model:{nbody:this.model.nbody,njnt:this.model.njnt,ngeom:this.model.ngeom,nu:this.model.nu,nq:this.model.nq,nv:this.model.nv}};
  }
  dispose(): void {if(this.disposed)return; this.disposed=true; this.clearExternalForce(); this.perturb.delete(); this.data.delete(); this.model.delete();}
}