From 9001935a5bf7a7cee39eea347e2f1fd156823771 Mon Sep 17 00:00:00 2001
From: minoring <minho1350@gmail.com>
Date: Thu, 2 May 2024 20:36:06 +0900
Subject: [PATCH] implement positional control for gripper action.

This can address levitation issues for the tabletop (e.g., scripted
agent).
---
 furniture_bench/data/data_collector.py    |  4 ++-
 furniture_bench/envs/furniture_sim_env.py | 30 +++++++++++++++++------
 furniture_bench/scripts/collect_data.py   |  3 +++
 3 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/furniture_bench/data/data_collector.py b/furniture_bench/data/data_collector.py
index 3f46ef9c..6017027e 100644
--- a/furniture_bench/data/data_collector.py
+++ b/furniture_bench/data/data_collector.py
@@ -39,6 +39,7 @@ def __init__(
         save_failure: bool = False,
         num_demos: int = 100,
         resize_sim_img: bool = False,
+        gripper_pos_control: bool = False,
     ):
         """
         Args:
@@ -73,7 +74,8 @@ def __init__(
                 channel_first=False,
                 randomness=randomness,
                 compute_device_id=compute_device_id,
-                graphics_device_id=graphics_device_id
+                graphics_device_id=graphics_device_id,
+                gripper_pos_control=gripper_pos_control,
             )
         else:
             if randomness == "med":
diff --git a/furniture_bench/envs/furniture_sim_env.py b/furniture_bench/envs/furniture_sim_env.py
index 82f23452..fc284e68 100644
--- a/furniture_bench/envs/furniture_sim_env.py
+++ b/furniture_bench/envs/furniture_sim_env.py
@@ -134,6 +134,7 @@ def __init__(
         self.last_grasp = torch.tensor([-1.0] * num_envs, device=self.device)
         self.grasp_margin = 0.02 - 0.001  # To prevent repeating open an close actions.
         self.max_gripper_width = config["robot"]["max_gripper_width"][furniture]
+        self.gripper_pos_control = kwargs.get("gripper_pos_control", False)
 
         self.save_camera_input = save_camera_input
         self.img_size = sim_config["camera"][
@@ -352,10 +353,15 @@ def create_envs(self):
             franka_dof_props["damping"][:7].fill(0.0)
             franka_dof_props["friction"][:7] = sim_config["robot"]["arm_frictions"]
             # Grippers
-            franka_dof_props["driveMode"][7:].fill(gymapi.DOF_MODE_EFFORT)
-            franka_dof_props["stiffness"][7:].fill(0)
-            franka_dof_props["damping"][7:].fill(0)
-            franka_dof_props["friction"][7:] = sim_config["robot"]["gripper_frictions"]
+            if self.gripper_pos_control:
+                franka_dof_props["driveMode"][7:].fill(gymapi.DOF_MODE_POS)
+                franka_dof_props["stiffness"][7:].fill(200.0)
+                franka_dof_props["damping"][7:].fill(60.0)
+            else:
+                franka_dof_props["driveMode"][7:].fill(gymapi.DOF_MODE_EFFORT)
+                franka_dof_props["stiffness"][7:].fill(0)
+                franka_dof_props["damping"][7:].fill(0)
+                franka_dof_props["friction"][7:] = sim_config["robot"]["gripper_frictions"]
             franka_dof_props["upper"][7:] = self.max_gripper_width / 2
 
             self.isaac_gym.set_actor_dof_properties(
@@ -786,11 +792,19 @@ def step(self, action):
                     "joint_torques"
                 ]
 
-                if grip_sep > 0:
-                    torque_action[env_idx, 7:9] = sim_config["robot"]["gripper_torque"]
+                if self.gripper_pos_control:
+                    grip_action[env_idx, -1] = grip_sep
                 else:
-                    torque_action[env_idx, 7:9] = -sim_config["robot"]["gripper_torque"]
-
+                    if grip_sep > 0:
+                        torque_action[env_idx, 7:9] = sim_config["robot"]["gripper_torque"]
+                    else:
+                        torque_action[env_idx, 7:9] = -sim_config["robot"]["gripper_torque"]
+            # Gripper action
+            if self.gripper_pos_control:
+                pos_action[:, 7:9] = grip_action
+                self.isaac_gym.set_dof_position_target_tensor(
+                    self.sim, gymtorch.unwrap_tensor(pos_action)
+                )
             self.isaac_gym.set_dof_actuation_force_tensor(
                 self.sim, gymtorch.unwrap_tensor(torque_action)
             )
diff --git a/furniture_bench/scripts/collect_data.py b/furniture_bench/scripts/collect_data.py
index 2cba55ff..73ae1f2e 100644
--- a/furniture_bench/scripts/collect_data.py
+++ b/furniture_bench/scripts/collect_data.py
@@ -75,6 +75,8 @@ def main():
     parser.add_argument("--num-demos", default=100, type=int)
 
     parser.add_argument("--resize-sim-img", action="store_true")
+    parser.add_argument("--gripper-pos-control", action="store_true")
+
 
     args = parser.parse_args()
 
@@ -104,6 +106,7 @@ def main():
         save_failure=args.save_failure,
         num_demos=args.num_demos,
         resize_sim_img=args.resize_sim_img,
+        gripper_pos_control=args.gripper_pos_control,
     )
     data_collector.collect()