extra loading function + fixes

Author: sWintermoor

src/bitbots_rl_walk/handlers/command_handler.py

@@ -4,8 +4,8 @@
 
 
 class CommandHandler(Handler):
-    def __init__(self, config_file: str):
-        super().__init__(config_file=config_file)
+    def __init__(self, config):
+        super().__init__(config)
 
         self._cmd_vel = None
 

src/bitbots_rl_walk/handlers/gravity_handler.py

@@ -6,8 +6,8 @@
 
 
 class GravityHandler(Handler):
-    def __init__(self, config_file: str):
-        super().__init__(config_file=config_file)
+    def __init__(self, config):
+        super().__init__(config)
 
         self._imu_data = None
         self._gravity = None

src/bitbots_rl_walk/handlers/gyro_handler.py

@@ -4,8 +4,8 @@
 
 
 class GyroHandler(Handler):
-    def __init__(self, config_file: str):
-        super().__init__(config_file=config_file)
+    def __init__(self, config):
+        super().__init__(config)
         self._imu_data = None
 
     # Callables

src/bitbots_rl_walk/handlers/handler.py

@@ -1,13 +1,6 @@
-import yaml
-
-
 class Handler:
-    def __init__(self, config_file: str):
-        self._config = self._load_config(self, config_file)
-
-    def _load_config(self, path: str):
-        with open(path) as f:
-            return yaml.safe_load(f)
+    def __init__(self, config):
+        self._config = config
 
     def get_data(self):
         pass

src/bitbots_rl_walk/handlers/joint_handler.py

@@ -5,8 +5,8 @@
 
 
 class JointHandler(Handler):
-    def __init__(self, config_file: str):
-        super().__init__(config_file=config_file)
+    def __init__(self, config):
+        super().__init__(config)
 
         self._ordered_relevant_joint_names = self._config["joints"]["ordered_relevant_joint_names"]
         self._walkready_state = self._config["joints"]["walkready_state"]

src/bitbots_rl_walk/handlers/phase_handler.py

@@ -7,8 +7,8 @@ class PhaseHandler(Handler):
     _phase: np.ndarray = np.array([0.0, np.pi], dtype=np.float32)
     _phase_dt: float
 
-    def __init__(self, config_file: str):
-        super().__init__(config_file=config_file)
+    def __init__(self, config):
+        super().__init__(config)
 
         self._control_dt = self._config["phase"]["control_dt"]
         self._gait_frequency = self._config["phase"]["gait_frequency"]

src/bitbots_rl_walk/nodes/rl_node.py

@@ -14,12 +14,15 @@
 # ==============================================================================
 """Deploy an MJX policy in ONNX format to C MuJoCo and play with it."""
 
+import os
 from pathlib import Path
 from typing import Callable, NamedTuple
 
 import numpy as np
 import onnx
 import onnxruntime as rt
+import yaml
+from ament_index_python import get_package_share_directory
 from rclpy.node import Node
 from rclpy.qos import QoSProfile
 from rclpy.subscription import Subscription
@@ -42,29 +45,14 @@ class SubscriptionParam(NamedTuple):
         callback: Callable
         qos_profile: int | QoSProfile
 
-    def __init__(self, path_to_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, providers=["CPUExecutionProvider"])
-        self._onnx_model = onnx.load(self._onnx_model_path)
-
-        self._onnx_input_name = []
-        for inp in self._onnx_model.graph.input:
-            self._onnx_input_name.append(inp)
-
-        self._onnx_output_name = []
-        for out in self._onnx_model.graph.output:
-            self._onnx_output_name.append(out)
-
-        self._timer_phase_config = None  # Should be implemented in the subclass
-
-        self._config = False
-
+    def __init__(self, config_path: str):
+        self._config = self._load_config(config_path)
         self._obs = None  # should be defined in subclass
-        self._timer_phase_config = None  # Should be defined in subclass
+        self._phase_handler = None  # shoul be defined in subclass
+
+    def _load_config(self, path: str):
+        with open(path) as f:
+            return yaml.safe_load(f)
 
     # TODO: fix
     def _timer_callback(self):
@@ -85,7 +73,7 @@ def _timer_callback(self):
 
             self._timer_phase_config.set_obs_phase(
                 np.array(
-                    [np.cos(self._timer_phase_config.get_phase()), np.sin(self._timer_phase_config.get_phase())],
+                    [np.cos(self._phase_handler.get_phase()), np.sin(self._phase_handler.get_phase())],
                     dtype=np.float32,
                 ).flatten()
             )
@@ -97,11 +85,31 @@ def _timer_callback(self):
 
             self.publisher(onnx_pred)
 
-            phase_tp1 = self._timer_phase_config.get_phase() + self._timer_phase_config.get_phase_dt()
-            self._timer_phase_config.set_phase(np.fmod(phase_tp1 + np.pi, 2 * np.pi) - np.pi)
+            phase_tp1 = self._phase_handler.get_phase() + self._phase_handler.get_phase_dt()
+            self._phase_handler.set_phase(np.fmod(phase_tp1 + np.pi, 2 * np.pi) - np.pi)
         else:
             raise ConfigError("Configuration is missing! Try to run self.config() in init.")
 
+    def load_model(self, model):
+        path_to_model = os.path.join(get_package_share_directory("bitbots_rl_walk"), "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, providers=["CPUExecutionProvider"])
+        self._onnx_model = onnx.load(self._onnx_model_path)
+
+        self._onnx_input_name = []
+        for inp in self._onnx_model.graph.input:
+            self._onnx_input_name.append(inp)
+
+        self._onnx_output_name = []
+        for out in self._onnx_model.graph.output:
+            self._onnx_output_name.append(out)
+
     def config(self):
         self._timer = self.create_timer(self._timer_phase_config.get_control_dt(), self._timer_callback)
         self.load_phase()

src/bitbots_rl_walk/nodes/walk_node.py

@@ -14,8 +14,12 @@
 
 
 class WalkNode(RLNode):
-    def __init__(self, config):
-        super().__init__(config)
+    def __init__(self, config_path: str):
+        super().__init__(config_path)
+
+        # loading model
+        model = self._config["models"]["walk_model"]
+        self.load_model(model)
 
         # publishers
         self._joint_command_pub = self.create_publisher(JointCommand, "walking_motor_goals", 10)
@@ -26,11 +30,11 @@ def __init__(self, config):
         self._cmd_vel_sub = self.create_subscription(Twist, "cmd_vel", self._cmd_vel_callback, 10)
 
         # handlers
-        self._gyro_handler = GyroHandler()
-        self._gravity_handler = GravityHandler()
-        self._joint_handler = JointHandler()
-        self._command_handler = CommandHandler()
-        self._phase_handler = PhaseHandler()
+        self._gyro_handler = GyroHandler(self._config)
+        self._gravity_handler = GravityHandler(self._config)
+        self._joint_handler = JointHandler(self._config)
+        self._command_handler = CommandHandler(self._config)
+        self._phase_handler = PhaseHandler(self._config)
 
         self.config()