Teleporting free joint bodies in Unity

[prnthp]

I've found only crumbs of information on how to "teleport" objects in MuJoCo and I'd thought I should share a snippet of how this can be done in the Unity plug-in.

public unsafe void TeleportBody(MjFreeJoint freeJoint, Vector3 position, Quaternion rotation)
{
    var data = MjScene.Instance.Data;

    var qposAddress = freeJoint.QposAddress;
    var dofAddress = freeJoint.DofAddress;
    
    MjEngineTool.SetMjTransform(
        data->qpos + qposAddress,
        position,
        rotation);
        
    // This is just for resetting velocity to zero
    MjEngineTool.SetMjVector3(
        data->qvel + dofAddress,
        Vector3.zero);
    MjEngineTool.SetMjVector3(
        data->qvel + dofAddress + 3,
        Vector3.zero);
}

[Balint-H]

Be careful as SetMjTransform sets 7 values total, while free joints only have 6 component velocities, so SetMjTransform on qvel will write into the next body's data! You can create your own function that takes two Vector3 inputs, or call SetMjVector3 twice.

In case someone also needs to preserve momentum, you can store the local velocity velocity and restore it on teleport. This script restores momentum based on the local frame of reference (e.g. if a flying arrow is teleported and rotated, it will keep going forwards in the new direction, as opposed to flying sideways in the old one):

public static unsafe void TeleportMjRoot(MjFreeJoint root, Vector3 unityPos, Quaternion unityRot)
{
    MujocoLib.mjData_* Data = mjScene.Data;

    Quaternion fromUnityRotation = MjEngineTool.UnityQuaternion(Data->qpos + root.DofAddress + 3); // Store original rotation

    Quaternion rotationOffset = unityRot * Quaternion.Inverse(fromUnityRotation); // Get difference to new one (we'll transform velocity with this)            
    Vector3 newUnityLinVel = rotationOffset * MjEngineTool.UnityVector3(Data->qvel + root.DofAddress);
    Vector3 newUnityAngVel = rotationOffset * MjEngineTool.UnityVector3(Data->qvel + root.DofAddress+3);

    MjEngineTool.SetMjTransform(Data->qpos + root.QposAddress, unityPos, unityRot); //Set position state
    MjEngineTool.SetMjVector3(Data->qvel + root.DofAddress, newUnityLinVel); // Set linear velocity
    MjEngineTool.SetMjVector3(Data->qvel + root.DofAddress + 3, newUnityAngVel); //Set angular velocity

    //Could update state potentially
    //MujocoLib.mjModel_* Model = mjScene.Model;
    //MujocoLib.mj_kinematics(Model, Data);
    //mjScene.SyncUnityToMjState();
   
}

[prnthp]

Thank you! This is great!

I'll update my post.