model.js

// Engine-independent MuJoCo model API. The same code runs in a Worker and Node tests.
let nextMount=0
export async function loadModel(config, progress=()=>{}) {
  if (!config.modelUrl) throw Error('modelUrl is required')
  progress('Loading MuJoCo WASM…')
  const moduleUrl = config.runtimeUrl || new URL('./runtime/mujoco.js',import.meta.url).href
  const {default:load}=await import(moduleUrl)
  const mj=await load({locateFile:p=>new URL(p,moduleUrl).href})
  const url=new URL(config.modelUrl)
  const entry=decodeURIComponent(url.pathname.split('/').at(-1))
  const checked=async path=>{const response=await fetch(path);if(!response.ok)throw Error(`Asset HTTP ${response.status}: ${path}`);return response}
  const files=config.manifestUrl ? await (await checked(config.manifestUrl)).json() : [entry]
  if(!Array.isArray(files)||!files.includes(entry))throw Error('Asset manifest must list the model XML and its dependencies')
  const mount=`/model-${++nextMount}`
  await Promise.all(files.map(async path=>{
    if(typeof path!=='string'||path.startsWith('/')||path.split('/').some(p=>p==='..'||p===''))throw Error('Manifest paths must be relative files')
    const bytes=new Uint8Array(await (await checked(new URL(path,url))).arrayBuffer())
    const dest=mount+'/'+path;mj.FS.mkdirTree(dest.slice(0,dest.lastIndexOf('/')));mj.FS.writeFile(dest,bytes)
  }))
  const model=mj.MjModel.from_xml_path(mount+'/'+entry)
  return {mj,model}
}
const vector=(value,size,label)=>{
  if(!value||value.length!==size||!Array.from(value).every(Number.isFinite))throw Error(`${label} requires ${size} finite numbers`)
  return value
}
export class ModelSimulation {
  constructor(mj,model){this.mj=mj;this.model=model;this.data=new mj.MjData(model);const opt=model.opt;this.timestep=opt.timestep;opt.delete();this.reset()}
  reset({qpos,qvel,keyframe}={}) {
    const {mj,model:m,data:d}=this
    if(qpos)vector(qpos,m.nq,'qpos');if(qvel)vector(qvel,m.nv,'qvel')
    if(keyframe!==undefined){
      const key=typeof keyframe==='string'?mj.mj_name2id(m,mj.mjtObj.mjOBJ_KEY.value,keyframe):keyframe
      if(!Number.isInteger(key)||key<0||key>=m.nkey)throw Error('Unknown keyframe')
      mj.mj_resetDataKeyframe(m,d,key)
    }else mj.mj_resetData(m,d)
    if(qpos)d.qpos.set(qpos);if(qvel)d.qvel.set(qvel)
    mj.mj_forward(m,d)
    return this.snapshot()
  }
  step({action,frames=1}={}) {
    const {mj,model:m,data:d}=this
    vector(action,m.nu,'action')
    if(!Number.isInteger(frames)||frames<1||frames>10000)throw Error('frames must be an integer from 1 to 10000')
    d.ctrl.set(action)
    for(let i=0;i<frames;i++)mj.mj_step(m,d)
    mj.mj_forward(m,d)
    if(!d.qpos.every(Number.isFinite)||!d.qvel.every(Number.isFinite))throw Error('Non-finite simulation state')
    return this.snapshot()
  }
  snapshot(){const d=this.data;return {time:d.time,qpos:d.qpos.slice(),qvel:d.qvel.slice(),ctrl:d.ctrl.slice(),sensors:d.sensordata.slice(),bodyPositions:d.xpos.slice(),bodyQuaternions:d.xquat.slice(),geomPositions:d.geom_xpos.slice(),geomMatrices:d.geom_xmat.slice(),contacts:d.ncon}}
  describe(){
    const {mj,model:m}=this
    const name=(type,i)=>mj.mj_id2name(m,type.value,i)||`${i}`
    const geometries=Array.from({length:m.ngeom},(_,i)=>{
      const meshId=m.geom_dataid[i],type=m.geom_type[i]
      const geom={id:i,name:name(mj.mjtObj.mjOBJ_GEOM,i),body:m.geom_bodyid[i],type,size:Array.from(m.geom_size.slice(i*3,i*3+3)),rgba:Array.from(m.geom_rgba.slice(i*4,i*4+4))}
      if(type===mj.mjtGeom.mjGEOM_MESH.value){
        const v=m.mesh_vertadr[meshId],f=m.mesh_faceadr[meshId]
        geom.vertices=m.mesh_vert.slice(v*3,(v+m.mesh_vertnum[meshId])*3)
        geom.faces=m.mesh_face.slice(f*3,(f+m.mesh_facenum[meshId])*3)
      }
      return geom
    })
    return {version:mj.mj_versionString(),nq:m.nq,nv:m.nv,nu:m.nu,nbody:m.nbody,ngeom:m.ngeom,timestep:this.timestep,
      bodies:Array.from({length:m.nbody},(_,i)=>name(mj.mjtObj.mjOBJ_BODY,i)),
      actuators:Array.from({length:m.nu},(_,i)=>({name:name(mj.mjtObj.mjOBJ_ACTUATOR,i),limited:!!m.actuator_ctrllimited[i],range:Array.from(m.actuator_ctrlrange.slice(i*2,i*2+2))})),geometries}
  }
  dispose(){if(!this.data)return;this.data.delete();this.model.delete();this.data=undefined;this.model=undefined}
}