refactor: privatize non-public solver internals (#1683)

Author: eric-heiden

asv/benchmarks/benchmark_mujoco.py

@@ -484,7 +484,8 @@ def print_trace(trace, indent, steps):
         print_trace(sub_trace, indent + 1, steps)
     if indent == 0:
         step_time = trace["step"][0]
-        mujoco_warp_step_time = trace["step"][1]["mujoco_warp_step"][0]
+        step_trace = trace["step"][1]
+        mujoco_warp_step_time = step_trace["_mujoco_warp_step"][0]
         overhead = 100.0 * (step_time - mujoco_warp_step_time) / step_time
         print("---------------------------------------------")
         print(f"Newton overhead:\t{overhead:.2f} %")

asv/benchmarks/simulation/bench_mujoco.py

@@ -181,7 +181,8 @@ def track_simulate(self, world_count):
                     trace = tracer.add_trace(trace, tracer.trace())
 
         step_time = trace["step"][0]
-        mujoco_warp_step_time = trace["step"][1]["mujoco_warp_step"][0]
+        step_trace = trace["step"][1]
+        mujoco_warp_step_time = step_trace["_mujoco_warp_step"][0]
         overhead = 100.0 * (step_time - mujoco_warp_step_time) / step_time
         return overhead
 

newton/_src/solvers/featherstone/solver_featherstone.py

@@ -122,8 +122,8 @@ def __init__(
 
         self._step = 0
 
-        self.compute_articulation_indices(model)
-        self.allocate_model_aux_vars(model)
+        self._compute_articulation_indices(model)
+        self._allocate_model_aux_vars(model)
 
         if self.use_tile_gemm:
             # create a custom kernel to evaluate the system matrix for this type
@@ -140,7 +140,7 @@ def __init__(
             # todo: should not be necessary?
             wp.load_module(device=wp.get_device())
 
-    def compute_articulation_indices(self, model):
+    def _compute_articulation_indices(self, model):
         # calculate total size and offsets of Jacobian and mass matrices for entire system
         if model.joint_count:
             self.J_size = 0
@@ -210,7 +210,7 @@ def compute_articulation_indices(self, model):
             self.articulation_dof_start = wp.array(articulation_dof_start, dtype=wp.int32, device=model.device)
             self.articulation_coord_start = wp.array(articulation_coord_start, dtype=wp.int32, device=model.device)
 
-    def allocate_model_aux_vars(self, model):
+    def _allocate_model_aux_vars(self, model):
         # allocate mass, Jacobian matrices, and other auxiliary variables pertaining to the model
         if model.joint_count:
             # system matrices
@@ -244,7 +244,7 @@ def allocate_model_aux_vars(self, model):
                 device=model.device,
             )
 
-    def allocate_state_aux_vars(self, model, target, requires_grad):
+    def _allocate_state_aux_vars(self, model, target, requires_grad):
         # allocate auxiliary variables that vary with state
         if model.body_count:
             # joints
@@ -302,7 +302,7 @@ def step(
         model = self.model
 
         if not getattr(state_aug, "_featherstone_augmented", False):
-            self.allocate_state_aux_vars(model, state_aug, requires_grad)
+            self._allocate_state_aux_vars(model, state_aug, requires_grad)
         if control is None:
             control = model.control(clone_variables=False)
 

newton/_src/solvers/mujoco/solver_mujoco.py

@@ -2091,34 +2091,34 @@ def __init__(
             self.mjw_model.opt.run_collision_detection = use_mujoco_contacts
 
     @event_scope
-    def mujoco_warp_step(self):
+    def _mujoco_warp_step(self):
         self._mujoco_warp.step(self.mjw_model, self.mjw_data)
 
     @event_scope
     @override
     def step(self, state_in: State, state_out: State, control: Control, contacts: Contacts, dt: float):
         if self.use_mujoco_cpu:
-            self.apply_mjc_control(self.model, state_in, control, self.mj_data)
+            self._apply_mjc_control(self.model, state_in, control, self.mj_data)
             if self.update_data_interval > 0 and self._step % self.update_data_interval == 0:
                 # XXX updating the mujoco state at every step may introduce numerical instability
-                self.update_mjc_data(self.mj_data, self.model, state_in)
+                self._update_mjc_data(self.mj_data, self.model, state_in)
             self.mj_model.opt.timestep = dt
             self._mujoco.mj_step(self.mj_model, self.mj_data)
-            self.update_newton_state(self.model, state_out, self.mj_data)
+            self._update_newton_state(self.model, state_out, self.mj_data)
         else:
             self.enable_rne_postconstraint(state_out)
-            self.apply_mjc_control(self.model, state_in, control, self.mjw_data)
+            self._apply_mjc_control(self.model, state_in, control, self.mjw_data)
             if self.update_data_interval > 0 and self._step % self.update_data_interval == 0:
-                self.update_mjc_data(self.mjw_data, self.model, state_in)
+                self._update_mjc_data(self.mjw_data, self.model, state_in)
             self.mjw_model.opt.timestep.fill_(dt)
             with wp.ScopedDevice(self.model.device):
                 if self.mjw_model.opt.run_collision_detection:
-                    self.mujoco_warp_step()
+                    self._mujoco_warp_step()
                 else:
-                    self.convert_contacts_to_mjwarp(self.model, state_in, contacts)
-                    self.mujoco_warp_step()
+                    self._convert_contacts_to_mjwarp(self.model, state_in, contacts)
+                    self._mujoco_warp_step()
 
-            self.update_newton_state(self.model, state_out, self.mjw_data)
+            self._update_newton_state(self.model, state_out, self.mjw_data)
         self._step += 1
         return state_out
 
@@ -2135,7 +2135,7 @@ def enable_rne_postconstraint(self, state_out: State):
                 print("Setting model.sensor_rne_postconstraint True")
             m.sensor_rne_postconstraint = True
 
-    def convert_contacts_to_mjwarp(self, model: Model, state_in: State, contacts: Contacts):
+    def _convert_contacts_to_mjwarp(self, model: Model, state_in: State, contacts: Contacts):
         # Ensure the inverse shape mapping exists (lazy creation)
         if self.newton_shape_to_mjc_geom is None:
             self._create_inverse_shape_mapping()
@@ -2194,16 +2194,16 @@ def convert_contacts_to_mjwarp(self, model: Model, state_in: State, contacts: Co
     @override
     def notify_model_changed(self, flags: int):
         if flags & SolverNotifyFlags.BODY_INERTIAL_PROPERTIES:
-            self.update_model_inertial_properties()
+            self._update_model_inertial_properties()
         if flags & SolverNotifyFlags.JOINT_PROPERTIES:
-            self.update_joint_properties()
+            self._update_joint_properties()
         if flags & SolverNotifyFlags.JOINT_DOF_PROPERTIES:
-            self.update_joint_dof_properties()
+            self._update_joint_dof_properties()
         if flags & SolverNotifyFlags.SHAPE_PROPERTIES:
-            self.update_geom_properties()
+            self._update_geom_properties()
             self.update_pair_properties()
         if flags & SolverNotifyFlags.MODEL_PROPERTIES:
-            self.update_model_properties()
+            self._update_model_properties()
         if flags & SolverNotifyFlags.CONSTRAINT_PROPERTIES:
             self.update_eq_properties()
             self.update_mimic_eq_properties()
@@ -2241,7 +2241,7 @@ def _data_is_mjwarp(data):
         # Check if the data is a mujoco_warp Data object
         return hasattr(data, "nworld")
 
-    def apply_mjc_control(self, model: Model, state: State, control: Control | None, mj_data: MjWarpData | MjData):
+    def _apply_mjc_control(self, model: Model, state: State, control: Control | None, mj_data: MjWarpData | MjData):
         if control is None or control.joint_f is None:
             if state.body_f is None:
                 return
@@ -2337,7 +2337,7 @@ def apply_mjc_control(self, model: Model, state: State, control: Control | None,
             mj_data.ctrl[:] = ctrl.numpy().flatten()
             mj_data.qfrc_applied[:] = qfrc.numpy()
 
-    def update_mjc_data(self, mj_data: MjWarpData | MjData, model: Model, state: State | None = None):
+    def _update_mjc_data(self, mj_data: MjWarpData | MjData, model: Model, state: State | None = None):
         is_mjwarp = SolverMuJoCo._data_is_mjwarp(mj_data)
         if is_mjwarp:
             # we have an MjWarp Data object
@@ -2380,7 +2380,7 @@ def update_mjc_data(self, mj_data: MjWarpData | MjData, model: Model, state: Sta
             mj_data.qpos[:] = qpos.numpy().flatten()[: len(mj_data.qpos)]
             mj_data.qvel[:] = qvel.numpy().flatten()[: len(mj_data.qvel)]
 
-    def update_newton_state(
+    def _update_newton_state(
         self,
         model: Model,
         state: State,
@@ -2468,7 +2468,7 @@ def update_newton_state(
             )
 
     @staticmethod
-    def find_body_collision_filter_pairs(
+    def _find_body_collision_filter_pairs(
         model: Model,
         selected_bodies: nparray,
         colliding_shapes: nparray,
@@ -2504,7 +2504,7 @@ def find_body_collision_filter_pairs(
         return body_exclude_pairs
 
     @staticmethod
-    def color_collision_shapes(
+    def _color_collision_shapes(
         model: Model, selected_shapes: nparray, visualize_graph: bool = False, shape_keys: list[str] | None = None
     ) -> nparray:
         """
@@ -3083,13 +3083,13 @@ def get_custom_attribute(name: str) -> nparray | None:
         colliding_shapes_per_world = len(colliding_shapes)
 
         # filter out non-colliding bodies using excludes
-        body_filters = self.find_body_collision_filter_pairs(
+        body_filters = self._find_body_collision_filter_pairs(
             model,
             selected_bodies,
             colliding_shapes,
         )
 
-        shape_color = self.color_collision_shapes(
+        shape_color = self._color_collision_shapes(
             model, colliding_shapes, visualize_graph=False, shape_keys=model.shape_key
         )
 
@@ -3819,9 +3819,9 @@ def get_body_name(body_idx: int) -> str:
         self.mj_model = spec.compile()
         self.mj_data = mujoco.MjData(self.mj_model)
 
-        self.update_mjc_data(self.mj_data, model, state)
+        self._update_mjc_data(self.mj_data, model, state)
 
-        # fill some MjWarp model fields that are outdated after update_mjc_data.
+        # fill some MjWarp model fields that are outdated after _update_mjc_data.
         # just setting qpos0 to d.qpos leads to weird behavior here, needs
         # to be investigated.
 
@@ -4090,7 +4090,7 @@ def get_body_name(body_idx: int) -> str:
             )
 
             # expand model fields that can be expanded:
-            self.expand_model_fields(self.mjw_model, nworld)
+            self._expand_model_fields(self.mjw_model, nworld)
 
             # update solver options from Newton model (only if not overridden by constructor)
             self._update_solver_options(overridden_options=overridden_options)
@@ -4099,7 +4099,7 @@ def get_body_name(body_idx: int) -> str:
             # now complete the data from the Newton model
             self.notify_model_changed(SolverNotifyFlags.ALL)
 
-    def expand_model_fields(self, mj_model: MjWarpModel, nworld: int):
+    def _expand_model_fields(self, mj_model: MjWarpModel, nworld: int):
         if nworld == 1:
             return
 
@@ -4309,7 +4309,7 @@ def get_option(name: str):
             device=self.model.device,
         )
 
-    def update_model_inertial_properties(self):
+    def _update_model_inertial_properties(self):
         if self.model.body_count == 0:
             return
 
@@ -4357,7 +4357,7 @@ def update_model_inertial_properties(self):
         # - cam_pos0, light_pos0, actuator_acc0: other derived quantities
         self._mujoco_warp.set_const(self.mjw_model, self.mjw_data)
 
-    def update_joint_dof_properties(self):
+    def _update_joint_dof_properties(self):
         """Update all joint DOF properties including effort limits, friction, armature, solimplimit, solref, passive stiffness and damping, and joint limit ranges in the MuJoCo model."""
         if self.model.joint_dof_count == 0:
             return
@@ -4482,7 +4482,7 @@ def update_joint_dof_properties(self):
         # - cam_pos0, light_pos0, actuator_acc0: other derived quantities
         self._mujoco_warp.set_const(self.mjw_model, self.mjw_data)
 
-    def update_joint_properties(self):
+    def _update_joint_properties(self):
         """Update joint properties including joint positions, joint axes, and relative body transforms in the MuJoCo model."""
         if self.model.joint_count == 0:
             return
@@ -4535,7 +4535,7 @@ def update_joint_properties(self):
                 device=self.model.device,
             )
 
-    def update_geom_properties(self):
+    def _update_geom_properties(self):
         """Update geom properties including collision radius, friction, and contact parameters in the MuJoCo model."""
 
         # Get number of geoms and worlds from MuJoCo model
@@ -4650,7 +4650,7 @@ def update_pair_properties(self):
                 device=self.model.device,
             )
 
-    def update_model_properties(self):
+    def _update_model_properties(self):
         """Update model properties including gravity in the MuJoCo model."""
         if self.use_mujoco_cpu:
             self.mj_model.opt.gravity[:] = np.array([*self.model.gravity.numpy()[0]])
@@ -4843,7 +4843,7 @@ def update_actuator_properties(self):
         """Update CTRL_DIRECT actuator properties (gainprm, biasprm) in the MuJoCo model.
 
         Only updates actuators that use CTRL_DIRECT mode. JOINT_TARGET actuators are
-        updated via update_joint_dof_properties() using joint_target_ke/kd.
+        updated via _update_joint_dof_properties() using joint_target_ke/kd.
         """
         if self.mjc_actuator_ctrl_source is None or self.mjc_actuator_to_newton_idx is None:
             return

newton/_src/solvers/vbd/particle_vbd_kernels.py

@@ -643,7 +643,7 @@ def compute_cofactor_derivative(F: wp.mat33, scale: float) -> mat99:
 
 
 @wp.kernel
-def count_num_adjacent_edges(
+def _count_num_adjacent_edges(
     edges_array: wp.array(dtype=wp.int32, ndim=2), num_vertex_adjacent_edges: wp.array(dtype=wp.int32)
 ):
     for edge_id in range(edges_array.shape[0]):
@@ -663,7 +663,7 @@ def count_num_adjacent_edges(
 
 
 @wp.kernel
-def fill_adjacent_edges(
+def _fill_adjacent_edges(
     edges_array: wp.array(dtype=wp.int32, ndim=2),
     vertex_adjacent_edges_offsets: wp.array(dtype=wp.int32),
     vertex_adjacent_edges_fill_count: wp.array(dtype=wp.int32),
@@ -703,7 +703,7 @@ def fill_adjacent_edges(
 
 
 @wp.kernel
-def count_num_adjacent_faces(
+def _count_num_adjacent_faces(
     face_indices: wp.array(dtype=wp.int32, ndim=2), num_vertex_adjacent_faces: wp.array(dtype=wp.int32)
 ):
     for face in range(face_indices.shape[0]):
@@ -717,7 +717,7 @@ def count_num_adjacent_faces(
 
 
 @wp.kernel
-def fill_adjacent_faces(
+def _fill_adjacent_faces(
     face_indices: wp.array(dtype=wp.int32, ndim=2),
     vertex_adjacent_faces_offsets: wp.array(dtype=wp.int32),
     vertex_adjacent_faces_fill_count: wp.array(dtype=wp.int32),
@@ -748,7 +748,7 @@ def fill_adjacent_faces(
 
 
 @wp.kernel
-def count_num_adjacent_springs(
+def _count_num_adjacent_springs(
     springs_array: wp.array(dtype=wp.int32), num_vertex_adjacent_springs: wp.array(dtype=wp.int32)
 ):
     num_springs = springs_array.shape[0] / 2
@@ -761,7 +761,7 @@ def count_num_adjacent_springs(
 
 
 @wp.kernel
-def fill_adjacent_springs(
+def _fill_adjacent_springs(
     springs_array: wp.array(dtype=wp.int32),
     vertex_adjacent_springs_offsets: wp.array(dtype=wp.int32),
     vertex_adjacent_springs_fill_count: wp.array(dtype=wp.int32),
@@ -784,7 +784,7 @@ def fill_adjacent_springs(
 
 
 @wp.kernel
-def count_num_adjacent_tets(
+def _count_num_adjacent_tets(
     tet_indices: wp.array(dtype=wp.int32, ndim=2), num_vertex_adjacent_tets: wp.array(dtype=wp.int32)
 ):
     for tet in range(tet_indices.shape[0]):
@@ -800,7 +800,7 @@ def count_num_adjacent_tets(
 
 
 @wp.kernel
-def fill_adjacent_tets(
+def _fill_adjacent_tets(
     tet_indices: wp.array(dtype=wp.int32, ndim=2),
     vertex_adjacent_tets_offsets: wp.array(dtype=wp.int32),
     vertex_adjacent_tets_fill_count: wp.array(dtype=wp.int32),

newton/_src/solvers/vbd/rigid_vbd_kernels.py

@@ -1054,7 +1054,7 @@ def evaluate_joint_force_hessian(
 # Utility kernels
 # -----------------------------
 @wp.kernel
-def count_num_adjacent_joints(
+def _count_num_adjacent_joints(
     joint_parent: wp.array(dtype=wp.int32),
     joint_child: wp.array(dtype=wp.int32),
     num_body_adjacent_joints: wp.array(dtype=wp.int32),
@@ -1072,7 +1072,7 @@ def count_num_adjacent_joints(
 
 
 @wp.kernel
-def fill_adjacent_joints(
+def _fill_adjacent_joints(
     joint_parent: wp.array(dtype=wp.int32),
     joint_child: wp.array(dtype=wp.int32),
     body_adjacent_joints_offsets: wp.array(dtype=wp.int32),