[software] 添加16DOF早期训练仿真与Sim2Real闭环

2026-07-21 16:15:14 +08:00
parent 9bd22225f9
commit e9e2c946b3
681 changed files with 137221 additions and 8 deletions
@@ -0,0 +1,181 @@
+/**
+ * RobotViewer3D — sim2real 3D 可视化（基于 robot_viewer 的 MJCFAdapter + Three.js）
+ *
+ * 加载 wheelleg.xml → MJCFAdapter.parse → Three.js 场景树
+ * 建立 jointName → THREE.Object3D 映射，通过 updateJoints(pos16) 实时更新。
+ * 支持 OrbitControls 旋转/缩放/平移。
+ *
+ * 用法：
+ *   const viewer = new RobotViewer3D(canvasElement);
+ *   await viewer.load('/mjcf/wheelleg.xml');
+ *   viewer.updateJoints(jointPositions16);
+ */
+import * as THREE from 'three';
+import { OrbitControls } from 'three/examples/jsm/controls/OrbitControls.js';
+import { MJCFAdapter } from './MJCFAdapter.js';
+import { STLLoader } from 'three/examples/jsm/loaders/STLLoader.js';
+
+// 16 关节的标准顺序（与 motor_mapping.py:SIM_JOINT_ORDER 对齐）
+const JOINT_ORDER = [
+    'fl_hip_abduction_joint', 'fl_hip_pitch_joint', 'fl_knee_joint',
+    'fr_hip_abduction_joint', 'fr_hip_pitch_joint', 'fr_knee_joint',
+    'rl_hip_abduction_joint', 'rl_hip_pitch_joint', 'rl_knee_joint',
+    'rr_hip_abduction_joint', 'rr_hip_pitch_joint', 'rr_knee_joint',
+    'fl_wheel_joint', 'fr_wheel_joint', 'rl_wheel_joint', 'rr_wheel_joint',
+];
+
+// MJCF → Three.js 坐标轴转换：让 MJCF 的 Z 轴(向上) 映射到 Three.js 的 Y 轴(向上)
+const MJCF_TO_THREE = new THREE.Matrix4().makeRotationX(-Math.PI / 2);
+// 或直接用 euler: (0, PI, 0)
+
+export class RobotViewer3D {
+    /**
+     * @param {HTMLCanvasElement} canvas
+     * @param {object} [opts]
+     * @param {string} [opts.meshBaseUrl='/meshes/']  STL mesh 文件的 HTTP 路径前缀
+     * @param {string} [opts.mjcfUrl='/mjcf/wheelleg.xml']
+     * @param {string} [opts.backgroundColor='#1a1d24']
+     */
+    constructor(canvas, opts = {}) {
+        this.canvas = canvas;
+        this.meshBaseUrl = opts.meshBaseUrl || '/meshes/';
+        this.mjcfUrl = opts.mjcfUrl || '/mjcf/wheelleg.xml';
+
+        // Three.js 核心
+        const w = canvas.clientWidth, h = canvas.clientHeight;
+        this.scene = new THREE.Scene();
+        // 移除背景色，使用透明背景，由 CSS 控制
+        // this.scene.background = new THREE.Color(opts.backgroundColor || '#1a1d24');
+
+        this.camera = new THREE.PerspectiveCamera(55, w / h, 0.05, 50);
+        this.camera.position.set(0.5, 0.35, 0.65);
+        this.camera.lookAt(0.2, 0, 0);
+
+        this.renderer = new THREE.WebGLRenderer({ canvas, antialias: true, alpha: true });
+        this.renderer.setSize(w, h);
+        this.renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2));
+        this.renderer.shadowMap.enabled = true;
+
+        // OrbitControls
+        this.controls = new OrbitControls(this.camera, canvas);
+        this.controls.target.set(0.15, 0.08, 0.0);
+        this.controls.enableDamping = true;
+        this.controls.dampingFactor = 0.12;
+        this.controls.update();
+
+        // 灯光
+        this._setupLights();
+
+        // 地面
+        const grid = new THREE.GridHelper(2, 20, 0x444444, 0x222222);
+        grid.position.y = -0.35;
+        this.scene.add(grid);
+
+        // 状态
+        this.model = null;
+        this.rootGroup = null;
+        this.jointMap = new Map();           // jointName → { joint, group }
+        this._isLoaded = false;
+        this._rafId = null;
+        this._stlCache = new Map();          // filename → BufferGeometry
+    }
+
+    _setupLights() {
+        const ambient = new THREE.AmbientLight(0x606060, 1.5);
+        this.scene.add(ambient);
+
+        const dir1 = new THREE.DirectionalLight(0xffffff, 2.5);
+        dir1.position.set(2, 3, 2);
+        this.scene.add(dir1);
+
+        const dir2 = new THREE.DirectionalLight(0x8899cc, 1.0);
+        dir2.position.set(-1, 1, -1);
+        this.scene.add(dir2);
+
+        const hemi = new THREE.HemisphereLight(0x8899cc, 0x334455, 1.2);
+        this.scene.add(hemi);
+    }
+
+    // ---- 加载模型 ----
+    async load(mjcfUrlOverride) {
+        const url = mjcfUrlOverride || this.mjcfUrl;
+        console.log('[RobotViewer3D] loading MJCF:', url);
+        const resp = await fetch(url);
+        if (!resp.ok) throw new Error(`MJCF 404: ${url}`);
+        const xmlText = await resp.text();
+
+        // 用 MJCFAdapter 解析 → UnifiedRobotModel
+        // fileMap 为空时不传；MeshLoader 会自动 fallback 到 URL 加载
+        const model = await MJCFAdapter.parse(xmlText, null);
+        this.model = model;
+        console.log('[RobotViewer3D] parsed:', model.links.size, 'links,', model.joints.size, 'joints');
+
+        // 取 rootGroup（MJCFAdapter.createThreeObject 已构建完整 hierarchy）
+        this.rootGroup = model.threeObject;
+        // 坐标轴转换：MJCF → Three.js
+        this.rootGroup.applyMatrix4(MJCF_TO_THREE);
+        this.scene.add(this.rootGroup);
+
+        // 遍历 joints，建立索引
+        this.jointMap.clear();
+        for (const [jointName, joint] of model.joints) {
+            if (joint.threeObject) {
+                this.jointMap.set(jointName, joint);
+            }
+        }
+        // 已建立映射的关节列表
+        const mapped = Array.from(this.jointMap.keys()).sort();
+        console.log('[RobotViewer3D] joint map:', mapped.length, 'joints');
+
+        this._isLoaded = true;
+        this._startRenderLoop();
+    }
+
+    // ---- 渲染循环（按需 + 持续） ----
+    _startRenderLoop() {
+        if (this._rafId) return;
+        const loop = () => {
+            this.controls.update();
+            this.renderer.render(this.scene, this.camera);
+            this._rafId = requestAnimationFrame(loop);
+        };
+        loop();
+    }
+
+    // ---- 实时更新关节角度 ----
+    /**
+     * @param {Float64Array|number[]} pos16 — 16 关节角度 (rad)，顺序同 SIM_JOINT_ORDER
+     *   索引 0-11: 腿关节 (fl_abd,fl_pitch,fl_knee,fr...,rl...,rr...)
+     *   索引 12-15: 轮子关节 (fl_wheel,fr_wheel,rl_wheel,rr_wheel)
+     */
+    updateJoints(pos16) {
+        if (!this._isLoaded) return;
+        for (let i = 0; i < JOINT_ORDER.length && i < pos16.length; i++) {
+            const name = JOINT_ORDER[i];
+            const joint = this.jointMap.get(name);
+            if (joint) {
+                MJCFAdapter.setJointAngle(joint, pos16[i]);
+            }
+        }
+    }
+
+    // ---- 重置相机 ----
+    resetCamera() {
+        this.camera.position.set(0.5, 0.35, 0.65);
+        this.controls.target.set(0.15, 0.08, 0.0);
+        this.controls.update();
+    }
+
+    // ---- 调整大小 ----
+    resize() {
+        const w = this.canvas.clientWidth, h = this.canvas.clientHeight;
+        this.camera.aspect = w / h;
+        this.camera.updateProjectionMatrix();
+        this.renderer.setSize(w, h);
+    }
+
+    dispose() {
+        if (this._rafId) cancelAnimationFrame(this._rafId);
+        this.renderer.dispose();
+    }
+}