Initial commit

unilabos/ros/nodes/presets/controller_node.py

@@ -0,0 +1,122 @@
+from typing import Callable, Dict
+from std_msgs.msg import Float64
+
+from unilabos.ros.nodes.base_device_node import BaseROS2DeviceNode, DeviceNodeResourceTracker
+
+
+class ControllerNode(BaseROS2DeviceNode):
+    namespace_prefix = "/controllers"
+
+    def __init__(
+        self,
+        device_id: str,
+        controller_func: Callable,
+        update_rate: float,
+        inputs: Dict[str, Dict[str, type | str]],
+        outputs: Dict[str, Dict[str, type]],
+        parameters: Dict,
+        resource_tracker: DeviceNodeResourceTracker,
+    ):
+        """
+        通用控制器节点
+
+        :param controller_id: 控制器的唯一标识符（作为命名空间的一部分）
+        :param update_rate: 控制器更新频率 (Hz)
+        :param controller_func: 控制器函数，接收 Python 格式的 inputs 和 parameters 返回 outputs
+        :param input_types: 输入话题及其消息类型的字典
+        :param output_types: 输出话题及其消息类型的字典
+        :param parameters: 控制器函数的额外参数
+        """
+        # 先准备所需的属性，以便在调用父类初始化前就可以使用
+        self.device_id = device_id
+        self.controller_func = controller_func
+        self.update_rate = update_rate
+        self.update_time = 1.0 / update_rate
+        self.parameters = parameters
+        self.inputs = {topic: None for topic in inputs.keys()}
+        self.control_input_subscribers = {}
+        self.control_output_publishers = {}
+        self.topic_mapping = {
+            **{input_info["topic"]: input for input, input_info in inputs.items()},
+            **{output_info["topic"]: output for output, output_info in outputs.items()},
+        }
+
+        # 调用BaseROS2DeviceNode初始化，使用自身作为driver_instance
+        status_types = {}
+        action_value_mappings = {}
+        hardware_interface = {}
+
+        # 使用短ID作为节点名，完整ID（带namespace_prefix）作为device_id
+        BaseROS2DeviceNode.__init__(
+            self,
+            driver_instance=self,
+            device_id=device_id,
+            status_types=status_types,
+            action_value_mappings=action_value_mappings,
+            hardware_interface=hardware_interface,
+            print_publish=False,
+            resource_tracker=resource_tracker
+        )
+
+        # 原始初始化逻辑
+        # 初始化订阅者
+        for input, input_info in inputs.items():
+            msg_type = input_info["type"]
+            topic = str(input_info["topic"])
+            self.control_input_subscribers[input] = self.create_subscription(
+                msg_type, topic, lambda msg, t=topic: self.input_callback(t, msg), 10
+            )
+
+        # 初始化发布者
+        for output, output_info in outputs.items():
+            self.lab_logger().info(f"Creating publisher for {output} {output_info}")
+            msg_type = output_info["type"]
+            topic = str(output_info["topic"])
+            self.control_output_publishers[output] = self.create_publisher(msg_type, topic, 10)
+
+        # 定时器，用于定期调用控制逻辑
+        self.timer = self.create_timer(self.update_time, self.control_loop)
+
+    def input_callback(self, topic: str, msg):
+        """
+        更新指定话题的输入数据，并将 ROS 消息转换为普通 Python 数据。
+        支持 `std_msgs` 类型消息。
+        """
+        self.inputs[self.topic_mapping[topic]] = msg.data
+        self.lab_logger().info(f"Received input on topic {topic}: {msg.data}")
+
+    def control_loop(self):
+        """主控制逻辑"""
+        # 检查所有输入是否已更新
+        if all(value is not None for value in self.inputs.values()):
+            self.lab_logger().info(
+                f"Calling controller function with inputs: {self.inputs}, parameters: {self.parameters}"
+            )
+            try:
+                # 调用控制器函数，传入 Python 格式的数据
+                outputs = self.controller_func(**self.inputs, **self.parameters)
+                self.lab_logger().info(f"Inputs: {self.inputs}, Outputs: {outputs}")
+                self.inputs = {topic: None for topic in self.inputs.keys()}
+            except Exception as e:
+                self.lab_logger().error(f"Controller function execution failed: {e}")
+                return
+
+            # 发布控制信号，将普通 Python 数据转换为 ROS 消息
+            if isinstance(outputs, dict):
+                for topic, value in outputs.items():
+                    if topic in self.control_output_publishers:
+                        # 支持 Float64 输出
+                        if isinstance(value, (float, int)):
+                            self.control_output_publishers[topic].publish(Float64(data=value))
+                        else:
+                            self.lab_logger().error(f"Unsupported output type for topic {topic}: {type(value)}")
+                    else:
+                        self.lab_logger().warning(f"Output topic {topic} is not defined in output_types.")
+            else:
+                publisher = list(self.control_output_publishers.values())[0]
+                if isinstance(outputs, (float, int)):
+                    publisher.publish(Float64(data=outputs))
+                else:
+                    self.lab_logger().error(f"Unsupported output type: {type(outputs)}")
+        else:
+            self.lab_logger().info("Waiting for all inputs to be received.")