Source code for go2.core.controller
import os
import signal
import threading
from typing import Callable, List, Dict
from .module import DogModule
from .registry import ModuleRegistry, ModuleType, ExecutionMode
from ..modules.input import InputSignal
from ..modules.audio import AudioModule
from ..modules.input import InputModule
from ..modules.movement import MovementModule
from ..modules.ocr import OCRModule
from ..modules.video import VideoModule
from ..modules.lidar import LIDARModule
from ..hardware.hardware_type import HardwareType
from ..hardware.hardware_interface_movement import HardwareInterfaceMovement
from ..hardware.native.native_hardware_movement import NativeHardwareMovement
from ..hardware.virtual.virtual_hardware_movement import VirtualHardwareMovement
from ..logging import get_logger
logger = get_logger(__name__)
# TTS: https://medium.com/@vndee.huynh/build-your-own-voice-assistant-and-run-it-locally-whisper-ollama-bark-c80e6f815cba
# Digging into Dog: https://www.darknavy.org/darknavy_insight/the_jailbroken_unitree_robot_dog
# Some very interesting turtorial with the G02: https://hackmd.io/@c12hQ00ySVi6JYIERU7bCg/ByAOr12qJg
# Update to JetPack 6.x on Dog's Jetson: https://theroboverse.com/unitree-go2-edu-jetpack-6-2-1-update/
# TETHERING + ROS2 SETUP: https://www.youtube.com/watch?v=Q_dqPLJDPms
[docs]
class Go2Controller:
"""
Primary control interface for the Unitree Go2 robot.
This class is the **main entry point** that users interact with.
It manages hardware initialization, module lifecycles, safety checks, and shutdown.
Modules are accessed via properties rather than direct instantiation.
Notes
-----
- All hardware access is routed through this controller.
- Modules creation, initialization, and shutdown is handled automatically.
- Supports both SDK-backed hardware and a Mujoco simulation.
See Also
--------
DogModule
ModuleRegistry
"""
def __init__(self, hardware_type: HardwareType, execution_mode: ExecutionMode=ExecutionMode.BASIC) -> None:
"""
Create a new controller instance.
Parameters
----------
use_sdk : bool or None
If ``True``, forces use of the Unitree SDK.
If ``False``, forces simulation mode.
Raises
------
RuntimeError
If hardware backend initialization fails.
"""
self._hardware_type = hardware_type
self._execution_mode = execution_mode
self._emit_execution_mode_info(execution_mode)
self._shutdown_event = threading.Event()
self._shutdown_lock = threading.Lock()
self._cleanup_callbacks_pre_module_shutdown: List[Callable[[], None]] = []
self._cleanup_callbacks_post_module_shutdown: List[Callable[[], None]] = []
self._install_signal_handlers()
self._hardware: HardwareInterfaceMovement = (
NativeHardwareMovement() if hardware_type == HardwareType.NATIVE else VirtualHardwareMovement()
)
self._hardware._initialize()
self._modules: Dict[ModuleType, DogModule] = {}
self._register_default_modules()
self._initialize_input_bindings()
logger.info(f"[Controller] Initialized in {'NATIVE' if hardware_type == HardwareType.NATIVE else 'SIMULATION'} mode\n")
def _emit_execution_mode_info(self, exec_mode: ExecutionMode):
logger.info(f"[Controller] Executing in {'BASIC' if exec_mode == ExecutionMode.BASIC else 'ADVANCED'} execution mode.")
if exec_mode == ExecutionMode.BASIC:
logger.warning("[Controller] LIDAR functionalities not available in 'BASIC' execution")
# Automatic shutdown on exception incase its not done by users
def _install_signal_handlers(self) -> None:
def handler(signum, frame):
self.safe_shutdown()
signal.signal(signum, signal.SIG_DFL) # Hands control back to default handler
os.kill(os.getpid(), signum)
signal.signal(signal.SIGINT, handler)
signal.signal(signal.SIGTERM, handler)
def _initialize_input_bindings(self) -> None:
if self._hardware_type == HardwareType.NATIVE:
self.input.register_callback(
InputSignal.BUTTON_A,
lambda _: self._shutdown_event.set(),
"emergency_stop"
)
def _register_default_modules(self) -> None:
self.add_module(ModuleType.MOVEMENT, hardware=self._hardware)
if self._hardware_type == HardwareType.NATIVE:
self.add_module(ModuleType.INPUT)
[docs]
def add_module(self, module_type: ModuleType, **kwargs) -> None:
"""
Add a module to the controller. Modules are initialized immediately upon addition.
Manual module initialization is discouraged and should not be attempted.
Duplicate module addition is prevented explicitly.
Modules are identified using :class:`ModuleType` enums.
Parameters
----------
module_type : ModuleType
The type of module to add.
**kwargs
Constructor arguments passed to the module implementation.
Raises
------
ValueError
If the module type is not registered.
ValueError
If the module type depends on native hardware support, but the controller is not set up to provide it.
"""
if module_type in self._modules:
logger.info(f"[Controller] Module type '{module_type.name}' is already loaded. Returning existing instance.")
return
descriptor = ModuleRegistry.get_descriptor(module_type)
if descriptor is None:
raise ValueError(f"[Controller] Module type '{module_type.name}' is not registered")
if descriptor._requires_native_hardware and self._hardware_type != HardwareType.NATIVE:
raise ValueError(f"[Controller] Module type '{module_type.name}' requires native hardware support")
if descriptor._requires_advanced_execution and self._execution_mode != ExecutionMode.ADVANCED:
raise ValueError(
f"[Controller] Module type '{module_type.name}' requires advanced execution mode. "
"You may need to install ROS2 Humble libraries.\n"
"[Controller] For more information see here: https://go2-control.readthedocs.io/en/latest/api/core.html#go2.core.registry.ExecutionMode"
)
module: DogModule = descriptor._create_instance(**kwargs)
self._modules[module_type] = module
module._set_hardware_internal(hardware_type=self._hardware_type)
module._initialize()
[docs]
def has_module(self, module_type: ModuleType) -> bool:
"""Check whether a module is currently loaded."""
return module_type in self._modules
@property
def video(self) -> VideoModule:
"""
Access the video capture module.
Raises
------
RuntimeError
If the module is not loaded or shutdown has been requested.
"""
module = self._modules.get(ModuleType.VIDEO)
if not isinstance(module, VideoModule):
raise RuntimeError("Video module not loaded")
if self.is_shutdown_requested():
raise RuntimeError("Cannot access Video module after shutdown has been requested")
return module
@property
def movement(self) -> MovementModule:
"""
Access the movement control module.
Raises
------
RuntimeError
If the module is not loaded or shutdown has been requested.
"""
module = self._modules.get(ModuleType.MOVEMENT)
if not isinstance(module, MovementModule):
raise RuntimeError("Movement module not loaded")
if self.is_shutdown_requested():
raise RuntimeError("Cannot access Movement module after shutdown has been requested")
return module
@property
def ocr(self) -> OCRModule:
"""
Access the ocr control module.
Raises
------
RuntimeError
If the module is not loaded or shutdown has been requested.
"""
module = self._modules.get(ModuleType.OCR)
if not isinstance(module, OCRModule):
raise RuntimeError("OCR module not loaded")
if self.is_shutdown_requested():
raise RuntimeError("Cannot access OCR module after shutdown has been requested")
return module
@property
def audio(self) -> AudioModule:
"""
Access the audio control module.
Raises
------
RuntimeError
If the module is not loaded or shutdown has been requested.
"""
module = self._modules.get(ModuleType.AUDIO)
if not isinstance(module, AudioModule):
raise RuntimeError("Audio module not loaded")
if self.is_shutdown_requested():
raise RuntimeError("Cannot access Audio module after shutdown has been requested")
return module
@property
def input(self) -> InputModule:
"""
Access the input control module.
Raises
------
RuntimeError
If the module is not loaded or shutdown has been requested.
"""
module = self._modules.get(ModuleType.INPUT)
if not isinstance(module, InputModule):
raise RuntimeError("Input module not loaded")
if self.is_shutdown_requested():
raise RuntimeError("Cannot access Input module after shutdown has been requested")
return module
@property
def lidar(self) -> LIDARModule:
"""
Access the lidar control module.
Raises
------
RuntimeError
If the module is not loaded or shutdown has been requested.
"""
module = self._modules.get(ModuleType.LIDAR)
if not isinstance(module, LIDARModule):
raise RuntimeError("LIDAR module not loaded")
if self.is_shutdown_requested():
raise RuntimeError("Cannot access LIDAR module after shutdown has been requested")
return module
[docs]
def register_cleanup_callback_pre_module_shutdown(self, callback: Callable[[], None]):
"""
Register a cleanup callback to be executed during safe shutdown before modules are stopped and before hardware is released.
Parameters
----------
callback : callable
Zero-argument function to execute during shutdown.
"""
self._cleanup_callbacks_pre_module_shutdown.append(callback)
[docs]
def register_cleanup_callback_post_module_shutdown(self, callback: Callable[[], None]):
"""
Register a cleanup callback to be executed during safe shutdown after modules are stopped, but before hardware is released.
Parameters
----------
callback : callable
Zero-argument function to execute during shutdown.
"""
self._cleanup_callbacks_post_module_shutdown.append(callback)
[docs]
def safe_shutdown(self):
"""
Perform a coordinated and safe shutdown.
This method:
- Stops movement immediately
- Signals shutdown to all subsystems
- Shuts down all modules
- Executes registered cleanup callbacks
- Releases hardware resources
Important
---------
Select modules depend on ROS2 nodes. These lifetime of these nodes must handled cleanly on process termination.
Therefore, it is **critical** that students **ALWAYS** call :meth:`safe_shutdown` to ensure proper shutdown of ROS2.
This called automatically on process crash (via `SIGINT` or `SIGTERM`). However, it is **fully** the
user's responsibility to call this method upon normal (error free) script exit.
Failing to follow this can (**and probably will**) result in zombie ROS2 processes, unreleased resources, and UB.
Notes
-----
- This method is **idempotent** and may be safely called multiple times.
"""
logger.info("\n[Controller] Starting safe shutdown...")
with self._shutdown_lock:
if self._shutdown_event.is_set():
return
self._shutdown_event.set()
for callback in self._cleanup_callbacks_pre_module_shutdown:
try:
callback()
except Exception:
logger.exception("[Controller] Cleanup callback failed")
for module_type, module in self._modules.items():
try:
module._shutdown()
except Exception:
logger.exception(f"[Controller] Failed to shutdown {module_type.name}")
for callback in self._cleanup_callbacks_post_module_shutdown:
try:
callback()
except Exception:
logger.exception("[Controller] Cleanup callback failed")
try:
self._hardware._shutdown()
except Exception:
logger.exception("[Controller] Hardware shutdown failed")
logger.info("[Controller] Shutdown complete")
[docs]
def is_shutdown_requested(self) -> bool:
"""Check if shutdown has been requested"""
return self._shutdown_event.is_set()