406 lines
16 KiB
TypeScript
406 lines
16 KiB
TypeScript
import type { MainModule } from '@mujoco/mujoco';
|
||
import type { ProjectFile, ProjectManifest } from '../project/types';
|
||
import { prepareProjectForMujoco } from '../project/importer';
|
||
import {
|
||
enhanceConvertedMjcf,
|
||
groundConvertedMjcf,
|
||
type UrdfBaseMode,
|
||
type UrdfEnhancementOptions,
|
||
} from '../project/urdfToMjcf';
|
||
import { MemfsWorkspace } from '../project/workspace';
|
||
import {
|
||
SimulationSession,
|
||
type ActuatorParameters,
|
||
type FrameResult,
|
||
type SimulationSnapshot,
|
||
} from './SimulationSession';
|
||
import type { ControllerCommand, ControllerStatus } from '../controller/types';
|
||
import type { RLCommand, RLPolicyStatus } from '../rl/types';
|
||
import type {
|
||
DataRecorderConfig,
|
||
DataRecorderStatus,
|
||
TelemetryChannel,
|
||
} from '../telemetry/DataRecorder';
|
||
import { composePhysicalMap } from '../map/physicalMap';
|
||
import { composeProjectMap } from '../map/MapComposer';
|
||
import { resolveProjectMap } from '../map/MapLoader';
|
||
import { DEFAULT_MAP_SELECTION, type MapSelection } from '../map/types';
|
||
|
||
export type UrdfLoadMode = 'mjcf' | 'native';
|
||
export type { UrdfBaseMode, UrdfEnhancementOptions };
|
||
|
||
export interface PhysicsLoadProgress {
|
||
value: number;
|
||
label: string;
|
||
}
|
||
|
||
export interface PhysicsLoadOptions {
|
||
urdfMode?: UrdfLoadMode;
|
||
baseMode?: UrdfBaseMode;
|
||
enhancements?: UrdfEnhancementOptions;
|
||
map?: MapSelection;
|
||
onProgress?: (progress: PhysicsLoadProgress) => void;
|
||
}
|
||
|
||
export interface PhysicsAdapter {
|
||
load(
|
||
manifest: ProjectManifest,
|
||
entryPath: string,
|
||
options?: PhysicsLoadOptions,
|
||
): Promise<SimulationSnapshot>;
|
||
advance(now: number): FrameResult;
|
||
snapshot(): SimulationSnapshot | null;
|
||
setPaused(paused: boolean): void;
|
||
setSpeed(speed: number): void;
|
||
reset(): void;
|
||
singleStep(): void;
|
||
setActuator(id: number, value: number): void;
|
||
setActuatorParameters(id: number, parameters: ActuatorParameters): boolean;
|
||
setJointPosition(id: number, value: number): boolean;
|
||
resetJoints(): void;
|
||
setIgnoreJointLimits(ignore: boolean): void;
|
||
setExternalForce(bodyId: number, force: [number, number, number]): void;
|
||
clearExternalForce(): void;
|
||
loadPythonController(source: string, path: string): Promise<ControllerStatus>;
|
||
setControllerEnabled(enabled: boolean): void;
|
||
sendControllerCommand(command: ControllerCommand): void;
|
||
removeController(): void;
|
||
loadRLPolicy(model: Uint8Array, path: string): Promise<RLPolicyStatus>;
|
||
setRLPolicyEnabled(enabled: boolean): void;
|
||
setRLCommand(command: RLCommand): void;
|
||
removeRLPolicy(): void;
|
||
configureDataRecorder(config: Partial<DataRecorderConfig>): DataRecorderStatus | undefined;
|
||
startDataRecording(): DataRecorderStatus | undefined;
|
||
stopDataRecording(): DataRecorderStatus | undefined;
|
||
clearDataRecording(): DataRecorderStatus | undefined;
|
||
registerDataChannel(channel: TelemetryChannel): (() => void) | undefined;
|
||
exportDataRecording(format: 'csv' | 'json'): Uint8Array;
|
||
cachedSupportFiles(): ProjectFile[];
|
||
releaseRetired(): void;
|
||
rollbackRetired(): void;
|
||
exportMjcf(): Uint8Array;
|
||
dispose(): void;
|
||
}
|
||
|
||
let modulePromise: Promise<MainModule> | undefined;
|
||
export function getMujocoModule(): Promise<MainModule> {
|
||
if (!modulePromise) {
|
||
console.info('[MuJoCo] 开始初始化单线程 WASM');
|
||
modulePromise = import('@mujoco/mujoco')
|
||
.then(({ default: loadMujoco }) => loadMujoco())
|
||
.then((module) => {
|
||
console.info('[MuJoCo] WASM 初始化完成');
|
||
return module;
|
||
})
|
||
.catch((error) => {
|
||
// 初始化失败后允许用户重试，而不是永久复用已 rejected 的 Promise。
|
||
modulePromise = undefined;
|
||
throw error;
|
||
});
|
||
}
|
||
return modulePromise;
|
||
}
|
||
|
||
export class MainThreadPhysicsAdapter implements PhysicsAdapter {
|
||
session: SimulationSession | null = null;
|
||
workspace: MemfsWorkspace | null = null;
|
||
private supportFiles: ProjectFile[] = [];
|
||
private loadGeneration = 0;
|
||
private disposed = false;
|
||
private retiredSession: SimulationSession | null = null;
|
||
private retiredWorkspace: MemfsWorkspace | null = null;
|
||
private retiredSupportFiles: ProjectFile[] = [];
|
||
|
||
async load(
|
||
manifest: ProjectManifest,
|
||
entryPath: string,
|
||
options: PhysicsLoadOptions = {},
|
||
): Promise<SimulationSnapshot> {
|
||
if (this.disposed) throw new Error('物理适配器已释放');
|
||
const generation = ++this.loadGeneration;
|
||
const urdfMode = options.urdfMode ?? 'mjcf';
|
||
const baseMode = options.baseMode ?? 'floating';
|
||
const enhancements = options.enhancements ?? {
|
||
addActuators: false,
|
||
addSensors: false,
|
||
sensorType: 'camera',
|
||
};
|
||
const mapSelection = options.map ?? DEFAULT_MAP_SELECTION;
|
||
const report = (value: number, label: string) => {
|
||
if (!this.disposed && generation === this.loadGeneration)
|
||
options.onProgress?.({ value, label });
|
||
};
|
||
report(0.08, '初始化 MuJoCo WebAssembly');
|
||
const module = await getMujocoModule();
|
||
if (this.disposed || generation !== this.loadGeneration) throw new Error('模型加载已取消');
|
||
report(0.32, '校验并准备模型资源');
|
||
const workspace = new MemfsWorkspace(module, `${manifest.id}_stage_${generation}`);
|
||
const prepared = await prepareProjectForMujoco(manifest, entryPath);
|
||
const supportFiles = prepared.manifest.files.filter(
|
||
(file) => !manifest.files.some((original) => original.path === file.path),
|
||
);
|
||
let nextSession: SimulationSession | null = null;
|
||
try {
|
||
report(0.46, '写入浏览器内存文件系统');
|
||
console.info('[MuJoCo] 写入 MEMFS', prepared.manifest.files.length);
|
||
workspace.mount(prepared.manifest);
|
||
const entry = prepared.manifest.entries.find((candidate) => candidate.path === entryPath);
|
||
let modelRelativePath = entryPath,
|
||
modelPath = workspace.path(modelRelativePath);
|
||
const warnings = [...prepared.warnings];
|
||
if (entry?.format === 'urdf' && urdfMode === 'mjcf') {
|
||
report(0.58, '转换并增强 URDF 模型');
|
||
console.info('[MuJoCo] 编译 URDF 中间模型', entryPath);
|
||
const intermediate = new SimulationSession(module, modelPath);
|
||
try {
|
||
const minimumZ = intermediate.minimumGeometryZ();
|
||
const slash = entryPath.lastIndexOf('/');
|
||
const directory = slash >= 0 ? entryPath.slice(0, slash + 1) : '';
|
||
const convertedPath = `${directory}.__mujoco_converted_${manifest.id.replace(/[^a-zA-Z0-9_-]/g, '_')}.xml`;
|
||
if (module.mj_saveLastXML(workspace.path(convertedPath), intermediate.model) === 0)
|
||
throw new Error('MuJoCo 无法导出中间 MJCF');
|
||
const grounded = groundConvertedMjcf(
|
||
new TextEncoder().encode(workspace.readText(convertedPath)),
|
||
minimumZ,
|
||
baseMode,
|
||
);
|
||
const enhanced = enhanceConvertedMjcf(grounded, enhancements);
|
||
workspace.writeGenerated(convertedPath, enhanced.data);
|
||
modelRelativePath = convertedPath;
|
||
modelPath = workspace.path(modelRelativePath);
|
||
warnings.push(
|
||
`URDF 已转换为 MJCF（${baseMode === 'floating' ? '浮动基座' : '固定基座'}），并整体平移 ${(-minimumZ).toFixed(4)} m，使最低点接触 z=0 地面`,
|
||
);
|
||
if (enhanced.actuatorCount)
|
||
warnings.push(
|
||
`已为 ${enhanced.actuatorCount} 个 hinge/slide 关节生成 motor 驱动器（控制输入不限幅；hinge 输出单位 N·m，slide 输出单位 N）`,
|
||
);
|
||
if (enhanced.unitreeGo2wTuned)
|
||
warnings.unshift(
|
||
'已识别 Unitree Go2-W，并补齐官方 MuJoCo 关节惯量/阻尼、力矩限幅和轮胎接触参数',
|
||
);
|
||
if (enhanced.imuAdded)
|
||
warnings.unshift(
|
||
'已在浮动基座添加6轴 IMU：imu_gyro（三轴角速度）和 imu_acc（三轴加速度）',
|
||
);
|
||
if (enhanced.cameraAdded)
|
||
warnings.push(
|
||
`已将 640×480 摄像头固连到 ${enhancements.cameraMountBody || '自动选择的头部/末端 body'}，局部位置 ${(enhancements.cameraPosition ?? [0.1, 0, 0.05]).join(' ')} m，朝向 ${enhancements.cameraDirection ?? '+X'}`,
|
||
);
|
||
} finally {
|
||
intermediate.dispose();
|
||
}
|
||
}
|
||
if (mapSelection.kind !== 'none') {
|
||
report(0.72, '组合机器人与物理地图');
|
||
if (entry?.format === 'urdf' && urdfMode === 'native')
|
||
throw new Error('原生 URDF 模式暂不支持地图，请切换为转换模式');
|
||
const slash = modelRelativePath.lastIndexOf('/');
|
||
const directory = slash >= 0 ? modelRelativePath.slice(0, slash + 1) : '';
|
||
const mapPath = `${directory}.__mujoco_map_scene_${manifest.id.replace(/[^a-zA-Z0-9_-]/g, '_')}.xml`;
|
||
if (mapSelection.kind === 'builtin') {
|
||
const composed = composePhysicalMap(
|
||
new TextEncoder().encode(workspace.readText(modelRelativePath)),
|
||
mapSelection.config,
|
||
);
|
||
workspace.writeGenerated(mapPath, composed.data);
|
||
if (composed.summary) warnings.push(composed.summary);
|
||
} else {
|
||
const resolvedMap = resolveProjectMap(prepared.manifest, mapSelection.descriptorPath);
|
||
const composed = composeProjectMap(
|
||
new TextEncoder().encode(workspace.readText(modelRelativePath)),
|
||
mapPath,
|
||
prepared.manifest,
|
||
resolvedMap,
|
||
mapSelection,
|
||
);
|
||
workspace.writeGenerated(mapPath, composed.data);
|
||
warnings.push(...composed.warnings, composed.summary);
|
||
}
|
||
modelRelativePath = mapPath;
|
||
modelPath = workspace.path(modelRelativePath);
|
||
}
|
||
report(0.84, '编译模型与物理数据');
|
||
console.info('[MuJoCo] 编译模型', modelPath);
|
||
nextSession = new SimulationSession(module, modelPath, warnings);
|
||
if (entry?.format === 'urdf' && urdfMode === 'native') {
|
||
const offset = nextSession.alignLowestPointToGround();
|
||
warnings.push(`原生 URDF 已整体平移 ${offset.toFixed(4)} m，使最低点位于 z=0`);
|
||
}
|
||
if (warnings.length) console.info('[MuJoCo] 兼容与地图处理', warnings);
|
||
console.info('[MuJoCo] 模型编译完成');
|
||
report(0.94, '生成初始仿真状态');
|
||
const snapshot = nextSession.snapshot();
|
||
if (this.disposed || generation !== this.loadGeneration) throw new Error('模型加载已取消');
|
||
this.releaseRetired();
|
||
this.retiredSession = this.session;
|
||
this.retiredWorkspace = this.workspace;
|
||
this.retiredSupportFiles = this.supportFiles;
|
||
this.session = nextSession;
|
||
this.workspace = workspace;
|
||
this.supportFiles = supportFiles;
|
||
nextSession = null;
|
||
console.info('[MuJoCo] 状态快照完成');
|
||
return snapshot;
|
||
} catch (error) {
|
||
nextSession?.dispose();
|
||
workspace.dispose();
|
||
throw new Error(
|
||
`模型编译失败（${entryPath}）：${error instanceof Error ? error.message : String(error)}`,
|
||
{ cause: error },
|
||
);
|
||
}
|
||
}
|
||
advance(now: number): FrameResult {
|
||
return this.session?.advance(now) ?? { steps: 0, stepMs: 0, overBudget: false };
|
||
}
|
||
snapshot(): SimulationSnapshot | null {
|
||
return this.session?.snapshot() ?? null;
|
||
}
|
||
setPaused(value: boolean): void {
|
||
this.session?.setPaused(value);
|
||
}
|
||
setSpeed(value: number): void {
|
||
this.session?.setSpeed(value);
|
||
}
|
||
reset(): void {
|
||
this.session?.reset();
|
||
}
|
||
singleStep(): void {
|
||
this.session?.singleStep();
|
||
}
|
||
setActuator(id: number, value: number): void {
|
||
this.session?.setActuator(id, value);
|
||
}
|
||
setActuatorParameters(id: number, parameters: ActuatorParameters): boolean {
|
||
return this.session?.setActuatorParameters(id, parameters) ?? false;
|
||
}
|
||
setJointPosition(id: number, value: number): boolean {
|
||
return this.session?.setJointPosition(id, value) ?? false;
|
||
}
|
||
resetJoints(): void {
|
||
this.session?.resetJoints();
|
||
}
|
||
setIgnoreJointLimits(ignore: boolean): void {
|
||
this.session?.setIgnoreJointLimits(ignore);
|
||
}
|
||
setExternalForce(bodyId: number, force: [number, number, number]): void {
|
||
this.session?.setExternalForce(bodyId, force);
|
||
}
|
||
clearExternalForce(): void {
|
||
this.session?.clearExternalForce();
|
||
}
|
||
async loadPythonController(source: string, path: string): Promise<ControllerStatus> {
|
||
if (!this.session) throw new Error('请先加载模型');
|
||
return this.session.loadPythonController(source, path);
|
||
}
|
||
setControllerEnabled(enabled: boolean): void {
|
||
this.session?.setControllerEnabled(enabled);
|
||
}
|
||
sendControllerCommand(command: ControllerCommand): void {
|
||
this.session?.sendControllerCommand(command);
|
||
}
|
||
removeController(): void {
|
||
this.session?.removeController();
|
||
}
|
||
async loadRLPolicy(model: Uint8Array, path: string): Promise<RLPolicyStatus> {
|
||
if (!this.session) throw new Error('请先加载模型');
|
||
return this.session.loadRLPolicy(model, path);
|
||
}
|
||
setRLPolicyEnabled(enabled: boolean): void {
|
||
this.session?.setRLPolicyEnabled(enabled);
|
||
}
|
||
setRLCommand(command: RLCommand): void {
|
||
this.session?.setRLCommand(command);
|
||
}
|
||
removeRLPolicy(): void {
|
||
this.session?.removeRLPolicy();
|
||
}
|
||
configureDataRecorder(config: Partial<DataRecorderConfig>): DataRecorderStatus | undefined {
|
||
return this.session?.configureDataRecorder(config);
|
||
}
|
||
startDataRecording(): DataRecorderStatus | undefined {
|
||
return this.session?.startDataRecording();
|
||
}
|
||
stopDataRecording(): DataRecorderStatus | undefined {
|
||
return this.session?.stopDataRecording();
|
||
}
|
||
clearDataRecording(): DataRecorderStatus | undefined {
|
||
return this.session?.clearDataRecording();
|
||
}
|
||
registerDataChannel(channel: TelemetryChannel): (() => void) | undefined {
|
||
return this.session?.registerDataChannel(channel);
|
||
}
|
||
exportDataRecording(format: 'csv' | 'json'): Uint8Array {
|
||
if (!this.session) throw new Error('请先加载模型');
|
||
return this.session.exportDataRecording(format);
|
||
}
|
||
cachedSupportFiles(): ProjectFile[] {
|
||
return this.supportFiles.map((file) => ({ ...file, data: file.data.slice() }));
|
||
}
|
||
releaseRetired(): void {
|
||
this.retiredSession?.dispose();
|
||
this.retiredSession = null;
|
||
this.retiredWorkspace?.dispose();
|
||
this.retiredWorkspace = null;
|
||
this.retiredSupportFiles = [];
|
||
}
|
||
rollbackRetired(): void {
|
||
this.session?.dispose();
|
||
this.workspace?.dispose();
|
||
this.session = this.retiredSession;
|
||
this.workspace = this.retiredWorkspace;
|
||
this.supportFiles = this.retiredSupportFiles;
|
||
this.retiredSession = null;
|
||
this.retiredWorkspace = null;
|
||
this.retiredSupportFiles = [];
|
||
}
|
||
exportMjcf(): Uint8Array {
|
||
if (!this.session || !this.workspace) throw new Error('尚未加载可导出的模型');
|
||
const relative = '.__platform_export__.xml';
|
||
if (this.session.module.mj_saveLastXML(this.workspace.path(relative), this.session.model) === 0)
|
||
throw new Error('MuJoCo 无法生成 MJCF');
|
||
const source = this.workspace.readText(relative),
|
||
document = new DOMParser().parseFromString(source, 'application/xml'),
|
||
actuatorSection = document.querySelector('mujoco > actuator');
|
||
if (document.querySelector('parsererror')) return new TextEncoder().encode(source);
|
||
const snapshot = this.session.snapshot();
|
||
if (actuatorSection) {
|
||
for (const info of snapshot.actuators) {
|
||
const element = Array.from(actuatorSection.children).find(
|
||
(candidate) => candidate.getAttribute('name') === info.name,
|
||
);
|
||
if (!element) continue;
|
||
element.setAttribute('ctrllimited', info.ctrlLimited ? 'true' : 'false');
|
||
element.setAttribute('forcelimited', info.forceLimited ? 'true' : 'false');
|
||
}
|
||
}
|
||
for (const info of snapshot.actuators) {
|
||
if (info.kind !== 'motor' || !info.jointName) continue;
|
||
const joint = Array.from(document.querySelectorAll('worldbody joint[name]')).find(
|
||
(candidate) => candidate.getAttribute('name') === info.jointName,
|
||
);
|
||
if (joint) {
|
||
joint.setAttribute('stiffness', String(info.kp));
|
||
joint.setAttribute('damping', String(info.kv));
|
||
}
|
||
}
|
||
const output = new TextEncoder().encode(new XMLSerializer().serializeToString(document));
|
||
this.workspace.writeGenerated(relative, output);
|
||
return output;
|
||
}
|
||
private releaseCurrent(): void {
|
||
this.releaseRetired();
|
||
this.session?.dispose();
|
||
this.session = null;
|
||
this.workspace?.dispose();
|
||
this.workspace = null;
|
||
this.supportFiles = [];
|
||
}
|
||
dispose(): void {
|
||
this.disposed = true;
|
||
this.loadGeneration += 1;
|
||
this.releaseCurrent();
|
||
}
|
||
}
|