ika

examples/IKSolver.js

@@ -0,0 +1,68 @@
+import { IK, IKChain, IKJoint, IKBallConstraint } from 'three-ik';
+import * as THREE from 'three';
+
+export class IKSolver {
+  constructor(model) {
+    this.model = model;
+    this.ik = new IK();
+    this.chains = {};
+  }
+
+  createChain(chainName, jointNames) {
+    const chain = new IKChain();
+    const joints = jointNames.map((name, index) => {
+      const bone = new THREE.Bone();
+      bone.name = name;
+      bone.position.set(0, index * 0.1, 0);
+      const joint = new IKJoint(bone);
+      return joint;
+    });
+
+    joints.forEach((joint, i) => {
+      chain.add(joint, { constraints: [new IKBallConstraint(180)] });
+    });
+
+    this.ik.add(chain);
+    this.chains[chainName] = chain;
+  }
+
+  solve(chainName, currentPositions, targetPosition) {
+    console.log('Entering solve method');
+    console.log('Chain name:', chainName);
+    console.log('Current positions:', currentPositions);
+    console.log('Target position:', targetPosition);
+
+    const chain = this.chains[chainName];
+    if (!chain) {
+      console.error(`Chain ${chainName} not found`);
+      return null;
+    }
+
+    console.log('Chain found');
+
+    // Set current positions
+    let prevPosition = null;
+    chain.joints.forEach((joint, i) => {
+      let position = currentPositions[i].clone();
+      if (prevPosition && position.distanceTo(prevPosition) < 0.01) {
+        // Add a significant offset if positions are too close
+        const offset = new THREE.Vector3(0, 0.01, 0);
+        position.add(offset);
+      }
+      joint.position.copy(position);
+      prevPosition = position;
+    });
+
+    // Set target
+    chain.target.position.copy(targetPosition);
+
+    // Solve IK
+    console.log('Before solving:', chain.joints.map(j => j.position.toArray()));
+    this.ik.solve();
+    console.log('After solving:', chain.joints.map(j => j.position.toArray()));
+
+
+    // Return new joint rotations
+    return chain.joints.map(joint => joint.rotation.toArray());
+  }
+}

examples/main.js

@@ -5,6 +5,7 @@ import { OrbitControls    } from '../node_modules/three/examples/jsm/controls/Or
 import { DragStateManager } from './utils/DragStateManager.js';
 import { setupGUI, downloadExampleScenesFolder, loadSceneFromURL, getPosition, getQuaternion, toMujocoPos, standardNormal } from './mujocoUtils.js';
 import   load_mujoco        from '../dist/mujoco_wasm.js';
+import { IKSolver } from './IKSolver.js';
 
 // Load the MuJoCo Module
 const mujoco = await load_mujoco();
@@ -68,10 +69,29 @@ export class MuJoCoDemo {
     this.controls.screenSpacePanning = true;
     this.controls.update();
 
+    // User actuator/model direct control
     this.actuatorNames = [];
     this.actuatorRanges = [];
     this.loadPPOModel();
 
+    // IKA
+    console.log('Initializing IK Solver');
+    this.ikSolver = new IKSolver(this.model);
+    console.log('IK Solver initialized');
+    this.endEffectors = {
+      'right_hand': { chain: ['right_shoulder', 'right_elbow', 'right_wrist'] },
+      'left_hand': { chain: ['left_shoulder', 'left_elbow', 'left_wrist'] },
+      // Add more end effectors as needed
+    };
+    this.bodyNameToId = this.createBodyNameToIdMap();
+
+    console.log('Initializing IK Chains');
+    // Initialize IK chains
+    for (const [effector, data] of Object.entries(this.endEffectors)) {
+      this.ikSolver.createChain(effector, data.chain);
+    }
+    console.log('Initializing IK Solver');
+
     window.addEventListener('resize', this.onWindowResize.bind(this));
 
     // Initialize the Drag State Manager.
@@ -93,6 +113,69 @@ export class MuJoCoDemo {
     // this.initializeActuators();
   }
 
+  /* IKA */
+  createBodyNameToIdMap() {
+    const map = {};
+    const textDecoder = new TextDecoder();
+    for (let i = 0; i < this.model.nbody; i++) {
+      const name = textDecoder.decode(
+        this.model.names.subarray(
+          this.model.name_bodyadr[i]
+        )
+      ).split('\0')[0];
+      map[name] = i;
+    }
+    return map;
+  }
+
+  getBodyIdByName(name) {
+    return this.bodyNameToId[name];
+  }
+
+  handleDrag(bodyName, targetPosition) {
+    if (this.endEffectors[bodyName]) {
+      const chain = this.endEffectors[bodyName].chain;
+      const currentJointPositions = this.getCurrentJointPositions(chain);
+      const newJointRotations = this.ikSolver.solve(bodyName, currentJointPositions, targetPosition);
+      this.applyJointRotations(chain, newJointRotations);
+    }
+  }
+
+  getCurrentJointPositions(chain) {
+    return chain.map(bodyName => {
+      const bodyId = this.getBodyIdByName(bodyName);
+      return new THREE.Vector3().fromArray(this.simulation.xpos.subarray(bodyId * 3, bodyId * 3 + 3));
+    });
+  }
+
+  applyJointRotations(chain, rotations) {
+    chain.forEach((bodyName, i) => {
+      const bodyId = this.getBodyIdByName(bodyName);
+      const jointId = this.model.body_jntadr[bodyId];
+      const qposAddr = this.model.jnt_qposadr[jointId];
+
+      // Apply rotation to qpos
+      // This assumes that the joint is a ball joint with 3 DoF
+      // You may need to adjust this based on your specific joint types
+      this.simulation.qpos[qposAddr] = rotations[i][0];
+      this.simulation.qpos[qposAddr + 1] = rotations[i][1];
+      this.simulation.qpos[qposAddr + 2] = rotations[i][2];
+    });
+
+    // Update forward kinematics
+    this.simulation.forward();
+  }
+
+  getBodyNameById(bodyId) {
+    const textDecoder = new TextDecoder();
+    return textDecoder.decode(
+      this.model.names.subarray(
+        this.model.name_bodyadr[bodyId]
+      )
+    ).split('\0')[0];
+  }
+
+  /* User actuator/model direct control */
   // does this ordering align with the ordering of the model output?
   initializeActuators() {
     const textDecoder = new TextDecoder();

examples/mujocoUtils.js

@@ -511,6 +511,7 @@ export async function loadSceneFromURL(mujoco, filename, parent) {
       } else if (type == mujoco.mjtGeom.mjGEOM_SPHERE.value) {
         geometry = new THREE.SphereGeometry(size[0]);
       } else if (type == mujoco.mjtGeom.mjGEOM_CAPSULE.value) {
+        // imported capsule geometry
         geometry = new THREE.CapsuleGeometry(size[0], size[1] * 2.0, 20, 20);
       } else if (type == mujoco.mjtGeom.mjGEOM_ELLIPSOID.value) {
         geometry = new THREE.SphereGeometry(1); // Stretch this below