Named joint states helper

examples/set_get_joints_example.py

@@ -0,0 +1,59 @@
+"""
+This example shows PyRoboPlan capabilities for getting and setting joint states easily with named configurations.
+"""
+
+from pinocchio.visualize import MeshcatVisualizer
+
+from pyroboplan.core.utils import (
+    get_random_collision_free_state,
+)
+from pyroboplan.models.two_dof import load_models, add_object_collisions
+from pyroboplan.visualization.set_get_joints import (
+    NamedJointConfigurationsOptions,
+    NamedJointConfigurations,
+)
+
+if __name__ == "__main__":
+    # Show an example saved joints states with the 2DoF robot
+    model, collision_model, visual_model = load_models()
+    add_object_collisions(model, collision_model, visual_model)
+    ee_name = "ee"
+    data = model.createData()
+    collision_data = collision_model.createData()
+
+    # Define the initial configuration
+    q = get_random_collision_free_state(model, collision_model, distance_padding=0.1)
+
+    # Initialize
+    named_joints_configuration = NamedJointConfigurations(model, collision_model)
+    named_joints_configuration["home"] = q
+
+    # Initialize visualizer
+    visualizer = MeshcatVisualizer(model, collision_model, visual_model, data=data)
+    visualizer.initViewer(open=True)
+    visualizer.loadViewerModel()
+    visualizer.display(q)
+
+    while True:
+        # get user input
+        user_input = input(
+            "Press 'Enter' to show another random state, enter a new name to save the current state as a named joint configuration, enter a used name to visualize that saved state, type 'h' or 'help' to see the existing saved joint configurations, or ctrl-c to quit.\n"
+        )
+        print()
+
+        if user_input:
+            # if the input isn't empty
+            user_input = user_input.lower()
+            if user_input == "h" or user_input == "help":
+                print("Stored states:")
+                print(named_joints_configuration)
+            if user_input in named_joints_configuration:
+                named_joints_configuration.visualize_state(visualizer, user_input)
+            else:
+                named_joints_configuration[user_input] = q
+        else:
+            # if the input is empty, make a new state
+            q = get_random_collision_free_state(
+                model, collision_model, distance_padding=0.1
+            )
+            visualizer.display(q)

src/pyroboplan/core/utils.py

@@ -319,6 +319,25 @@ def check_within_limits(model, q):
     )
 
 
+def check_valid_pose(model, q):
+    """
+    Checks whether a particular joint configuration is a valid configuration for the given model.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model for which to check the joint configuration's validity.
+        q : array-like
+            The joint configuration to check.
+
+    Returns
+    -------
+        bool
+            True if the configuration q is a possible configuration of the given model
+    """
+    return len(q) == model.nq and check_within_limits(model, q)
+
+
 def extract_cartesian_pose(model, target_frame, q, data=None):
     """
     Extracts the Cartesian pose of a specified model frame given a joint configuration.

src/pyroboplan/visualization/set_get_joints.py

@@ -0,0 +1,98 @@
+from pyroboplan.core.utils import (
+    check_collisions_at_state,
+    check_valid_pose,
+    check_within_limits,
+)
+
+
+class NamedJointConfigurationsOptions:
+    """Options for named joint configurations"""
+
+    def __init__(self, allow_collisions=False):
+        """
+        Initializes a set of named joint configuration options.
+
+        Parameters
+        ----------
+            allow_collisions : bool
+                Whether to allow in-collision states to be stored.
+        """
+        self.allow_collisions = allow_collisions
+
+
+class NamedJointConfigurations:
+    """Named joint configurations helper
+
+    This is a simple helper to get, set, and visualize joint states for a given model.
+    Each model should have its own NamedJointConfigurations, and validation is performed to
+    ensure joint states are dimensionally valid and, optionally, within joint limits.
+
+    """
+
+    def __init__(
+        self, model, collision_model, options=NamedJointConfigurationsOptions()
+    ):
+        """
+        Creates an instance of a named joint configurations object.
+
+        Parameters
+        ----------
+            model : `pinocchio.Model`
+                The model to use for this named joint configurations.
+            collision_model : `pinocchio.Model`
+                The model to use for collision checking.
+            options : `NamedJointConfigurationsOptions`, optional
+                The options to use for this named joint configurations object.
+                  If not specified, default options are used.
+        """
+        self.model = model
+        self.collision_model = collision_model
+        self.data = self.model.createData()
+        self.collision_data = self.collision_model.createData()
+        self.options = options
+        self.configuration_dict = {}
+
+    def __setitem__(self, key, state):
+        """
+        Set the configuration at key to state
+        """
+        # first, check if the state is valid
+        if not check_valid_pose(self.model, state):
+            if not check_within_limits(self.model, state):
+                raise ValueError("Tried to add state outside of model's joint limits")
+            raise ValueError(
+                f"Tried to add state of dimension {len(state)} to model with number of joints {self.model.nq}"
+            )
+
+        # then, check collisions if necessary
+        if not self.options.allow_collisions:
+            if check_collisions_at_state(self.model, self.collision_model, state):
+                raise ValueError("Tried to add state in collision")
+
+        # finally, add to the dict; if the key is of the wrong type then the error message from this should be sufficiently explanatory
+        self.configuration_dict[key] = state
+
+    def __getitem__(self, key):
+        """
+        Get the configuration stored at key
+        """
+        # just return the dict at the key. Errors from accessing the dict incorrectly (invalid key type, no entry) should be sufficiently explanatory
+        return self.configuration_dict[key]
+
+    def __contains__(self, key):
+        """
+        See whether this object contains a configuration at key
+        """
+        return key in self.configuration_dict
+
+    def __str__(self):
+        """
+        Print the states in this object
+        """
+        return str([k for k in self.configuration_dict])
+
+    def visualize_state(self, visualizer, name):
+        """
+        Visualize the joint state stored in key `name`
+        """
+        visualizer.display(self[name])