Files
Mujoco_WASM/web_platform/src/project/urdfToMjcf.ts
T
chenlin f4b415c54f
web-platform-ci / TypeScript、Lint、Unit、Build (push) Has been cancelled
web-platform-ci / Playwright E2E (push) Has been cancelled
web-platform-ci / TypeScript、Lint、Unit、Build (pull_request) Has been cancelled
web-platform-ci / Playwright E2E (pull_request) Has been cancelled
refactor(web-platform): release V0.6 精简代码
2026-08-28 14:10:16 +08:00

102 lines
9.0 KiB
TypeScript

const decoder=new TextDecoder('utf-8');
const encoder=new TextEncoder();
function numbers(value:string|undefined,count:number):number[]{
const parsed=(value??'').trim().split(/\s+/).filter(Boolean).map(Number);
return Array.from({length:count},(_,index)=>Number.isFinite(parsed[index])?parsed[index]:0);
}
export type UrdfBaseMode='floating'|'fixed';
export type CameraDirection='+X'|'-X'|'+Y'|'-Y'|'+Z'|'-Z';
export interface UrdfEnhancementOptions {addActuators:boolean;addSensors:boolean;sensorType:'camera';cameraMountBody?:string;cameraPosition?:[number,number,number];cameraDirection?:CameraDirection;}
export interface UrdfEnhancementResult {data:Uint8Array;actuatorCount:number;cameraAdded:boolean;imuAdded:boolean;unitreeGo2wTuned:boolean;}
function uniqueName(document:Document,selector:string,base:string):string {
const occupied=new Set(Array.from(document.querySelectorAll(selector)).map(element=>element.getAttribute('name')).filter(Boolean));
if(!occupied.has(base))return base;
for(let index=2;;index+=1)if(!occupied.has(`${base}_${index}`))return `${base}_${index}`;
}
function tuneUnitreeGo2w(document:Document):boolean {
const prefixes=['FL','FR','RL','RR'],legParts=['hip','thigh','calf'];
const expected=[...prefixes.flatMap(prefix=>legParts.map(part=>`${prefix}_${part}_joint`)),...prefixes.map(prefix=>`${prefix}_foot_joint`)];
const joints=new Map(Array.from(document.querySelectorAll('worldbody joint[name]')).map(joint=>[joint.getAttribute('name')!,joint]));
if(!expected.every(name=>joints.has(name)))return false;
const root=document.querySelector('mujoco');if(!root)return false;
let option=Array.from(root.children).find(element=>element.tagName==='option');if(!option){option=document.createElement('option');root.prepend(option);}
option.setAttribute('cone','elliptic');option.setAttribute('impratio','100');
for(const name of expected){const joint=joints.get(name)!;joint.setAttribute('armature','0.01');joint.setAttribute('damping','0.1');joint.setAttribute('frictionloss','0.2');}
for(const motor of Array.from(document.querySelectorAll('actuator > *[joint]'))){const jointName=motor.getAttribute('joint')??'';if(!expected.includes(jointName))continue;const limit=jointName.includes('_calf_')?45.43:jointName.endsWith('_foot_joint')?15:23.7;motor.setAttribute('forcelimited','true');motor.setAttribute('forcerange',`${-limit} ${limit}`);}
for(const geom of Array.from(document.querySelectorAll('worldbody body geom'))){if(geom.getAttribute('contype')==='0'||geom.getAttribute('group')==='1')continue;geom.setAttribute('friction','0.4 0.005 0.0001');geom.setAttribute('margin','0.001');geom.setAttribute('condim','1');const body=geom.closest('body'),name=body?.getAttribute('name')??'';if(/_(?:foot|wheel)(?:_link)?$/i.test(name)){geom.setAttribute('friction','0.8 0.02 0.01');geom.setAttribute('condim','6');geom.setAttribute('priority','1');}}
return true;
}
/** 为转换后的 MJCF 按需补充可直接控制的关节电机和一台自动取景相机。 */
export function enhanceConvertedMjcf(data:Uint8Array,options:UrdfEnhancementOptions):UrdfEnhancementResult {
const document=new DOMParser().parseFromString(decoder.decode(data),'application/xml');
if(document.querySelector('parsererror'))throw new Error('MuJoCo 导出的 MJCF XML 无法解析');
const root=document.querySelector('mujoco'),worldbody=document.querySelector('mujoco > worldbody');
if(!root||!worldbody)throw new Error('MuJoCo 导出的 MJCF 缺少 worldbody');
let actuatorCount=0;
if(options.addActuators){
let actuator=Array.from(root.children).find(element=>element.tagName==='actuator');
if(!actuator){actuator=document.createElement('actuator');root.append(actuator);}
const driven=new Set(Array.from(actuator.children).flatMap(element=>[element.getAttribute('joint'),element.getAttribute('jointinparent')]).filter(Boolean));
for(const joint of Array.from(worldbody.querySelectorAll('joint'))){
const type=joint.getAttribute('type')??'hinge';
if(type!=='hinge'&&type!=='slide')continue;
let jointName=joint.getAttribute('name');
if(!jointName){jointName=uniqueName(document,'joint[name]','__platform_joint__');joint.setAttribute('name',jointName);}
if(driven.has(jointName))continue;
const motor=document.createElement('motor');
motor.setAttribute('name',uniqueName(document,'actuator > *[name]',`${jointName}_motor`));
motor.setAttribute('joint',jointName);motor.setAttribute('ctrllimited','false');motor.setAttribute('forcelimited','true');motor.setAttribute('forcerange','-100 100');
if(!joint.hasAttribute('stiffness'))joint.setAttribute('stiffness','0');if(!joint.hasAttribute('damping'))joint.setAttribute('damping','0');
actuator.append(motor);driven.add(jointName);actuatorCount+=1;
}
if(!actuator.children.length)actuator.remove();
}
const unitreeGo2wTuned=tuneUnitreeGo2w(document);
let imuAdded=false;
if(options.addSensors&&!document.querySelector('sensor > gyro[name="imu_gyro"], sensor > gyro[name="__platform_imu_gyro__"]')){
const bodies=Array.from(worldbody.querySelectorAll('body')),mount=bodies.find(body=>/^(?:base|base_link|trunk)$/i.test(body.getAttribute('name')??''))??Array.from(worldbody.children).find(element=>element.tagName==='body');
if(mount){
const site=document.createElement('site'),siteName=uniqueName(document,'site[name]','imu');site.setAttribute('name',siteName);site.setAttribute('pos','0 0 0');site.setAttribute('size','0.005');site.setAttribute('rgba','0.2 0.8 1 0.5');mount.append(site);
let sensor=Array.from(root.children).find(element=>element.tagName==='sensor');if(!sensor){sensor=document.createElement('sensor');root.append(sensor);}
const gyro=document.createElement('gyro');gyro.setAttribute('name',uniqueName(document,'sensor > *[name]','imu_gyro'));gyro.setAttribute('site',siteName);gyro.setAttribute('noise','0');sensor.append(gyro);
const accelerometer=document.createElement('accelerometer');accelerometer.setAttribute('name',uniqueName(document,'sensor > *[name]','imu_acc'));accelerometer.setAttribute('site',siteName);accelerometer.setAttribute('noise','0');sensor.append(accelerometer);imuAdded=true;
}
}
let cameraAdded=false;
if(options.addSensors&&options.sensorType==='camera'&&!document.querySelector('camera[name^="__platform_camera__"]')){
const bodies=Array.from(worldbody.querySelectorAll('body')),preferred=options.cameraMountBody?bodies.find(body=>body.getAttribute('name')===options.cameraMountBody):undefined;
const mount=preferred??bodies.find(body=>/(head|camera|sensor|neck|头)/i.test(body.getAttribute('name')??''))??bodies.at(-1);
if(mount){
const directions:Record<CameraDirection,string>={'+X':'0 -1 0 0 0 1','-X':'0 1 0 0 0 1','+Y':'1 0 0 0 0 1','-Y':'-1 0 0 0 0 1','+Z':'0 1 0 1 0 0','-Z':'0 -1 0 1 0 0'},position=options.cameraPosition??[.1,0,.05],direction=options.cameraDirection??'+X';
const camera=document.createElement('camera');camera.setAttribute('name',uniqueName(document,'camera[name]','__platform_camera__'));camera.setAttribute('mode','fixed');
camera.setAttribute('pos',position.map(value=>Number.isFinite(value)?value:0).join(' '));camera.setAttribute('xyaxes',directions[direction]);camera.setAttribute('fovy','45');camera.setAttribute('resolution','640 480');mount.append(camera);cameraAdded=true;
}
}
return {data:encoder.encode(new XMLSerializer().serializeToString(document)),actuatorCount,cameraAdded,imuAdded,unitreeGo2wTuned};
}
/** 给 MuJoCo 从 URDF 导出的 MJCF 添加地面、设置基座类型,并整体抬升根 body。 */
export function groundConvertedMjcf(data:Uint8Array,minimumZ:number,baseMode:UrdfBaseMode='fixed'):Uint8Array {
const document=new DOMParser().parseFromString(decoder.decode(data),'application/xml');
if(document.querySelector('parsererror'))throw new Error('MuJoCo 导出的 MJCF XML 无法解析');
const worldbody=document.querySelector('mujoco > worldbody');
if(!worldbody)throw new Error('MuJoCo 导出的 MJCF 缺少 worldbody');
const lift=Number.isFinite(minimumZ)?-minimumZ:0;
const rootBodies=Array.from(worldbody.children).filter(element=>element.tagName==='body');
for(const [index,body] of rootBodies.entries()){
const pos=numbers(body.getAttribute('pos')??undefined,3);pos[2]+=lift;body.setAttribute('pos',pos.join(' '));
if(baseMode==='floating'&&!Array.from(body.children).some(element=>element.tagName==='freejoint'||element.tagName==='joint')){
const freejoint=document.createElement('freejoint');freejoint.setAttribute('name',rootBodies.length===1?'__platform_base_freejoint__':`__platform_base_freejoint_${index}__`);body.prepend(freejoint);
}
}
const floor=document.createElement('geom');
floor.setAttribute('name','__platform_ground__');floor.setAttribute('type','plane');floor.setAttribute('size','1 1 0.1');floor.setAttribute('pos','0 0 0');floor.setAttribute('rgba','0.12 0.16 0.22 1');floor.setAttribute('group','5');floor.setAttribute('friction','1 0.005 0.0001');
worldbody.prepend(floor);
return encoder.encode(new XMLSerializer().serializeToString(document));
}