Repository: kevinzakka/mink

Issue:
Allow posture task to take in a cost array.
Example use case: assigning different costs to an arm/hand combination.

Discussion:

Patch:
diff --git a/mink/tasks/posture_task.py b/mink/tasks/posture_task.py
index fef3074..3a8da48 100644
--- a/mink/tasks/posture_task.py
+++ b/mink/tasks/posture_task.py
@@ -55,8 +55,8 @@ class PostureTask(Task):
         cost = np.atleast_1d(cost)
         if cost.ndim != 1 or cost.shape[0] not in (1, self.k):
             raise TaskDefinitionError(
-                f"{self.__class__.__name__} cost should be a vector of shape 1"
-                "(aka identical cost for all dofs) or ({self.k},) but got {cost.shape}"
+                f"{self.__class__.__name__} cost must be a vector of shape (1,) "
+                f"(aka identical cost for all dofs) or ({self.k},). Got {cost.shape}"
             )
         if not np.all(cost >= 0.0):
             raise TaskDefinitionError(f"{self.__class__.__name__} cost should be >= 0")


Test patch:
diff --git a/tests/test_posture_task.py b/tests/test_posture_task.py
index 3428691..59b2fd5 100644
--- a/tests/test_posture_task.py
+++ b/tests/test_posture_task.py
@@ -28,6 +28,24 @@ class TestPostureTask(absltest.TestCase):
             PostureTask(model=self.model, cost=(-1))
         self.assertEqual(str(cm.exception), "PostureTask cost should be >= 0")
 
+    def test_cost_correctly_broadcast(self):
+        task = PostureTask(model=self.model, cost=5.0)
+        np.testing.assert_array_equal(task.cost, np.ones((self.model.nv,)) * 5.0)
+        task = PostureTask(model=self.model, cost=[5.0])
+        np.testing.assert_array_equal(task.cost, np.ones((self.model.nv,)) * 5.0)
+        cost = np.random.random(size=(self.model.nv,))
+        task = PostureTask(model=self.model, cost=cost)
+        np.testing.assert_array_equal(task.cost, cost)
+
+    def test_cost_invalid_shape(self):
+        with self.assertRaises(TaskDefinitionError) as cm:
+            PostureTask(model=self.model, cost=(0.5, 2.0))
+        expected_error_message = (
+            "PostureTask cost must be a vector of shape (1,) (aka identical cost for "
+            f"all dofs) or ({self.model.nv},). Got (2,)"
+        )
+        self.assertEqual(str(cm.exception), expected_error_message)
+
     def test_task_raises_error_if_target_is_invalid(self):
         task = PostureTask(model=self.model, cost=1.0)
         with self.assertRaises(InvalidTarget) as cm: