YikaiFu-cart/acRealman_xr
Template
1
1
Fork 1
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add reset ability in mujoco

Browse Source
This commit is contained in:
YikaiFu-cart
2026-07-02 14:02:44 +08:00
parent 5678cc8e69
commit 0d7e30cb7b
7 changed files with 176 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
ik_qp/src/rm75_ik/mujoco_robot.py
View File

@ -1,8 +1,9 @@
from __future__ import annotations
from contextlib import nullcontext
from dataclasses import dataclass
from time import sleep
from typing import Collection, Dict, Mapping, Optional
from typing import Callable, Collection, Dict, Mapping, Optional
import numpy as np
import pinocchio as pin
@ -47,6 +48,7 @@ class MujocoRobot:
self._viewer = None
self._closed = False
self._stopped_arms: set[ArmName] = set()
self._initial_joint_positions: Optional[Dict[ArmName, np.ndarray]] = None
self.initial_pose_diagnostics: Dict[ArmName, InitialPoseDiagnostic] = {}
if initialize_from_tcp:
self._initialize_from_configured_tcp()
@ -93,6 +95,9 @@ class MujocoRobot:
configured_joint_orientation_error_rad=configured_error[1],
)
self.scene.set_dual_configuration(solved["left"], solved["right"])
self._initial_joint_positions = {
arm: solved[arm].copy() for arm in ("left", "right")
}
for arm in ("left", "right"):
self.set_target_tcp_pose(arm, self.profiles[arm].initial_tcp)
@ -140,6 +145,23 @@ class MujocoRobot:
q, self.profiles[arm].tool_from_flange
)
def reset_to_initial(self) -> None:
"""Restore the solved startup state without recreating the MuJoCo model."""
self._require_open()
if self._initial_joint_positions is None:
raise RuntimeError("MuJoCo initial joint positions are unavailable")
viewer_lock = self._viewer.lock() if self._viewer is not None else nullcontext()
with viewer_lock:
self.scene.set_dual_configuration(
self._initial_joint_positions["left"],
self._initial_joint_positions["right"],
)
self.scene.data.qvel[:] = 0.0
for arm in ("left", "right"):
self.set_target_tcp_pose(arm, self.profiles[arm].initial_tcp)
self._stopped_arms.update(("left", "right"))
def stop(self, arms: Collection[ArmName] = ("left", "right")) -> None:
self._require_open()
selected = set(arms)
@ -148,13 +170,23 @@ class MujocoRobot:
self.scene.data.qvel[:] = 0.0
self._stopped_arms.update(selected)
def open_viewer(self):
def open_viewer(self, key_callback: Optional[Callable[[int], None]] = None):
self._require_open()
if self._viewer is None:
import mujoco.viewer
self._viewer = mujoco.viewer.launch_passive(
self.scene.model, self.scene.data
self.scene.model,
self.scene.data,
key_callback=key_callback,
)
self._viewer.set_texts(
(
None,
mujoco.mjtGridPos.mjGRID_BOTTOMLEFT,
"Reset dual arms",
"R / Home",
)
)
return self._viewer

2
ik_qp/src/rm75_ik/robot_backend.py
View File

@ -38,6 +38,8 @@ class RobotBackend(Protocol):
def set_target_tcp_pose(self, arm: ArmName, target_tcp: pin.SE3) -> None: ...
def reset_to_initial(self) -> None: ...
def stop(self, arms: Collection[ArmName] = ("left", "right")) -> None: ...
def close(self) -> None: ...

49
ik_qp/src/rm75_ik/teleop_control.py
View File

@ -17,6 +17,7 @@ from .types import IkOptions, teleop_joint_limits
class SafetyState(str, Enum):
IDLE = "idle"
ACTIVE = "active"
RESETTING = "resetting"
FAULT = "fault"
@ -297,6 +298,7 @@ class DualArmQpTeleopController:
}
self._latest: Dict[ArmName, tuple[ControllerSample, float]] = {}
self._fault_reason = ""
self._reset_waiting_for_release = False
self._closed = False
self.robot.connect()
@ -304,6 +306,8 @@ class DualArmQpTeleopController:
def safety_state(self) -> SafetyState:
if self._fault_reason:
return SafetyState.FAULT
if self._reset_waiting_for_release:
return SafetyState.RESETTING
if any(mapper.active for mapper in self.mappers.values()):
return SafetyState.ACTIVE
return SafetyState.IDLE
@ -326,6 +330,27 @@ class DualArmQpTeleopController:
def reject_input(self, reason: str) -> None:
self._trip_fault(f"invalid controller input: {reason}")
def reset_to_initial(self) -> ControlCycleResult:
"""Reset the backend and require fresh dual-grip release before rearming."""
if self._closed:
raise RuntimeError("controller is closed")
for mapper in self.mappers.values():
mapper.reset()
try:
self.robot.stop(("left", "right"))
self.robot.reset_to_initial()
except Exception as exc:
self._trip_fault(f"reset/backend failure: {exc}")
raise RuntimeError(f"failed to reset robot backend: {exc}") from exc
self._latest.clear()
self._fault_reason = ""
self._reset_waiting_for_release = True
return ControlCycleResult(
SafetyState.RESETTING,
reason="reset complete; waiting for fresh dual-grip release",
)
def step(self, timestamp_sec: float) -> ControlCycleResult:
if self._closed:
raise RuntimeError("controller is closed")
@ -336,6 +361,15 @@ class DualArmQpTeleopController:
self._fault_reason = ""
return ControlCycleResult(SafetyState.IDLE, reason="fault cleared")
return ControlCycleResult(SafetyState.FAULT, reason=self._fault_reason)
if self._reset_waiting_for_release:
if self._can_finish_reset(timestamp_sec):
self._reset_waiting_for_release = False
self._latest.clear()
return ControlCycleResult(SafetyState.IDLE, reason="reset rearmed")
return ControlCycleResult(
SafetyState.RESETTING,
reason="waiting for fresh dual-grip release",
)
try:
state = self.robot.read_joint_positions()
@ -404,8 +438,22 @@ class DualArmQpTeleopController:
mapper.reset()
return True
def _can_finish_reset(self, timestamp_sec: float) -> bool:
if set(self._latest) != {"left", "right"}:
return False
for arm, (sample, sample_time) in self._latest.items():
if sample.grip:
return False
age = timestamp_sec - sample_time
if age < -1e-6 or age > self.profiles[arm].command_timeout_sec:
return False
for mapper in self.mappers.values():
mapper.reset()
return True
def _trip_fault(self, reason: str) -> ControlCycleResult:
self._fault_reason = str(reason)
self._reset_waiting_for_release = False
for mapper in self.mappers.values():
mapper.reset()
try:
@ -415,6 +463,7 @@ class DualArmQpTeleopController:
return ControlCycleResult(SafetyState.FAULT, reason=self._fault_reason)
def stop(self) -> None:
self._reset_waiting_for_release = False
for mapper in self.mappers.values():
mapper.reset()
self.robot.stop(("left", "right"))