Fix buggy FR3 state machine (#49)

* temporarily don't use the distance sensors in the fr3_pick task because they seem broken

* add cases for negative, 0, and positive initial slider values

Author: alberthli

judo/gui.py

@@ -88,7 +88,12 @@ def _get_gui_element(
     # Create a slider for an integer-valued param.
     if field.type is int:
         if "ui_config" not in field.metadata:
-            min_int, max_int, step_int = 1, 2 * init_value, DEFAULT_SLIDER_STEP_INT
+            if init_value < 0:
+                min_int, max_int, step_int = 2 * init_value, -1, DEFAULT_SLIDER_STEP_INT
+            elif init_value == 0:
+                min_int, max_int, step_int = -5, 5, DEFAULT_SLIDER_STEP_INT
+            else:
+                min_int, max_int, step_int = 1, 2 * init_value, DEFAULT_SLIDER_STEP_INT
         else:
             min_int, max_int, step_int = field.metadata["ui_config"]
 
@@ -108,11 +113,12 @@ def _get_gui_element(
     # Create a slider for a float-valued param.
     elif field.type is float:
         if "ui_config" not in field.metadata:
-            min_float, max_float, step_float = (
-                0.0,
-                2 * init_value,
-                DEFAULT_SLIDER_STEP_FLOAT,
-            )
+            if init_value < 0.0:
+                min_float, max_float, step_float = 2 * init_value, 0.0, DEFAULT_SLIDER_STEP_FLOAT
+            elif init_value == 0.0:
+                min_float, max_float, step_float = -5.0, 5.0, DEFAULT_SLIDER_STEP_FLOAT
+            else:
+                min_float, max_float, step_float = 0.0, 2 * init_value, DEFAULT_SLIDER_STEP_FLOAT
         else:
             min_float, max_float, step_float = field.metadata["ui_config"]
 

judo/tasks/fr3_pick.py

@@ -204,12 +204,16 @@ def pre_rollout(self, curr_state: np.ndarray, config: FR3PickConfig) -> None:
         self._data.qpos[:] = curr_state[: self.model.nq]
         self._data.qvel[:] = curr_state[self.model.nq : self.model.nq + self.model.nv]
         mujoco.mj_forward(self.model, self._data)
-        curr_sensor = self._data.sensordata  # (total_sensor_dim,)
+
+        # BUG: mujoco distance sensor seems to be broken and doesn't always return signed distance, so here we instead
+        # check the object z position
+        # curr_sensor = self._data.sensordata  # (total_sensor_dim,)
 
         phase = Phase.LIFT  # default phase
 
         # check whether the phase is MOVE
-        obj_in_air = curr_sensor[self.obj_table_adr] > 0  # object is not touching the table
+        # obj_in_air = curr_sensor[self.obj_table_adr] > 0  # object is not touching the table
+        obj_in_air = curr_state[self.obj_pos_adr + 2] > 0.02 + 1e-3  # object z position is above the table
         if obj_in_air:
             phase = Phase.MOVE  # if the object is in the air, we are in lift phase
 
@@ -219,8 +223,11 @@ def pre_rollout(self, curr_state: np.ndarray, config: FR3PickConfig) -> None:
             phase = Phase.PLACE  # if the object is in the goal xy, we are in place phase
 
         # check whether the phase is HOMING
-        obj_table_dist = curr_sensor[self.obj_table_adr]
-        if in_goal_xy and obj_table_dist <= 0:
+        # obj_table_dist = curr_sensor[self.obj_table_adr]
+        # if in_goal_xy and obj_table_dist <= 0:
+        #     phase = Phase.HOMING
+        obj_z_pos = curr_state[self.obj_pos_adr + 2]  # z position of the object
+        if in_goal_xy and obj_z_pos <= 0.02 + 1e-3:  # the cube is 4cm wide and we allow a tolerance
             phase = Phase.HOMING
 
         self.phase = phase