mujoco-react 8.0.0 → 8.1.0

package/dist/index.js

@@ -956,7 +956,7 @@ function MujocoSimProvider({
     const arr = values instanceof Float64Array ? values : new Float64Array(values);
     data.qvel.set(arr.subarray(0, Math.min(arr.length, mjModelRef.current?.nv ?? 0)));
   }, []);
-  const getQpos = useCallback(() => {
+  const getQpos2 = useCallback(() => {
     return mjDataRef.current ? new Float64Array(mjDataRef.current.qpos) : new Float64Array(0);
   }, []);
   const getQvel = useCallback(() => {
@@ -976,7 +976,7 @@ function MujocoSimProvider({
       }
     }
   }, []);
-  const getCtrl = useCallback(() => {
+  const getCtrl2 = useCallback(() => {
     return mjDataRef.current ? new Float64Array(mjDataRef.current.ctrl) : new Float64Array(0);
   }, []);
   const applyForce = useCallback((bodyName, force, point) => {
@@ -1384,10 +1384,10 @@ function MujocoSimProvider({
       restoreState,
       setQpos,
       setQvel,
-      getQpos,
+      getQpos: getQpos2,
       getQvel,
       setCtrl,
-      getCtrl,
+      getCtrl: getCtrl2,
       applyForce,
       applyTorque: applyTorqueApi,
       setExternalForce,
@@ -1430,10 +1430,10 @@ function MujocoSimProvider({
       restoreState,
       setQpos,
       setQvel,
-      getQpos,
+      getQpos2,
       getQvel,
       setCtrl,
-      getCtrl,
+      getCtrl2,
       applyForce,
       applyTorqueApi,
       setExternalForce,
@@ -3098,60 +3098,133 @@ function ContactListener({
   });
   return null;
 }
+function getQpos(input, idx) {
+  const item = input[idx];
+  if (!item) return null;
+  if (Array.isArray(item)) return item;
+  return item.qpos;
+}
+function getCtrl(input, idx) {
+  const item = input[idx];
+  if (!item || Array.isArray(item)) return null;
+  return item.ctrl ?? null;
+}
 function useTrajectoryPlayer(trajectory, options = {}) {
   const { mjModelRef, mjDataRef, mujocoRef, pausedRef } = useMujocoContext();
-  const fps = options.fps ?? 30;
-  const loop = options.loop ?? false;
-  const playingRef = useRef(false);
+  const optionsRef = useRef(options);
+  optionsRef.current = options;
+  const stateRef = useRef("idle");
   const frameRef = useRef(0);
   const lastFrameTimeRef = useRef(0);
+  const speedRef = useRef(options.speed ?? 1);
+  const wasPausedRef = useRef(false);
+  const trajectoryRef = useRef(trajectory);
+  trajectoryRef.current = trajectory;
+  const setState = useCallback((next) => {
+    if (stateRef.current === next) return;
+    stateRef.current = next;
+    optionsRef.current.onStateChange?.(next);
+  }, []);
   const play = useCallback(() => {
-    playingRef.current = true;
-    pausedRef.current = true;
+    const traj = trajectoryRef.current;
+    if (traj.length === 0) return;
+    const mode = optionsRef.current.mode ?? "kinematic";
+    if (stateRef.current === "completed") {
+      frameRef.current = 0;
+    }
+    if (mode === "kinematic") {
+      wasPausedRef.current = pausedRef.current;
+      pausedRef.current = true;
+    }
     lastFrameTimeRef.current = performance.now();
-  }, [pausedRef]);
+    setState("playing");
+  }, [pausedRef, setState]);
   const pause = useCallback(() => {
-    playingRef.current = false;
-  }, []);
+    if (stateRef.current !== "playing") return;
+    setState("paused");
+  }, [setState]);
   const seek = useCallback((frameIdx) => {
-    frameRef.current = Math.max(0, Math.min(frameIdx, trajectory.length - 1));
+    const traj = trajectoryRef.current;
+    if (traj.length === 0) return;
+    frameRef.current = Math.max(0, Math.min(frameIdx, traj.length - 1));
     const model = mjModelRef.current;
     const data = mjDataRef.current;
-    if (!model || !data || !trajectory[frameRef.current]) return;
-    const qpos = trajectory[frameRef.current];
+    if (!model || !data) return;
+    const qpos = getQpos(traj, frameRef.current);
+    if (!qpos) return;
     for (let i = 0; i < Math.min(qpos.length, model.nq); i++) {
       data.qpos[i] = qpos[i];
     }
     mujocoRef.current.mj_forward(model, data);
-  }, [trajectory, mjModelRef, mjDataRef, mujocoRef]);
+  }, [mjModelRef, mjDataRef, mujocoRef]);
   const reset = useCallback(() => {
+    const mode = optionsRef.current.mode ?? "kinematic";
+    if (mode === "kinematic" && stateRef.current !== "idle") {
+      pausedRef.current = wasPausedRef.current;
+    }
     frameRef.current = 0;
-    playingRef.current = false;
-    pausedRef.current = false;
-  }, [pausedRef]);
+    setState("idle");
+  }, [pausedRef, setState]);
+  const setSpeed = useCallback((s) => {
+    speedRef.current = s;
+  }, []);
+  const complete = useCallback(() => {
+    const mode = optionsRef.current.mode ?? "kinematic";
+    if (mode === "kinematic") {
+      pausedRef.current = wasPausedRef.current;
+    }
+    setState("completed");
+    optionsRef.current.onComplete?.();
+  }, [pausedRef, setState]);
   useFrame(() => {
-    if (!playingRef.current || trajectory.length === 0) return;
+    if (stateRef.current !== "playing") return;
+    if ((optionsRef.current.mode ?? "kinematic") !== "kinematic") return;
+    const traj = trajectoryRef.current;
+    if (traj.length === 0) return;
     const now = performance.now();
+    const fps = optionsRef.current.fps ?? 30;
+    const frameInterval = 1e3 / (fps * speedRef.current);
     const elapsed = now - lastFrameTimeRef.current;
-    const frameInterval = 1e3 / fps;
     if (elapsed < frameInterval) return;
     lastFrameTimeRef.current = now;
     const model = mjModelRef.current;
     const data = mjDataRef.current;
     if (!model || !data) return;
-    const qpos = trajectory[frameRef.current];
+    const qpos = getQpos(traj, frameRef.current);
     if (!qpos) return;
     for (let i = 0; i < Math.min(qpos.length, model.nq); i++) {
       data.qpos[i] = qpos[i];
     }
     mujocoRef.current.mj_forward(model, data);
     frameRef.current++;
-    if (frameRef.current >= trajectory.length) {
-      if (loop) {
+    if (frameRef.current >= traj.length) {
+      if (optionsRef.current.loop) {
+        frameRef.current = 0;
+      } else {
+        complete();
+      }
+    }
+  });
+  useBeforePhysicsStep((model, data) => {
+    if (stateRef.current !== "playing") return;
+    if ((optionsRef.current.mode ?? "kinematic") !== "physics") return;
+    const traj = trajectoryRef.current;
+    if (traj.length === 0) return;
+    const fps = optionsRef.current.fps ?? 30;
+    const targetFrame = Math.floor(data.time * fps * speedRef.current);
+    frameRef.current = Math.min(targetFrame, traj.length - 1);
+    const ctrl = getCtrl(traj, frameRef.current);
+    if (ctrl) {
+      for (let i = 0; i < Math.min(ctrl.length, model.nu); i++) {
+        data.ctrl[i] = ctrl[i];
+      }
+    }
+    if (frameRef.current >= traj.length - 1) {
+      if (optionsRef.current.loop) {
+        data.time = 0;
         frameRef.current = 0;
       } else {
-        playingRef.current = false;
-        pausedRef.current = false;
+        complete();
       }
     }
   });
@@ -3160,14 +3233,21 @@ function useTrajectoryPlayer(trajectory, options = {}) {
     pause,
     seek,
     reset,
+    setSpeed,
+    get state() {
+      return stateRef.current;
+    },
     get frame() {
       return frameRef.current;
     },
     get playing() {
-      return playingRef.current;
+      return stateRef.current === "playing";
     },
     get totalFrames() {
       return trajectory.length;
+    },
+    get progress() {
+      return trajectory.length > 1 ? frameRef.current / (trajectory.length - 1) : 0;
     }
   };
 }
@@ -3176,14 +3256,28 @@ function useTrajectoryPlayer(trajectory, options = {}) {
 function TrajectoryPlayer({
   trajectory,
   fps = 30,
+  speed = 1,
   loop = false,
   playing = false,
-  onFrame
+  mode = "kinematic",
+  onFrame,
+  onComplete,
+  onStateChange
 }) {
-  const player = useTrajectoryPlayer(trajectory, { fps, loop });
+  const player = useTrajectoryPlayer(trajectory, {
+    fps,
+    speed,
+    loop,
+    mode,
+    onComplete,
+    onStateChange
+  });
   const onFrameRef = useRef(onFrame);
   onFrameRef.current = onFrame;
   const lastReportedFrameRef = useRef(-1);
+  useEffect(() => {
+    player.setSpeed(speed);
+  }, [speed, player]);
   useEffect(() => {
     if (playing) {
       player.play();