refactoring obs

Author: sWintermoor

src/bitbots_motion/bitbots_rl_motion/configs/wolfgang_config.yaml

@@ -1,6 +1,6 @@
 models: 
   walk_model: wolfang_walk_ppo.onnx
-  kick_model: wolfgang_kick_ppo.onnx
+  kick_model: wolfgang_forward_kick_better_ball_ppo.onnx
 
 joints:
   ordered_relevant_joint_names: [

src/bitbots_motion/bitbots_rl_motion/nodes/kick_node.py

@@ -12,7 +12,7 @@
 
 class KickNode(RLNode):
     def __init__(self, config_path: str):
-        super().__init__(config_path, "kick_node")
+        super().__init__(config_path, node_name="kick_node")
 
         # publishers
         self._joint_command_pub = self.create_publisher(JointCommand, "walking_motor_goals", 10)
@@ -28,26 +28,11 @@ def __init__(self, config_path: str):
         self._joint_handler = JointHandler(self._config)
         self._ball_handler = BallHandler(self._config)
 
-        # observations
-
-        self._obs = np.hstack(
-            [
-                self._gyro_handler.get_gyro(),
-                self._gravity_handler.get_gravity(),
-                self._joint_handler.get_velocity_data(),
-                self._joint_handler.get_angle_data(),
-                self._joint_handler.get_previous_action(),
-                self._phase.get_phase(),
-                self._ball_handler.get_ball_pos(),
-            ]
-        ).astype(np.float32)
-
-         # loading model
+        # loading model
         model = self._config["models"]["kick_model"]
         self.load_model(model)
 
     # callback functions
-
     def _imu_callback(self, msg):
         self._gyro_handler.imu_callback(msg)
         self._gravity_handler.imu_callback(msg)
@@ -58,13 +43,25 @@ def _joint_state_callback(self, msg):
     def _ball_pos_callback(self, msg):
         self._ball_handler.ball_pos_callback(msg)
 
-    # load phase function
+    # observations
+    def obs(self):
+        return np.hstack(
+            [
+                self._gyro_handler.get_gyro(),
+                self._gravity_handler.get_gravity(),
+                self._joint_handler.get_velocity_data(),
+                self._joint_handler.get_angle_data(),
+                self._joint_handler.get_previous_action(),
+                self._phase.get_phase(),
+                self._ball_handler.get_ball_pos(),
+            ]
+        ).astype(np.float32)
 
+    # load phase function
     def load_phase(self):
         pass
 
     # publisher function
-
     def publisher(self, onnx_pred):
         joint_command = self._joint_handler.get_joint_commands(onnx_pred)
         self._joint_command_pub.publish(joint_command)

src/bitbots_motion/bitbots_rl_motion/nodes/rl_node.py

@@ -16,7 +16,6 @@
 
 import os
 from pathlib import Path
-from typing import Callable, NamedTuple
 
 from abc import ABC, abstractmethod
 import numpy as np
@@ -26,7 +25,6 @@
 from ament_index_python import get_package_share_directory
 from bitbots_rl_motion.phase import PhaseObject
 from rclpy.node import Node
-from rclpy.qos import QoSProfile
 from rclpy.subscription import Subscription
 
 from handlers.handler import Handler
@@ -40,7 +38,6 @@ def __init__(self, config_path: str, node_name: str):
 
         self._config = self._load_config(config_path)
         self._phase = PhaseObject(self._config)
-        self._obs = None  # should be defined in subclass
 
     def _load_config(self, path: str):
         with open(path) as f:
@@ -64,7 +61,7 @@ def _timer_callback(self):
         )
 
         # Run the ONNX model
-        onnx_input = {self._onnx_input_name[0]: self._obs.reshape(1, -1)}  # TODO: Improve input
+        onnx_input = {self._onnx_input_name[0]: self.obs().reshape(1, -1)}  # TODO: Improve input
         onnx_pred = self._onnx_session.run(self._onnx_output_name, onnx_input)[0][0]
         self._previous_action = onnx_pred
 
@@ -114,6 +111,9 @@ def publisher(self, action):
     def load_phase(self):
         pass
 
+    @abstractmethod
+    def obs(self):
+        pass
 
 class ConfigError(Exception):
     pass

src/bitbots_motion/bitbots_rl_motion/nodes/walk_node.py

@@ -14,7 +14,7 @@
 
 class WalkNode(RLNode):
     def __init__(self, config_path: str):
-        super().__init__(config_path, "walk_node")
+        super().__init__(config_path, node_name="walk_node")
 
         # publishers
         self._joint_command_pub = self.create_publisher(JointCommand, "walking_motor_goals", 10)
@@ -30,26 +30,11 @@ def __init__(self, config_path: str):
         self._joint_handler = JointHandler(self._config)
         self._command_handler = CommandHandler(self._config)
 
-        # observations
-
-        self._obs = np.hstack(
-            [
-                self._gyro_handler.get_gyro(),
-                self._gravity_handler.get_gravity(),
-                self._command_handler.get_command(),
-                self._joint_handler.get_velocity_data(),
-                self._joint_handler.get_angle_data(),
-                self._joint_handler.get_previous_action(),
-                self._phase.get_phase(),
-            ]
-        ).astype(np.float32)
-
         # loading model
         model = self._config["models"]["walk_model"]
         self.load_model(model)
 
     # callback functions
-
     def _imu_callback(self, msg):
         self._gyro_handler.imu_callback(msg)
         self._gravity_handler.imu_callback(msg)
@@ -60,15 +45,27 @@ def _joint_state_callback(self, msg):
     def _cmd_vel_callback(self, msg):
         self._command_handler.cmd_vel_callback(msg)
 
-    # load phase function
+    # observations
+    def obs(self):
+        return np.hstack(
+            [
+                self._gyro_handler.get_gyro(),
+                self._gravity_handler.get_gravity(),
+                self._command_handler.get_command(),
+                self._joint_handler.get_velocity_data(),
+                self._joint_handler.get_angle_data(),
+                self._joint_handler.get_previous_action(),
+                self._phase.get_phase(),
+            ]
+        ).astype(np.float32)
 
+    # load phase function
     def load_phase(self):
         walkready_command = self._joint_handler.get_walkready_joint_command()
         self._joint_command_pub.publish(walkready_command)
         time.sleep(10)
 
     # publisher function
-
     def publisher(self, onnx_pred):
         joint_command = self._joint_handler.get_joint_commands(onnx_pred)
         self._joint_command_pub.publish(joint_command)