Orbit DifferentialIKController for velocity control

Hello, I’m trying to use the Orbit framework to do velocity control on a Franka Panda arm end-effector. I saw this differential IK controller from the following tutorial, but the example isn’t using the provided relative mode: Using a task-space controller — orbit documentation

I tried to use the relative mode so that I could send velocity commands to the controller, instead of desired pose, and avoid dealing with quaternions. I started with adapting the example code and sending zero delta change of ee pose, but it shows in the simulation that the end-effector keeps moving although the zero delta change command was constantly sent.
Could someone maybe check my code and tell me where I got mistakes? Thank you!

Ps. As I could not find it, are there functions for the Franks arm to compute its forward and inverse kinematics, instead of a controller like DifferentialIKController? Like directly getting ee pose from the joint configuration and directly getting joint positions from the desired pose of end-effector.

from __future__ import annotations

"""Launch Isaac Sim Simulator first."""

import argparse

from omni.isaac.orbit.app import AppLauncher

# add argparse arguments
parser = argparse.ArgumentParser(description="Tutorial on using the differential IK controller.")
parser.add_argument("--robot", type=str, default="franka_panda", help="Name of the robot.")
parser.add_argument("--num_envs", type=int, default=1, help="Number of environments to spawn.")
# append AppLauncher cli args
AppLauncher.add_app_launcher_args(parser)
# parse the arguments
args_cli = parser.parse_args()

# launch omniverse app
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app

"""Rest everything follows."""

import torch
import traceback

import carb

import omni.isaac.orbit.sim as sim_utils
from omni.isaac.orbit.assets import AssetBaseCfg
from omni.isaac.orbit.controllers import DifferentialIKController, DifferentialIKControllerCfg
from omni.isaac.orbit.managers import SceneEntityCfg
from omni.isaac.orbit.markers import VisualizationMarkers
from omni.isaac.orbit.markers.config import FRAME_MARKER_CFG
from omni.isaac.orbit.scene import InteractiveScene, InteractiveSceneCfg
from omni.isaac.orbit.utils import configclass
from omni.isaac.orbit.utils.assets import ISAAC_NUCLEUS_DIR
from omni.isaac.orbit.utils.math import subtract_frame_transforms

##
# Pre-defined configs
##
from omni.isaac.orbit_assets import FRANKA_PANDA_HIGH_PD_CFG, UR10_CFG  # isort:skip


@configclass
class TableTopSceneCfg(InteractiveSceneCfg):
    """Configuration for a cart-pole scene."""

    # ground plane
    ground = AssetBaseCfg(
        prim_path="/World/defaultGroundPlane",
        spawn=sim_utils.GroundPlaneCfg(),
        init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, -1.05)),
    )

    # lights
    dome_light = AssetBaseCfg(
        prim_path="/World/Light", spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75))
    )

    # mount
    table = AssetBaseCfg(
        prim_path="{ENV_REGEX_NS}/Table",
        spawn=sim_utils.UsdFileCfg(
            usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/Stand/stand_instanceable.usd", scale=(5.0, 5.0, 2.0)
        ),
        # init_state=AssetBaseCfg.InitialStateCfg(pos=(2.0, 0.0, 0.0)),
    )

    # articulation
    if args_cli.robot == "franka_panda":
        robot = FRANKA_PANDA_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
    elif args_cli.robot == "ur10":
        robot = UR10_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
    else:
        raise ValueError(f"Robot {args_cli.robot} is not supported. Valid: franka_panda, ur10")
    

# get robot pose in robot frame
def ee_pose_in_robot_frame(scene, robot, robot_entity_cfg) -> torch.Tensor:
    ee_pos_w, ee_quat_w = ee_pose_in_world_frame(scene, robot, robot_entity_cfg)
    root_pose_w = root_pose_in_world_frame(scene, robot, robot_entity_cfg)
    # joint_pos = robot.data.joint_pos[:, robot_entity_cfg.joint_ids]
    ee_pos_r, ee_quat_r = subtract_frame_transforms(
                    root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pos_w, ee_quat_w
                )
    return ee_pos_r, ee_quat_r

# get robot pose in world frame from simulation buffer
def ee_pose_in_world_frame(scene, robot, robot_entity_cfg) -> torch.Tensor:
    ee_pose_w = robot.data.body_state_w[:, robot_entity_cfg.body_ids[0], 0:7]
    ee_pos_w = ee_pose_w[:, 0:3]
    ee_quat_w = ee_pose_w[:, 3:7]
    # root_pose_w = robot.data.root_state_w[:, 0:7]
    return ee_pos_w, ee_quat_w

# get root pose in world frame from simulation buffer
def root_pose_in_world_frame(scene, robot, robot_entity_cfg) -> torch.Tensor:
    root_pose_w = robot.data.root_state_w[:, 0:7]
    return root_pose_w

def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene):
    """Runs the simulation loop."""
    # Extract scene entities
    # note: we only do this here for readability.
    robot = scene["robot"]

    # Create a differential IK controller
    diff_ik_cfg = DifferentialIKControllerCfg(command_type="pose", use_relative_mode=True, ik_method="dls")
    diff_ik_controller = DifferentialIKController(diff_ik_cfg, num_envs=scene.num_envs, device=sim.device)

    # Create a marker for the end-effector
    frame_marker_cfg = FRAME_MARKER_CFG.copy()
    frame_marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1)
    ee_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_current"))
    goal_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_goal"))

    # Specify robot-specific parameters
    if args_cli.robot == "franka_panda":
        robot_entity_cfg = SceneEntityCfg("robot", joint_names=["panda_joint.*"], body_names=["panda_hand"])
    elif args_cli.robot == "ur10":
        robot_entity_cfg = SceneEntityCfg("robot", joint_names=[".*"], body_names=["ee_link"])
    else:
        raise ValueError(f"Robot {args_cli.robot} is not supported. Valid: franka_panda, ur10") 
    
    # Resolving the scene entities
    robot_entity_cfg.resolve(scene)
    # Obtain the frame index of the end-effector
    # For a fixed base robot, the frame index is one less than the body index. This is because
    # the root body is not included in the returned Jacobians.
    if robot.is_fixed_base:
        ee_jacobi_idx = robot_entity_cfg.body_ids[0] - 1
    else:
        ee_jacobi_idx = robot_entity_cfg.body_ids[0]

    ik_commands = torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]], device=sim.device)

    # Define simulation stepping
    sim_dt = sim.get_physics_dt()
    count = 0
    # Simulation loop
    while simulation_app.is_running():
        # reset
        if count % 150 == 0:
            # reset time
            count = 0
            # reset joint state
            joint_pos_default = robot.data.default_joint_pos.clone()
            joint_vel_default = robot.data.default_joint_vel.clone()
            robot.write_joint_state_to_sim(joint_pos_default, joint_vel_default)
            robot.reset()
            # reset actions
            # ik_commands = ee_vel_goal_zero
            joint_pos_des = joint_pos_default[:, robot_entity_cfg.joint_ids].clone()
            joint_vel_des = joint_vel_default[:, robot_entity_cfg.joint_ids].clone()
            # joint_vel_des = torch.zeros((scene.num_envs, 7), device=sim.device)
            joint_pos = robot.data.joint_pos[:, robot_entity_cfg.joint_ids]
            ee_pos_b, ee_quat_b = ee_pose_in_world_frame(scene, robot, robot_entity_cfg)
            # reset controller
            diff_ik_controller.reset()
            diff_ik_controller.set_command(ik_commands, ee_pos= ee_pos_b, ee_quat = ee_quat_b) # delta change of pose
            
        else: 
            # obtain quantities from simulation
            jacobian = robot.root_physx_view.get_jacobians()[:, ee_jacobi_idx, :, robot_entity_cfg.joint_ids]
            ee_pos_b, ee_quat_b = ee_pose_in_world_frame(scene, robot, robot_entity_cfg)
            diff_ik_controller.set_command(ik_commands, ee_pos= ee_pos_b, ee_quat = ee_quat_b) # delta change of pose
            # compute the joint commands
            joint_pos = robot.data.joint_pos[:, robot_entity_cfg.joint_ids]
            joint_pos_des = diff_ik_controller.compute(ee_pos_b, ee_quat_b, jacobian, joint_pos)
            joint_vel_des = (joint_pos_des - joint_pos) / sim_dt
            
        # apply actions
        robot.set_joint_velocity_target(joint_vel_des, joint_ids=robot_entity_cfg.joint_ids, env_ids = list(range(scene.num_envs)))
        scene.write_data_to_sim()
        # perform step
        sim.step()
        # update sim-time
        count += 1
        # update buffers
        scene.update(sim_dt)
        # obtain quantities from simulation
        ee_pose_w = robot.data.body_state_w[:, robot_entity_cfg.body_ids[0], 0:7]
        # update marker positions
        ee_marker.visualize(ee_pose_w[:, 0:3], ee_pose_w[:, 3:7])
        # goal_marker.visualize(ee_pos_ref, ee_quat_ref)

def main(): 
    # Load kit helper
    sim_cfg = sim_utils.SimulationCfg(dt=0.01)
    sim = sim_utils.SimulationContext(sim_cfg)
    # Set main camera
    sim.set_camera_view([2.5, 2.5, 2.5], [0.0, 0.0, 0.0])
    # Design scene
    scene_cfg = TableTopSceneCfg(num_envs=args_cli.num_envs, env_spacing=2.0)
    scene = InteractiveScene(scene_cfg)

    robot = scene["robot"]
    robot.data

    # Play the simulator
    sim.reset()
    # Now we are ready!
    print("[INFO]: Setup complete...")
    # Run the simulator
    run_simulator(sim, scene)

if __name__ == "__main__":
    try:
        # run the main execution
        main()
    except Exception as err:
        carb.log_error(err)
        carb.log_error(traceback.format_exc())
        raise
    finally:
        #close sim app
        simulation_app.close()