# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------
"""
Gizmo utility functions for EmbodiSim.
This module provides utility functions for creating gizmo transform callbacks.
"""
from typing import Callable
from typing import TYPE_CHECKING
from dexsim.types import TransformMask
if TYPE_CHECKING:
from embodichain.lab.sim.objects import Robot
[docs]
def create_gizmo_callback() -> Callable:
"""Create a standard gizmo transform callback function.
This callback handles local pose for gizmo controls.
It applies transformations directly to the node when gizmo controls are manipulated.
Returns:
Callable: A callback function that can be used with gizmo.node.set_flush_transform_callback()
"""
def gizmo_transform_callback(node, local_pose, flag):
if node is not None:
node.set_transform(local_pose, flag)
return gizmo_transform_callback
[docs]
def run_gizmo_robot_control_loop(
robot: object | str, control_part: str = "arm", end_link_name: str | None = None
):
"""Run a control loop for testing gizmo controls on a robot.
This function implements a control loop that allows users to manipulate a robot
using gizmo controls with keyboard input for additional commands.
Args:
robot (Robot | str): The robot to control with the gizmo.
control_part (str, optional): The part of the robot to control. Defaults to "arm".
end_link_name (str | None, optional): The name of the end link for FK calculations. Defaults to None.
Keyboard Controls:
Q/ESC: Exit the control loop
P: Print current robot state (joint positions, end-effector pose)
G: Toggle gizmo visibility
R: Reset robot to initial pose
I: Print control information
"""
import select
import sys
import tty
import termios
import time
import numpy as np
np.set_printoptions(precision=5, suppress=True)
from embodichain.lab.sim import SimulationManager
from embodichain.lab.sim.objects import Robot
from embodichain.lab.sim.solvers import PinkSolverCfg
from embodichain.utils.logger import log_info, log_warning, log_error
sim = SimulationManager.get_instance()
if isinstance(robot, str):
robot = sim.get_robot(uid=robot)
# Enter auto-update mode.
sim.set_manual_update(False)
# Replace robot's default solver with PinkSolver for gizmo control.
robot_solver = robot.get_solver(name=control_part)
control_part_link_names = robot.get_control_part_link_names(name=control_part)
end_link_name = (
control_part_link_names[-1] if end_link_name is None else end_link_name
)
pink_solver_cfg = PinkSolverCfg(
urdf_path=robot.cfg.fpath,
end_link_name=end_link_name,
root_link_name=robot_solver.root_link_name,
pos_eps=1e-2,
rot_eps=5e-2,
max_iterations=300,
dt=0.1,
)
robot.init_solver(cfg={control_part: pink_solver_cfg})
# Enable gizmo for the robot
gizmo = sim.enable_gizmo(uid=robot.uid, control_part=control_part)
# Store initial robot configuration
initial_qpos = robot.get_qpos(name=control_part)
gizmo_visible = True
log_info("\n=== Gizmo Robot Control ===")
log_info("Gizmo Controls:")
log_info(" Use the 3D gizmo to drag and manipulate the robot")
log_info("\nKeyboard Controls:")
log_info(" Q/ESC: Exit control loop")
log_info(" P: Print current robot state")
log_info(" G: Toggle gizmo visibility")
log_info(" R: Reset robot to initial pose")
log_info(" I: Print this information again")
# Save terminal settings
old_settings = termios.tcgetattr(sys.stdin)
tty.setcbreak(sys.stdin.fileno())
def get_key():
"""Non-blocking keyboard input."""
if select.select([sys.stdin], [], [], 0)[0]:
return sys.stdin.read(1)
return None
try:
while True:
time.sleep(0.033) # ~30Hz
sim.update_gizmos()
# Check for keyboard input
key = get_key()
if key:
# Exit controls
if key in ["q", "Q", "\x1b"]: # Q or ESC
log_info("Exiting gizmo control loop...")
sim.disable_gizmo(uid=robot.uid, control_part=control_part)
if robot_solver:
robot.init_solver(
cfg={control_part: robot_solver.cfg}
) # Restore original solver
break
# Print robot state
elif key in ["p", "P"]:
current_qpos = robot.get_qpos(name=control_part)
eef_pose = robot.compute_fk(name=control_part, qpos=current_qpos)
log_info(f"\n=== Robot State ===")
log_info(f"Control part: {control_part}")
log_info(f"Joint positions: {current_qpos.squeeze().tolist()}")
log_info(f"End-effector pose:\n{eef_pose.squeeze().numpy()}")
if eef_pose is None:
log_info(
"End-effector pose unavailable: compute_fk returned None "
f"for control part '{control_part}'."
)
else:
eef_pose_np = eef_pose.detach().cpu().numpy().squeeze()
log_info(f"End-effector pose:\n{eef_pose_np}")
elif key in ["g", "G"]:
if gizmo_visible:
sim.set_gizmo_visibility(
uid=robot.uid, control_part=control_part, visible=False
)
log_info("Gizmo hidden")
gizmo_visible = False
else:
sim.set_gizmo_visibility(
uid=robot.uid, control_part=control_part, visible=True
)
log_info("Gizmo shown")
gizmo_visible = True
# Reset to initial pose
elif key in ["r", "R"]:
# TODO: Workaround for reset. Gizmo pose should be fixed in the future.
sim.disable_gizmo(uid=robot.uid, control_part=control_part)
robot.clear_dynamics()
robot.set_qpos(qpos=initial_qpos, name=control_part, target=False)
sim.enable_gizmo(uid=robot.uid, control_part=control_part)
log_info("Robot reset to initial pose")
# Print info
elif key in ["i", "I"]:
log_info("\n=== Gizmo Robot Control ===")
log_info("Gizmo Controls:")
log_info(" Use the 3D gizmo to drag and manipulate the robot")
log_info("\nKeyboard Controls:")
log_info(" Q/ESC: Exit control loop")
log_info(" P: Print current robot state")
log_info(" G: Toggle gizmo visibility")
log_info(" R: Reset robot to initial pose")
log_info(" I: Print this information again")
except KeyboardInterrupt:
sim.disable_gizmo(uid=robot.uid, control_part=control_part)
if robot_solver:
robot.init_solver(
cfg={control_part: robot_solver.cfg}
) # Restore original solver
log_info("\nControl loop interrupted by user (Ctrl+C)")
finally:
try:
# Restore terminal settings
termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings)
except:
pass
log_info("Gizmo control loop terminated")