Use base_names in base matchers (#69)

# Pull Request

### What change is being made

Use the base_names metadata in base matchers.

### Why this change is being made

Allow arbitrary base names.

### Tested

Adapted tests.

GitOrigin-RevId: dd7d3e298abe386a88eb9824e43e4a9ee54150e6

Author: awoll-bdai

control/context.cpp

@@ -26,6 +26,21 @@ std::optional<int> parseUpdateRate(OnnxRuntime& onnx_model) {
   return static_cast<int>(std::stod(maybe_update_rate.value()));
 }
 
+std::unordered_map<std::string, std::string> parseBaseNames(const OnnxRuntime& onnx_model) {
+  std::unordered_map<std::string, std::string> base_names;
+  const auto maybe_base_names = onnx_model.getCustomMetadata("base_names");
+  if (!maybe_base_names.has_value()) return base_names;
+  try {
+    auto json_base_names = json::parse(maybe_base_names.value());
+    for (auto it = json_base_names.begin(); it != json_base_names.end(); ++it) {
+      base_names[it.key()] = it.value().get<std::string>();
+    }
+  } catch (const json::exception& e) {
+    LOG_STREAM(ERROR, "Failed to parse base_names metadata: " << e.what());
+  }
+  return base_names;
+}
+
 }  // namespace
 
 // Registration methods
@@ -56,10 +71,13 @@ bool OnnxContext::createContext(OnnxRuntime& onnx_model, bool strict) {
   if (!maybe_update_rate.has_value()) return false;
   update_rate_ = maybe_update_rate.value();
 
+  base_names_ = parseBaseNames(onnx_model);
+
   for (const auto& input_name : onnx_model.inputNames()) {
     Match maybe_match{
         .name = input_name,
         .metadata = onnx_model.getCustomMetadata(input_name),
+        .base_names = base_names_,
     };
     bool found_match = false;
     for (auto& group_matchers : group_matchers_) {
@@ -81,6 +99,7 @@ bool OnnxContext::createContext(OnnxRuntime& onnx_model, bool strict) {
     Match maybe_match{
         .name = output_name,
         .metadata = onnx_model.getCustomMetadata(output_name),
+        .base_names = base_names_,
     };
     bool found_match = false;
     for (auto& group_matchers : group_matchers_) {

control/context.hpp

@@ -81,8 +81,9 @@ class OnnxContext {
   std::vector<std::unique_ptr<Output>> outputs_;  ///< Output components for writing robot commands.
   std::vector<std::unique_ptr<Matcher>> matchers_;  ///< Registered single-tensor matchers.
   std::vector<std::unique_ptr<GroupMatcher>>
-      group_matchers_;  ///< Registered multi-tensor matchers.
-  int update_rate_{0};  ///< Control loop update rate in Hz.
+      group_matchers_;                                       ///< Registered multi-tensor matchers.
+  int update_rate_{0};                                       ///< Control loop update rate in Hz.
+  std::unordered_map<std::string, std::string> base_names_;  ///< Map of base names.
 };
 
 }  // namespace exploy::control

control/matcher.cpp

@@ -1,7 +1,9 @@
 // Copyright (c) 2026 Robotics and AI Institute LLC dba RAI Institute. All rights reserved.
 
 #include <fmt/core.h>
+#include <fmt/ranges.h>
 #include <optional>
+#include <ranges>
 #include <regex>
 #include <string>
 #include <string_view>
@@ -13,6 +15,21 @@
 
 namespace exploy::control {
 
+namespace {
+
+// Builds a regex pattern for base tensor matching.
+// Returns std::nullopt when base_names is empty, causing matchers to reject the tensor.
+std::optional<std::regex> buildBasePattern(
+    const std::unordered_map<std::string, std::string>& base_names, std::string_view field) {
+  if (base_names.empty()) return std::nullopt;
+  auto pairs = base_names | std::views::transform([](const auto& p) {
+                 return fmt::format(R"({}\.{})", p.first, p.second);
+               });
+  return std::regex(fmt::format(R"(obj\.({})\.{})", fmt::join(pairs, "|"), field));
+}
+
+}  // namespace
+
 // ---------------  Joint matchers --------------------------------
 bool JointMatcher::matches(const Match& maybe_match) {
   std::smatch match;
@@ -53,8 +70,8 @@ std::vector<std::unique_ptr<Input>> JointMatcher::createInputs() const {
 
 // ---------------  Base matchers --------------------------------
 bool BasePositionMatcher::matches(const Match& maybe_match) {
-  std::regex pattern = std::regex(fmt::format("obj\\.({})\\.base\\.pos_b_rt_w_in_w", alphanumeric));
-  if (std::regex_match(maybe_match.name, pattern)) {
+  auto maybe_pattern = buildBasePattern(maybe_match.base_names, "pos_b_rt_w_in_w");
+  if (maybe_pattern.has_value() && std::regex_match(maybe_match.name, maybe_pattern.value())) {
     found_matches_[maybe_match.name] = maybe_match;
     return true;
   }
@@ -70,8 +87,8 @@ std::vector<std::unique_ptr<Input>> BasePositionMatcher::createInputs() const {
 }
 
 bool BaseOrientationMatcher::matches(const Match& maybe_match) {
-  std::regex pattern = std::regex(fmt::format("obj\\.({})\\.base\\.w_Q_b", alphanumeric));
-  if (std::regex_match(maybe_match.name, pattern)) {
+  auto maybe_pattern = buildBasePattern(maybe_match.base_names, "w_Q_b");
+  if (maybe_pattern.has_value() && std::regex_match(maybe_match.name, maybe_pattern.value())) {
     found_matches_[maybe_match.name] = maybe_match;
     return true;
   }
@@ -87,9 +104,8 @@ std::vector<std::unique_ptr<Input>> BaseOrientationMatcher::createInputs() const
 }
 
 bool BaseLinearVelocityMatcher::matches(const Match& maybe_match) {
-  std::regex pattern =
-      std::regex(fmt::format("obj\\.({})\\.base\\.lin_vel_b_rt_w_in_b", alphanumeric));
-  if (std::regex_match(maybe_match.name, pattern)) {
+  auto maybe_pattern = buildBasePattern(maybe_match.base_names, "lin_vel_b_rt_w_in_b");
+  if (maybe_pattern.has_value() && std::regex_match(maybe_match.name, maybe_pattern.value())) {
     found_matches_[maybe_match.name] = maybe_match;
     return true;
   }
@@ -105,9 +121,8 @@ std::vector<std::unique_ptr<Input>> BaseLinearVelocityMatcher::createInputs() co
 }
 
 bool BaseAngularVelocityMatcher::matches(const Match& maybe_match) {
-  std::regex pattern =
-      std::regex(fmt::format("obj\\.({})\\.base\\.ang_vel_b_rt_w_in_b", alphanumeric));
-  if (std::regex_match(maybe_match.name, pattern)) {
+  auto maybe_pattern = buildBasePattern(maybe_match.base_names, "ang_vel_b_rt_w_in_b");
+  if (maybe_pattern.has_value() && std::regex_match(maybe_match.name, maybe_pattern.value())) {
     found_matches_[maybe_match.name] = maybe_match;
     return true;
   }

control/matcher.hpp

@@ -28,8 +28,9 @@ constexpr std::string_view alphanumeric = "[a-zA-Z0-9_]+";
  * from the ONNX model.
  */
 struct Match {
-  std::string name{};                     ///< ONNX tensor name.
-  std::optional<std::string> metadata{};  ///< Optional JSON metadata string.
+  std::string name{};                                       ///< ONNX tensor name.
+  std::optional<std::string> metadata{};                    ///< Optional JSON metadata string.
+  std::unordered_map<std::string, std::string> base_names;  ///< Map of base names.
 };
 
 /**

control/test/components_test.cpp

@@ -53,7 +53,7 @@ TEST_F(OnnxComponentsTest, JointPositionInput_InitAndRead) {
 }
 
 TEST_F(OnnxComponentsTest, BasePositionInput_InitAndRead) {
-  BasePositionInput base_input("obj.robot1.base.pos_b_rt_w_in_w");
+  BasePositionInput base_input("obj.robot1.base_name.pos_b_rt_w_in_w");
 
   // Test successful initialization
   EXPECT_CALL(state_mock_, initBasePosW()).WillOnce(Return(true));
@@ -67,13 +67,13 @@ TEST_F(OnnxComponentsTest, BasePositionInput_InitAndRead) {
 }
 
 TEST_F(OnnxComponentsTest, BasePositionInput_InitFailure) {
-  BasePositionInput base_input("obj.robot1.base.pos_b_rt_w_in_w");
+  BasePositionInput base_input("obj.robot1.base_name.pos_b_rt_w_in_w");
   EXPECT_CALL(state_mock_, initBasePosW()).WillOnce(Return(false));
   EXPECT_FALSE(base_input.init(state_mock_, command_mock_));
 }
 
 TEST_F(OnnxComponentsTest, BaseOrientationInput_InitAndRead) {
-  BaseOrientationInput base_input("obj.robot1.base.w_Q_b");
+  BaseOrientationInput base_input("obj.robot1.base_name.w_Q_b");
 
   // Test initialization
   EXPECT_CALL(state_mock_, initBaseQuatW()).WillOnce(Return(true));
@@ -87,7 +87,7 @@ TEST_F(OnnxComponentsTest, BaseOrientationInput_InitAndRead) {
 }
 
 TEST_F(OnnxComponentsTest, BaseLinearVelocityInput_InitAndRead) {
-  BaseLinearVelocityInput base_input("obj.robot1.base.lin_vel_b_rt_w_in_b");
+  BaseLinearVelocityInput base_input("obj.robot1.base_name.lin_vel_b_rt_w_in_b");
 
   // Test initialization
   EXPECT_CALL(state_mock_, initBaseLinVelB()).WillOnce(Return(true));
@@ -101,7 +101,7 @@ TEST_F(OnnxComponentsTest, BaseLinearVelocityInput_InitAndRead) {
 }
 
 TEST_F(OnnxComponentsTest, BaseAngularVelocityInput_InitAndRead) {
-  BaseAngularVelocityInput base_input("obj.robot1.base.ang_vel_b_rt_w_in_b");
+  BaseAngularVelocityInput base_input("obj.robot1.base_name.ang_vel_b_rt_w_in_b");
 
   // Test initialization
   EXPECT_CALL(state_mock_, initBaseAngVelB()).WillOnce(Return(true));

control/test/testdata/test_onnx_generator.py

@@ -17,10 +17,10 @@
     "obj.robot1.joints.pos",
     "obj.robot1.joints.vel",
     # base
-    "obj.robot1.base.pos_b_rt_w_in_w",
-    "obj.robot1.base.w_Q_b",
-    "obj.robot1.base.lin_vel_b_rt_w_in_b",
-    "obj.robot1.base.ang_vel_b_rt_w_in_b",
+    "obj.robot1.base_name.pos_b_rt_w_in_w",
+    "obj.robot1.base_name.w_Q_b",
+    "obj.robot1.base_name.lin_vel_b_rt_w_in_b",
+    "obj.robot1.base_name.ang_vel_b_rt_w_in_b",
     # commands
     "cmd.se2_velocity.vel",
     "cmd.se2_velocity.vel_with_range",
@@ -216,6 +216,7 @@ def get_env_metadata() -> dict:
     return {
         "exploy_version": "0.1.0",
         "update_rate": 10.0,
+        "base_names": {"robot1": "base_name"},
     }
 
 

pixi.toml

@@ -25,6 +25,7 @@ pip = "*"
 # FEATURE: Python Development - Python tools, linting, formatting, testing
 [feature.python.dependencies]
 python = "3.11.*"
+onnx = "*"
 onnxscript = "*"
 pybind11 = ">=2.10"
 ruff = "*"