small bug fixes

Author: sWintermoor

src/bitbots_motion/bitbots_rl_motion/bitbots_rl_motion/phase.py

@@ -8,10 +8,9 @@ class PhaseObject:
     _phase_dt: float
 
     def __init__(self, config):
-        super().__init__(config)
 
-        self._control_dt = self._config["phase"]["control_dt"]
-        self._gait_frequency = self._config["phase"]["gait_frequency"]
+        self._control_dt = config["phase"]["control_dt"]
+        self._gait_frequency = config["phase"]["gait_frequency"]
         self._phase_dt = 2 * np.pi * self._gait_frequency * self._control_dt
 
     def set_phase(self, new_phase):

src/bitbots_motion/bitbots_rl_motion/bitbots_rl_motion/policy_nodes.py

@@ -8,7 +8,7 @@
 def main():
     rclpy.init()
 
-    wolfgang_config = os.path.join(get_package_share_directory("bitbots_rl_motion"), "config", "wolfgang_config.yaml")
+    wolfgang_config = os.path.join(get_package_share_directory("bitbots_rl_motion"), "configs", "wolfgang_config.yaml")
 
     # walk_node = WalkNode(wolfgang_config)
     kick_node = KickNode(wolfgang_config)

src/bitbots_motion/bitbots_rl_motion/handlers/ball_handler.py

@@ -1,4 +1,4 @@
-from bitbots_rl_motion.handlers.handler import Handler
+from handlers.handler import Handler
 
 
 class BallHandler(Handler):
@@ -9,5 +9,9 @@ def __init__(self, config):
     def ball_pos_callback(self, msg):
         self.ball_pos = msg
 
+
     def get_ball_pos(self):
-        return self.ball_pos
+        try:
+            return self.ball_pos
+        except:
+            return None

src/bitbots_motion/bitbots_rl_motion/handlers/command_handler.py

@@ -1,5 +1,5 @@
 import numpy as np
-from bitbots_rl_motion.handlers.handler import Handler
+from handlers.handler import Handler
 
 
 class CommandHandler(Handler):
@@ -9,8 +9,12 @@ def __init__(self, config):
         self._cmd_vel = None
 
     def get_command(self):
-        command = np.array([self._cmd_vel.linear.x, self._cmd_vel.linear.y, self._cmd_vel.angular.z], dtype=np.float32)
-        return command
+        try:
+            command = np.array([self._cmd_vel.linear.x, self._cmd_vel.linear.y, self._cmd_vel.angular.z], dtype=np.float32)
+            return command
+        except: 
+            return None
 
     def cmd_vel_callback(self, msg):
         self._cmd_vel = msg
+    

src/bitbots_motion/bitbots_rl_motion/handlers/gravity_handler.py

@@ -1,5 +1,5 @@
 import numpy as np
-from bitbots_rl_motion.handlers.handler import Handler
+from handlers.handler import Handler
 from transforms3d.euler import euler2mat
 from transforms3d.quaternions import quat2mat
 
@@ -15,17 +15,21 @@ def __init__(self, config):
     def imu_callback(self, msg):
         self._imu_data = msg
 
+
     def get_gravity(self):
-        gravity = (
-            quat2mat(
-                [
-                    self._imu_data.orientation.w,
-                    self._imu_data.orientation.x,
-                    self._imu_data.orientation.y,
-                    self._imu_data.orientation.z,
-                ]
-            )
-            @ euler2mat(0, -0.0, 0)
-        ).T @ np.array([0, 0, -1], dtype=np.float32)
+        try:
+            gravity = (
+                quat2mat(
+                    [
+                        self._imu_data.orientation.w,
+                        self._imu_data.orientation.x,
+                        self._imu_data.orientation.y,
+                        self._imu_data.orientation.z,
+                    ]
+                )
+                @ euler2mat(0, -0.0, 0)
+            ).T @ np.array([0, 0, -1], dtype=np.float32)
+            return gravity
+        except:
+            return None
 
-        return gravity

src/bitbots_motion/bitbots_rl_motion/handlers/gyro_handler.py

@@ -1,5 +1,5 @@
 import numpy as np
-from bitbots_rl_motion.handlers.handler import Handler
+from handlers.handler import Handler
 
 
 class GyroHandler(Handler):
@@ -10,15 +10,19 @@ def __init__(self, config):
     # Callables
     def imu_callback(self, msg):
         self._imu_data = msg
+        
 
     def get_gyro(self):
-        gyro = np.array(
-            [
-                self._imu_data.angular_velocity.x,
-                self._imu_data.angular_velocity.y,
-                self._imu_data.angular_velocity.z,
-            ],
-            dtype=np.float32,
-        )
+        try:
+            gyro = np.array(
+                [
+                    self._imu_data.angular_velocity.x,
+                    self._imu_data.angular_velocity.y,
+                    self._imu_data.angular_velocity.z,
+                ],
+                dtype=np.float32,
+            )
 
-        return gyro
+            return gyro
+        except:
+            return None

src/bitbots_motion/bitbots_rl_motion/handlers/joint_handler.py

@@ -1,5 +1,5 @@
 import numpy as np
-from bitbots_rl_motion.handlers.handler import Handler
+from handlers.handler import Handler
 
 from bitbots_msgs.msg import JointCommand
 
@@ -12,77 +12,77 @@ def __init__(self, config):
         self._walkready_state = self._config["joints"]["walkready_state"]
         self._previous_action: np.ndarray = np.zeros(len(self._ordered_relevant_joint_names), dtype=np.float32)
         self._joint_state = None
-        self._obs_phase = None
-        self._phase = None
+
 
     def joint_state_callback(self, msg):
         self._joint_state = msg
 
-    def set_obs_phase(self, phase):
-        self._obs_phase = phase
+    def get_angle_data(self):
+        try:
+            joint_angles = (
+                np.array(
+                    [
+                        self._joint_state.position[self._joint_state.name.index(name)]
+                        for name in self._ordered_relevant_joint_names
+                    ],
+                    dtype=np.float32,
+                )
+                - self._walkready_state
+            )
 
-    def set_phase(self, phase):
-        self._phase = phase
+            return joint_angles
+        except: 
+            return None
 
-    def get_angle_data(self):
-        joint_angles = (
-            np.array(
+    def get_velocity_data(self):
+        try:
+            joint_velocities = np.array(
                 [
-                    self._joint_state.position[self._joint_state.name.index(name)]
+                    self._joint_state.velocity[self._joint_state.name.index(name)]
                     for name in self._ordered_relevant_joint_names
                 ],
                 dtype=np.float32,
             )
-            - self._walkready_state
-        )
 
-        return joint_angles
-
-    def get_velocity_data(self):
-        joint_velocities = np.array(
-            [
-                self._joint_state.velocity[self._joint_state.name.index(name)]
-                for name in self._ordered_relevant_joint_names
-            ],
-            dtype=np.float32,
-        )
-
-        return joint_velocities
+            return joint_velocities
+        except:
+            return None
 
     def get_data(self):
         return self.get_angle_data(), self.get_velocity_data()
 
     def get_walkready_joint_command(self, timestamp):
-        joint_command = JointCommand()
-        joint_command.joint_names = self._ordered_relevant_joint_names
-        joint_command.velocities = [0.2] * len(self._ordered_relevant_joint_names)
-        joint_command.accelerations = [-1.0] * len(self._ordered_relevant_joint_names)
-        joint_command.max_currents = [-1.0] * len(self._ordered_relevant_joint_names)  # -1.0 means no limit
-        joint_command.header.stamp = timestamp.to_msg()
-        joint_command.positions = self._walkready_state
+        try:
+            joint_command = JointCommand()
+            joint_command.joint_names = self._ordered_relevant_joint_names
+            joint_command.velocities = [0.2] * len(self._ordered_relevant_joint_names)
+            joint_command.accelerations = [-1.0] * len(self._ordered_relevant_joint_names)
+            joint_command.max_currents = [-1.0] * len(self._ordered_relevant_joint_names)  # -1.0 means no limit
+            joint_command.header.stamp = timestamp.to_msg()
+            joint_command.positions = self._walkready_state
 
-        self._previous_action = joint_command
+            self._previous_action = joint_command
 
-        return joint_command
+            return joint_command
+        except:
+            return None
 
     def get_joint_commands(self, onnx_pred):
-        joint_command = JointCommand()
-        joint_command.header.stamp = self._joint_state.header.stamp
-        joint_command.joint_names = self._ordered_relevant_joint_names
-        joint_command.positions = onnx_pred * 0.5 + self._walkready_state
-        joint_command.velocities = [-1.0] * len(self._ordered_relevant_joint_names)
-        joint_command.accelerations = [-1.0] * len(self._ordered_relevant_joint_names)
-        joint_command.max_currents = [-1.0] * len(self._ordered_relevant_joint_names)
+        try:
+            joint_command = JointCommand()
+            joint_command.header.stamp = self._joint_state.header.stamp
+            joint_command.joint_names = self._ordered_relevant_joint_names
+            joint_command.positions = onnx_pred * 0.5 + self._walkready_state
+            joint_command.velocities = [-1.0] * len(self._ordered_relevant_joint_names)
+            joint_command.accelerations = [-1.0] * len(self._ordered_relevant_joint_names)
+            joint_command.max_currents = [-1.0] * len(self._ordered_relevant_joint_names)
 
-        self._previous_action = joint_command
+            self._previous_action = joint_command
 
-        return joint_command
+            return joint_command
+        except:
+            return None
 
     def get_previous_action(self):
         return self._previous_action
 
-    def get_obs_phase(self):
-        return self._obs_phase
-
-    def get_phase(self):
-        return self._phase

src/bitbots_motion/bitbots_rl_motion/nodes/kick_node.py

@@ -1,7 +1,7 @@
 import numpy as np
-from bitbots_rl_motion.nodes.rl_node import RLNode
-from bitbots_rl_motion.handlers.joint_handler import JointHandler
-from bitbots_rl_motion.handlers.ball_handler import BallHandler
+from nodes.rl_node import RLNode
+from handlers.joint_handler import JointHandler
+from handlers.ball_handler import BallHandler
 from geometry_msgs.msg import PoseStamped
 from handlers.gravity_handler import GravityHandler
 from handlers.gyro_handler import GyroHandler
@@ -12,11 +12,7 @@
 
 class KickNode(RLNode):
     def __init__(self, config_path: str):
-        super().__init__(config_path)
-
-        # loading model
-        model = self._config["models"]["kick_model"]
-        self.load_model(model)
+        super().__init__(config_path, "kick_node")
 
         # publishers
         self._joint_command_pub = self.create_publisher(JointCommand, "walking_motor_goals", 10)
@@ -46,6 +42,10 @@ def __init__(self, config_path: str):
             ]
         ).astype(np.float32)
 
+         # loading model
+        model = self._config["models"]["kick_model"]
+        self.load_model(model)
+
     # callback functions
 
     def _imu_callback(self, msg):

src/bitbots_motion/bitbots_rl_motion/nodes/rl_node.py

@@ -23,7 +23,7 @@
 import onnxruntime as rt
 import yaml
 from ament_index_python import get_package_share_directory
-from bitbots_rl_motion.bitbots_rl_motion.phase import PhaseObject
+from bitbots_rl_motion.phase import PhaseObject
 from rclpy.node import Node
 from rclpy.qos import QoSProfile
 from rclpy.subscription import Subscription
@@ -43,20 +43,13 @@ class SubscriptionParam(NamedTuple):
         callback: Callable
         qos_profile: int | QoSProfile
 
-    def __init__(self, config_path: str):
+    def __init__(self, config_path: str, node_name: str):
+        super().__init__(f"{node_name}")
+        
         self._config = self._load_config(config_path)
         self._phase = PhaseObject(self._config)
         self._obs = None  # should be defined in subclass
 
-        self._timer = self.create_timer(self._config["phase"]["control_dt"], self._timer_callback)
-        self.load_phase()
-
-        self._subs = []
-
-        for key, value in self.__dict__.values():
-            if type(value) is Subscription:
-                self._subs.append(key)
-
     def _load_config(self, path: str):
         with open(path) as f:
             return yaml.safe_load(f)
@@ -101,12 +94,9 @@ def load_model(self, model):
         path_to_model = os.path.join(get_package_share_directory("bitbots_rl_motion"), "models", model)
 
         self._onnx_model_path = Path(path_to_model)
-        model_name = self._onnx_model_path.stem
-
-        super().__init__(f"{model_name}")
 
         # Load the ONNX model
-        self._onnx_session = rt.InferenceSession(self._onnx_model_path, self._config["providers"])
+        self._onnx_session = rt.InferenceSession(self._onnx_model_path, providers=self._config["providers"])
         self._onnx_model = onnx.load(self._onnx_model_path)
 
         self._onnx_input_name = []
@@ -117,6 +107,16 @@ def load_model(self, model):
         for out in self._onnx_model.graph.output:
             self._onnx_output_name.append(out)
 
+        self._timer = self.create_timer(self._config["phase"]["control_dt"], self._timer_callback)
+
+        self._subs = []
+
+        for (key, value) in self.__dict__.values():
+            if type(value) is Subscription:
+                self._subs.append(key)
+
+        self.load_phase()
+
     def obs(self):
         # Should be defined in subclass
         pass