Merge pull request #9 from rail-berkeley/env_cleanup

Env cleanup

serl_robot_infra/franka_env/envs/cable_env/config.py

@@ -46,3 +46,43 @@ class CableEnvConfig(DefaultEnvConfig):
             TARGET_POSE[5] + RANDOM_RZ_RANGE,
         ]
     )
+    COMPLIANCE_PARAM = {
+        "translational_stiffness": 2000,
+        "translational_damping": 89,
+        "rotational_stiffness": 150,
+        "rotational_damping": 7,
+        "translational_Ki": 0,
+        "translational_clip_x": 0.005,
+        "translational_clip_y": 0.005,
+        "translational_clip_z": 0.005,
+        "translational_clip_neg_x": 0.005,
+        "translational_clip_neg_y": 0.005,
+        "translational_clip_neg_z": 0.005,
+        "rotational_clip_x": 0.05,
+        "rotational_clip_y": 0.05,
+        "rotational_clip_z": 0.02,
+        "rotational_clip_neg_x": 0.05,
+        "rotational_clip_neg_y": 0.05,
+        "rotational_clip_neg_z": 0.02,
+        "rotational_Ki": 0,
+    }
+    PRECISION_PARAM = {
+        "translational_stiffness": 3000,
+        "translational_damping": 89,
+        "rotational_stiffness": 300,
+        "rotational_damping": 9,
+        "translational_Ki": 0.1,
+        "translational_clip_x": 0.01,
+        "translational_clip_y": 0.01,
+        "translational_clip_z": 0.01,
+        "translational_clip_neg_x": 0.01,
+        "translational_clip_neg_y": 0.01,
+        "translational_clip_neg_z": 0.01,
+        "rotational_clip_x": 0.05,
+        "rotational_clip_y": 0.05,
+        "rotational_clip_z": 0.05,
+        "rotational_clip_neg_x": 0.05,
+        "rotational_clip_neg_y": 0.05,
+        "rotational_clip_neg_z": 0.05,
+        "rotational_Ki": 0.1,
+    }

serl_robot_infra/franka_env/envs/cable_env/franka_cable_route.py

@@ -1,5 +1,4 @@
 import numpy as np
-import gymnasium as gym
 import time
 import requests
 import copy
@@ -15,23 +14,28 @@ class FrankaCableRoute(FrankaEnv):
     def __init__(self, **kwargs):
         super().__init__(**kwargs, config=CableEnvConfig)
 
-    def go_to_rest(self, jpos=False):
+    def go_to_rest(self, joint_reset=False):
         self.update_currpos()
-        self._send_pos_command(self.clip_safety_box(self.currpos))
+        self._send_pos_command(self.currpos)
         time.sleep(0.5)
 
-        requests.post(self.url + "precision_mode")
-        time.sleep(0.5)  # wait for mode switching
-
+        # Move up to clear the slot
         self.update_currpos()
         reset_pose = copy.deepcopy(self.currpos)
         reset_pose[2] += 0.05
         self.interpolate_move(reset_pose, timeout=1)
 
-        reset_pose = self.resetpos.copy()
-        self.interpolate_move(reset_pose, timeout=1)
+        # Change to precision mode for reset
+        requests.post(self.url + "update_param", json=self.config.PRECISION_PARAM)
+        time.sleep(0.5)
+
+        # Perform joint reset if needed
+        if joint_reset:
+            print("JOINT RESET")
+            requests.post(self.url + "jointreset")
+            time.sleep(0.5)
 
-        # perform random reset
+        # Perform Carteasian reset
         if self.randomreset:  # randomize reset position in xy plane
             reset_pose[:2] += np.random.uniform(
                 -self.random_xy_range, self.random_xy_range, (2,)
@@ -42,13 +46,9 @@ def go_to_rest(self, jpos=False):
             )
             reset_pose[3:] = euler_2_quat(euler_random)
             self.interpolate_move(reset_pose, timeout=1.5)
+        else:
+            reset_pose = self.resetpos.copy()
+            self.interpolate_move(reset_pose, timeout=1.5)
 
-        if jpos:
-            requests.post(self.url + "precision_mode")
-            print("JOINT RESET")
-            requests.post(self.url + "jointreset")
-            time.sleep(0.5)
-            self.interpolate_move(self.resetpos, timeout=5)
-
-        requests.post(self.url + "cable_wrap_compliance_mode")
-        return True
+        # Change to compliance mode
+        requests.post(self.url + "update_param", json=self.config.COMPLIANCE_PARAM)

serl_robot_infra/franka_env/envs/franka_env.py

@@ -57,6 +57,8 @@ class DefaultEnvConfig:
     RANDOM_RZ_RANGE = (0.0,)
     ABS_POSE_LIMIT_HIGH = np.zeros((6,))
     ABS_POSE_LIMIT_LOW = np.zeros((6,))
+    COMPLIANCE_PARAM = {}
+    PRECISION_PARAM = {}
 
 
 ##############################################################################
@@ -308,40 +310,35 @@ def _get_obs(self):
 
         return copy.deepcopy(dict(images=images, state=state_observation))
 
-    def interpolate_move(self, goal, timeout, dt=0.1):
-        steps = int(timeout / dt)
+    def interpolate_move(self, goal, timeout):
+        steps = int(timeout * self.hz)
         self.update_currpos()
         path = np.linspace(self.currpos, goal, steps)
         for p in path:
             self._send_pos_command(p)
-            time.sleep(dt)
+            time.sleep(1 / self.hz)
         self.update_currpos()
 
-    def go_to_rest(self, jpos=False):
+    def go_to_rest(self, joint_reset=False):
         """
         The concrete steps to perform reset should be
         implemented each subclass for the specific task.
         """
         raise NotImplementedError
 
-    def reset(self, jpos=False, gripper=None, **kwargs):
+    def reset(self, joint_reset=False, **kwargs):
+        requests.post(self.url + "update_param", json=self.config.COMPLIANCE_PARAM)
         if self.save_video:
             self.save_video_recording()
 
         self.cycle_count += 1
         if self.cycle_count % self.joint_reset_cycle == 0:
             self.cycle_count = 0
-            jpos = True
+            joint_reset = True
 
-        success = self.go_to_rest(jpos=jpos)
-        self.update_currpos()
-        self.curr_path_length = 0
+        self.go_to_rest(joint_reset=joint_reset)
         self.recover()
-
-        if jpos:
-            self.go_to_rest(jpos=False)
-            self.update_currpos()
-            self.recover()
+        self.curr_path_length = 0
 
         self.update_currpos()
         o = self._get_obs()

serl_robot_infra/franka_env/envs/pcb_env/config.py

@@ -46,3 +46,43 @@ class PCBEnvConfig(DefaultEnvConfig):
             TARGET_POSE[5] + RANDOM_RZ_RANGE,
         ]
     )
+    COMPLIANCE_PARAM = {
+        "translational_stiffness": 3000,
+        "translational_damping": 180,
+        "rotational_stiffness": 150,
+        "rotational_damping": 7,
+        "translational_clip_neg_x": 0.002,
+        "translational_clip_neg_y": 0.002,
+        "translational_clip_neg_z": 0.003,
+        "translational_clip_x": 0.0025,
+        "translational_clip_y": 0.0015,
+        "translational_clip_z": 0.002,
+        "rotational_clip_neg_x": 0.025,
+        "rotational_clip_neg_y": 0.007,
+        "rotational_clip_neg_z": 0.01,
+        "rotational_clip_x": 0.025,
+        "rotational_clip_y": 0.007,
+        "rotational_clip_z": 0.01,
+        "translational_Ki": 0,
+        "rotational_Ki": 0,
+    }
+    PRECISION_PARAM = {
+        "translational_stiffness": 3000,
+        "translational_damping": 89,
+        "rotational_stiffness": 300,
+        "rotational_damping": 9,
+        "translational_Ki": 0.1,
+        "translational_clip_x": 0.01,
+        "translational_clip_y": 0.01,
+        "translational_clip_z": 0.01,
+        "translational_clip_neg_x": 0.01,
+        "translational_clip_neg_y": 0.01,
+        "translational_clip_neg_z": 0.01,
+        "rotational_clip_x": 0.05,
+        "rotational_clip_y": 0.05,
+        "rotational_clip_z": 0.05,
+        "rotational_clip_neg_x": 0.05,
+        "rotational_clip_neg_y": 0.05,
+        "rotational_clip_neg_z": 0.05,
+        "rotational_Ki": 0.1,
+    }

serl_robot_infra/franka_env/envs/pcb_env/franka_pcb_insert.py

@@ -22,24 +22,28 @@ def crop_image(self, name, image):
         else:
             return ValueError(f"Camera {name} not recognized in cropping")
 
-    def go_to_rest(self, jpos=False):
-        requests.post(self.url + "pcb_compliance_mode")
+    def go_to_rest(self, joint_reset=False):
         self.update_currpos()
-        self._send_pos_command(self.clip_safety_box(self.currpos))
+        self._send_pos_command(self.currpos)
         time.sleep(0.5)
 
-        requests.post(self.url + "pcb_compliance_mode")
+        # Move up to clear the slot
         self.update_currpos()
         reset_pose = copy.deepcopy(self.currpos)
         reset_pose[2] += 0.03
-        self.interpolate_move(reset_pose, timeout=1.5)
-
-        requests.post(self.url + "precision_mode")
-        time.sleep(1)  # wait for mode switching
-        reset_pose = self.resetpos.copy()
         self.interpolate_move(reset_pose, timeout=1)
 
-        # perform random reset
+        # Change to precision mode for reset
+        requests.post(self.url + "update_param", json=self.config.PRECISION_PARAM)
+        time.sleep(0.5)
+
+        # Perform joint reset if needed
+        if joint_reset:
+            print("JOINT RESET")
+            requests.post(self.url + "jointreset")
+            time.sleep(0.5)
+
+        # Perform Carteasian reset
         if self.randomreset:  # randomize reset position in xy plane
             reset_pose[:2] += np.random.uniform(
                 -self.random_xy_range, self.random_xy_range, (2,)
@@ -50,13 +54,9 @@ def go_to_rest(self, jpos=False):
             )
             reset_pose[3:] = euler_2_quat(euler_random)
             self.interpolate_move(reset_pose, timeout=1.5)
+        else:
+            reset_pose = self.resetpos.copy()
+            self.interpolate_move(reset_pose, timeout=1.5)
 
-        if jpos:
-            requests.post(self.url + "precision_mode")
-            print("JOINT RESET")
-            requests.post(self.url + "jointreset")
-            time.sleep(0.5)
-            self.interpolate_move(self.resetpos, timeout=5)
-
-        requests.post(self.url + "pcb_compliance_mode")
-        return True
+        # Change to compliance mode
+        requests.post(self.url + "update_param", json=self.config.COMPLIANCE_PARAM)

serl_robot_infra/franka_env/envs/peg_env/config.py

@@ -46,3 +46,43 @@ class PegEnvConfig(DefaultEnvConfig):
             TARGET_POSE[5] + RANDOM_RZ_RANGE,
         ]
     )
+    COMPLIANCE_PARAM = {
+        "translational_stiffness": 2000,
+        "translational_damping": 89,
+        "rotational_stiffness": 150,
+        "rotational_damping": 7,
+        "translational_Ki": 0,
+        "translational_clip_x": 0.003,
+        "translational_clip_y": 0.003,
+        "translational_clip_z": 0.01,
+        "translational_clip_neg_x": 0.003,
+        "translational_clip_neg_y": 0.003,
+        "translational_clip_neg_z": 0.01,
+        "rotational_clip_x": 0.02,
+        "rotational_clip_y": 0.02,
+        "rotational_clip_z": 0.02,
+        "rotational_clip_neg_x": 0.02,
+        "rotational_clip_neg_y": 0.02,
+        "rotational_clip_neg_z": 0.02,
+        "rotational_Ki": 0,
+    }
+    PRECISION_PARAM = {
+        "translational_stiffness": 3000,
+        "translational_damping": 89,
+        "rotational_stiffness": 300,
+        "rotational_damping": 9,
+        "translational_Ki": 0.1,
+        "translational_clip_x": 0.01,
+        "translational_clip_y": 0.01,
+        "translational_clip_z": 0.01,
+        "translational_clip_neg_x": 0.01,
+        "translational_clip_neg_y": 0.01,
+        "translational_clip_neg_z": 0.01,
+        "rotational_clip_x": 0.05,
+        "rotational_clip_y": 0.05,
+        "rotational_clip_z": 0.05,
+        "rotational_clip_neg_x": 0.05,
+        "rotational_clip_neg_y": 0.05,
+        "rotational_clip_neg_z": 0.05,
+        "rotational_Ki": 0.1,
+    }

serl_robot_infra/franka_env/envs/peg_env/franka_peg_insert.py

@@ -13,21 +13,28 @@ class FrankaPegInsert(FrankaEnv):
     def __init__(self, **kwargs):
         super().__init__(**kwargs, config=PegEnvConfig)
 
-    def go_to_rest(self, jpos=False):
-        requests.post(self.url + "peg_compliance_mode")
+    def go_to_rest(self, joint_reset=False):
+        self.update_currpos()
+        self._send_pos_command(self.currpos)
+        time.sleep(0.5)
+
+        # Move up to clear the slot
         self.update_currpos()
         reset_pose = copy.deepcopy(self.currpos)
-        reset_pose[2] += 0.03
-        self.interpolate_move(reset_pose, timeout=1.5)
+        reset_pose[2] += 0.10
+        self.interpolate_move(reset_pose, timeout=1)
 
-        # time.sleep(2)
-        requests.post(self.url + "precision_mode")
-        time.sleep(1)  # wait for mode switching
+        # Change to precision mode for reset
+        requests.post(self.url + "update_param", json=self.config.PRECISION_PARAM)
+        time.sleep(0.5)
 
-        reset_pose = self.resetpos.copy()
-        self.interpolate_move(reset_pose, timeout=1)
+        # Perform joint reset if needed
+        if joint_reset:
+            print("JOINT RESET")
+            requests.post(self.url + "jointreset")
+            time.sleep(0.5)
 
-        # perform random reset
+        # Perform Carteasian reset
         if self.randomreset:  # randomize reset position in xy plane
             reset_pose[:2] += np.random.uniform(
                 -self.random_xy_range, self.random_xy_range, (2,)
@@ -37,12 +44,10 @@ def go_to_rest(self, jpos=False):
                 -self.random_rz_range, self.random_rz_range
             )
             reset_pose[3:] = euler_2_quat(euler_random)
-            self.interpolate_move(reset_pose, timeout=1)
-
-        if jpos:
-            requests.post(self.url + "precision_mode")
-            print("JOINT RESET")
-            requests.post(self.url + "jointreset")
+            self.interpolate_move(reset_pose, timeout=1.5)
+        else:
+            reset_pose = self.resetpos.copy()
+            self.interpolate_move(reset_pose, timeout=1.5)
 
-        requests.post(self.url + "peg_compliance_mode")
-        return True
+        # Change to compliance mode
+        requests.post(self.url + "update_param", json=self.config.COMPLIANCE_PARAM)