Building Controllers

Custom controllers, IK solvers, and composing library hooks into robot-specific logic

A controller is a React component that calls useBeforePhysicsStep to write data.ctrl each frame and renders null.

The useIkController() hook follows this same pattern. You can use it, swap in your own IK solver, or write your own controller from scratch.

Pattern: Simple Keyboard Bindings

For robots where arm control comes from <IkGizmo />, the controller only adds extra bindings (gripper, etc.):

import { ModelActuators, useKeyboardTeleop } from "mujoco-react";

export function FrankaController() {
  useKeyboardTeleop({
    bindings: {
      v: { actuator: ModelActuators.franka.gripper, toggle: [0, 255] },
    },
  });
  return null;
}

Drop it into your scene as a child of <MujocoCanvas>:

import { IkGizmo, ModelSites, useIkController } from "mujoco-react";

function FrankaScene() {
  const ik = useIkController({ siteName: ModelSites.franka.tcp });
  return (
    <>
      {ik && <IkGizmo controller={ik} />}
      <FrankaController />
    </>
  );
}

<MujocoCanvas config={frankaConfig}>
  <FrankaScene />
</MujocoCanvas>

Pattern: Custom Physics-Step Control

For more complex control (IK solvers, velocity control, state machines), use useBeforePhysicsStep to write directly to data.ctrl each frame.

Keyboard State

Read keyboard input via window event listeners and a ref:

import { useEffect, useRef } from "react";
import { ModelActuators, useBeforePhysicsStep, useCtrl } from "mujoco-react";

function MyController() {
  const keys = useRef<Record<string, boolean>>({});
  const shoulder = useCtrl(ModelActuators.franka.actuator1);

  useEffect(() => {
    const down = (e: KeyboardEvent) => { keys.current[e.code] = true; };
    const up = (e: KeyboardEvent) => { keys.current[e.code] = false; };
    window.addEventListener("keydown", down);
    window.addEventListener("keyup", up);
    return () => {
      window.removeEventListener("keydown", down);
      window.removeEventListener("keyup", up);
    };
  }, []);

  useBeforePhysicsStep(({ data }) => {
    const k = keys.current;
    const next =
      shoulder.read() + (k["KeyW"] ? 0.01 : 0) - (k["KeyS"] ? 0.01 : 0);
    shoulder.write(next);
  });

  return null;
}

Config-Driven Arm Controller

A generic hook that accepts a static config object makes it easy to support multiple robots. Each robot is a different config:

interface ArmConfig {
  actuators: string[];      // Generated actuator names for this arm
  keys: string[];           // Key codes for movement
  initialJoints?: number[]; // Starting joint positions
}

interface ArmControllerConfig {
  numActuators: number;
  arms: ArmConfig[];
  base?: { ... };  // Mobile base drive
  head?: { ... };  // Pan/tilt head
}

The hook reads keyboard state and writes to the correct actuator handles each frame:

function useArmController(config: ArmControllerConfig) {
  const keys = useRef<Record<string, boolean>>({});
  const sim = useMujoco();

  // ... keyboard listeners ...

  useBeforePhysicsStep(({ data }) => {
    if (!sim.isReady) return;
    for (const arm of config.arms) {
      // Read keys, solve IK, then:
      // sim.api.setCtrl(arm.actuators[i], value)
    }
  });
}

Then each robot controller is just a config:

const SO101_CONFIG: ArmControllerConfig = {
  numActuators: 6,
  arms: [{
    actuators: [
      ModelActuators.so101.Rotation,
      ModelActuators.so101.Pitch,
      ModelActuators.so101.Elbow,
      ModelActuators.so101.Wrist_Pitch,
      ModelActuators.so101.Wrist_Roll,
    ],
    keys: ["KeyD", "KeyA", "KeyW", "KeyS", "KeyQ", "KeyE",
           "KeyR", "KeyF", "KeyZ", "KeyC", "KeyV"],
    initialJoints: [0.0158, 2.052, 2.1307, -0.0845, 1.5857, -0.3745],
  }],
};

export function SO101Controller() {
  useArmController(SO101_CONFIG);
  return null;
}

Custom IK Solvers

Three options for IK:

1. Use the built-in solver

The default useIkController() uses Damped Least-Squares:

import { IkGizmo, ModelSites, useIkController } from "mujoco-react";

const ik = useIkController({ siteName: ModelSites.franka.tcp });
return ik ? <IkGizmo controller={ik} /> : null;

2. Plug in your own solver

Pass ikSolveFn to replace the built-in solver while keeping the gizmo, reset handling, and context:

import { IkGizmo, ModelSites, useIkController } from "mujoco-react";
import type { IKSolveFn } from "mujoco-react";

const myIK: IKSolveFn = ({ position, currentQ }) => {
  return myAnalyticalSolver(position, currentQ); // return joint angles or null
};

const ik = useIkController({ siteName: ModelSites.franka.tcp, ikSolveFn: myIK });
return ik ? <IkGizmo controller={ik} /> : null;

3. Skip useIkController entirely

Solve IK yourself inside useBeforePhysicsStep with full access to the model and data:

function MyIKController({ targetRef }) {
  const sim = useMujoco();

  useBeforePhysicsStep(({ model, data }) => {
    const target = targetRef.current;
    if (!target || !sim.isReady) return;

    const joints = myCustomIKSolve(model, data, target);
    if (joints) {
      sim.api.setCtrl({
        [ModelActuators.franka.actuator1]: joints[0],
        [ModelActuators.franka.actuator2]: joints[1],
        [ModelActuators.franka.actuator3]: joints[2],
        [ModelActuators.franka.actuator4]: joints[3],
        [ModelActuators.franka.actuator5]: joints[4],
        [ModelActuators.franka.actuator6]: joints[5],
        [ModelActuators.franka.actuator7]: joints[6],
      });
    }
  });
  return null;
}

This gives you full access to model/data for whatever solver you want.

Pattern: Reusable Plugins with createControllerHook

For reusable controllers with typed config and default merging, use the createControllerHook factory. It stabilizes config references (so inline objects don’t cause re-renders), merges defaults, and supports disabling via null.

import { createControllerHook, ModelActuators, useBeforePhysicsStep, useCtrl } from "mujoco-react";

interface MyConfig {
  gain: number;
  actuator: string;
  frequency?: number;
}

export const useMyController = createControllerHook<MyConfig, { amplitude: number }>(
  { name: "useMyController", defaultConfig: { gain: 1.0, frequency: 1.0 } },
  (config) => {
    const amplitudeRef = useRef(0);
    const actuator = useCtrl(config?.actuator ?? ModelActuators.franka.actuator1);

    useBeforePhysicsStep(({ data }) => {
      if (!config) return;
      const freq = config.frequency ?? 1.0;
      amplitudeRef.current = config.gain * Math.sin(data.time * freq);
      actuator.write(amplitudeRef.current);
    });

    if (!config) return null;
    return { amplitude: amplitudeRef.current };
  },
);

// const result = useMyController({ gain: 2.0, actuator: ModelActuators.franka.actuator1 });
// const disabled = useMyController(null); // returns null, no-ops

createControllerHook API

function createControllerHook<TConfig, TValue>(
  options: { name: string; defaultConfig?: Partial<TConfig> },
  useImpl: (config: TConfig | null) => TValue | null,
): (config: TConfig | null) => TValue | null;

Pass null to disable the controller without breaking the rules of hooks — useImpl is always called, it just receives null and should no-op.

Pattern: Reusable Plugins with createController

The createController factory is the component equivalent — same config stabilization and default merging, but returns a component that can render children:

import { createController, ModelActuators, useBeforePhysicsStep, useCtrl } from "mujoco-react";

interface MyConfig {
  gain: number;
  actuator: string;
  frequency?: number;
}

function MyControllerImpl({ config, children }: { config: MyConfig; children?: React.ReactNode }) {
  const actuator = useCtrl(config.actuator);

  useBeforePhysicsStep(({ data }) => {
    const freq = config.frequency ?? 1.0;
    actuator.write(config.gain * Math.sin(data.time * freq));
  });
  return <>{children}</>;
}

export const MyController = createController<MyConfig>(
  { name: "MyController", defaultConfig: { gain: 1.0, frequency: 1.0 } },
  MyControllerImpl,
);

// <MyController config={{ gain: 2.0, actuator: ModelActuators.franka.actuator1 }}>
//   <Debug showJoints />
// </MyController>

createController API

function createController<TConfig>(
  options: { name: string; defaultConfig?: Partial<TConfig> },
  Impl: React.FC<{ config: TConfig; children?: React.ReactNode }>,
): ControllerComponent<TConfig>;

The returned component accepts config (merged with defaults) and optional children. It also exposes static metadata: MyController.controllerName and MyController.defaultConfig.

Providing Context to Children

Controllers can provide state to descendants via React context:

const MyContext = createContext<MyContextValue | null>(null);

function MyControllerImpl({ config, children }) {
  const value = useMemo(() => ({ /* state + methods */ }), []);
  return <MyContext.Provider value={value}>{children}</MyContext.Provider>;
}

Listening for Resets

Register a callback to reset your controller state when the simulation resets:

function MyControllerImpl({ config, children }) {
  const { resetCallbacks } = useMujoco();

  useEffect(() => {
    const cb = () => { /* reset your state */ };
    resetCallbacks.current.add(cb);
    return () => { resetCallbacks.current.delete(cb); };
  }, [resetCallbacks]);

  return <>{children}</>;
}

The library’s useIkController() hook demonstrates all these patterns: reset handling, useBeforePhysicsStep for solving, and useFrame for gizmo animation.

Coexisting with IK Gizmo

When a robot supports both gizmo drag and keyboard control, the controller needs to:

1. Accept ik as a prop (the IkContextValue from useIkController())
2. Sync state on transition: when the user switches from gizmo to keyboard, read data.ctrl to avoid a position jump
3. Disable IK via ik.setIkEnabled(false) when taking over

import { useBeforePhysicsStep } from "mujoco-react";
import type { IkContextValue } from "mujoco-react";

function MyArmController({ ik }: { ik?: IkContextValue | null }) {
  useBeforePhysicsStep(({ data }) => {
    const anyKeyPressed = /* check keyboard state */;

    if (anyKeyPressed && ik?.ikEnabledRef.current) {
      // Sync from current gizmo position
      for (let i = 0; i < arm.indices.length; i++) {
        targetJoints[i] = data.ctrl[arm.indices[i]];
      }
      ik.setIkEnabled(false);
    }

    if (!ik?.ikEnabledRef.current) {
      // Keyboard is in control, solve IK and write ctrl
      for (let i = 0; i < arm.indices.length; i++) {
        data.ctrl[arm.indices[i]] = targetJoints[i];
      }
    }
  });
}

Pass the ik value from useIkController() to the controller as a prop.

The gizmo re-enables IK automatically when dragged.

Composing Controllers in Your Scene

Controllers are React children. Swap them based on state:

function SceneChildren({ modelKey, ikConfig, showGizmo }) {
  const ik = useIkController(ikConfig);

  return (
    <>
      {ik && showGizmo && <IkGizmo controller={ik} />}
      <DragInteraction />

      {modelKey === "franka" && <FrankaController />}
      {modelKey === "so101" && <SO101Controller ik={ik} />}
      {modelKey === "xlerobot" && <XLeRobotController ik={ik} />}
    </>
  );
}

<MujocoCanvas config={entry.config}>
  <SceneChildren modelKey={modelKey} ikConfig={ikConfig} showGizmo={showGizmo} />
</MujocoCanvas>

Performance Tips

• Use for loops instead of .forEach / .map in useBeforePhysicsStep (it runs every physics tick)
• Store keyboard state in a useRef, not useState (avoids re-renders at 60fps)
• Cache actuator IDs once (not every frame) using findActuatorByName
• Keep the callback closure stable; avoid creating new functions each render