Keep GRU input hooks drive-local

Remove the optional control-input protocol from DriveBase and implement it only on DriveNeuralGRU. ModelBuilder and Actuator continue to discover the private hooks generically without changing public APIs.

Author: vidurv-nvidia

newton/_src/actuators/drives/base.py

@@ -50,22 +50,6 @@ def reset(self, mask: wp.array[wp.bool] | None = None) -> None:
             """
 
     SHARED_PARAMS: ClassVar[set[str]] = set()
-    _control_input_attribute_keys: ClassVar[tuple[str, ...]] = ()
-
-    def _required_control_attributes(self) -> tuple[str, ...]:
-        """Return JOINT_DOF float32 custom Control attributes required by this drive."""
-        return self._control_input_attribute_keys
-
-    def _bind_control_inputs(self, control: Any) -> None:
-        """Bind required colon-addressed attributes from the current Control object."""
-        self._bound_control_inputs = {}
-        bound_inputs = {}
-        for key in self._required_control_attributes():
-            value = control
-            for component in key.split(":"):
-                value = getattr(value, component)
-            bound_inputs[key] = value
-        self._bound_control_inputs = bound_inputs
 
     @classmethod
     def resolve_arguments(cls, args: dict[str, Any]) -> dict[str, Any]:

newton/_src/actuators/drives/drive_neural_gru.py

@@ -394,8 +394,24 @@ def __init__(self, model_path: str):
         self._gru_layers: list[Any] = []
         self._head: Any = None
         self._activation: Any = None
+        self._bound_control_inputs: dict[str, wp.array[float]] = {}
         self._next_hidden: list[wp.array2d[float]] | None = None
 
+    def _required_control_attributes(self) -> tuple[str, ...]:
+        """Return the optional custom Control input selected by metadata."""
+        return self._control_input_attribute_keys
+
+    def _bind_control_inputs(self, control: Any) -> None:
+        """Bind the optional bias input from the current Control object."""
+        self._bound_control_inputs = {}
+        bound_inputs = {}
+        for key in self._required_control_attributes():
+            value = control
+            for component in key.split(":"):
+                value = getattr(value, component)
+            bound_inputs[key] = value
+        self._bound_control_inputs = bound_inputs
+
     def finalize(self, device: wp.Device, num_actuators: int) -> None:
         """Create the Warp-NN layers and inference buffers.
 
@@ -504,7 +520,7 @@ def compute(
         device: wp.Device | None = None,
     ) -> None:
         """Evaluate one GRU sample and write physical effort."""
-        control_inputs = getattr(self, "_bound_control_inputs", {})
+        control_inputs = self._bound_control_inputs
         self._bound_control_inputs = {}
         bias_force = control_inputs.get(_DYNAMIC_BIAS_CONTROL_ATTRIBUTE)
         self._next_hidden = None