Added support for controlling the gripper with the pygame interface of gamepad

Minor modifications in gym_manipulator to quantize the gripper actions
clamped the observations after F.resize in ConvertToLeRobotObservation wrapper due to a bug in F.resize, images were returned exceeding the maximum value of 1.0
This commit is contained in:
Michel Aractingi 2025-03-28 18:12:03 +01:00
parent 4d5ecb082e
commit 8eb3c1510c
2 changed files with 26 additions and 11 deletions

View File

@ -258,11 +258,25 @@ class GamepadController(InputController):
elif event.button == 0:
self.episode_end_status = "rerecord_episode"
# RB button (6) for opening gripper
elif event.button == 6:
self.open_gripper_command = True
# LT button (7) for closing gripper
elif event.button == 7:
self.close_gripper_command = True
# Reset episode status on button release
elif event.type == pygame.JOYBUTTONUP:
if event.button in [0, 2, 3]:
self.episode_end_status = None
elif event.button == 6:
self.open_gripper_command = False
elif event.button == 7:
self.close_gripper_command = False
# Check for RB button (typically button 5) for intervention flag
if self.joystick.get_button(5):
self.intervention_flag = True
@ -298,10 +312,6 @@ class GamepadController(InputController):
logging.error("Error reading gamepad. Is it still connected?")
return 0.0, 0.0, 0.0
def should_intervene(self):
"""Return True if intervention flag was set."""
return self.intervention_flag
class GamepadControllerHID(InputController):
"""Generate motion deltas from gamepad input using HIDAPI."""
@ -471,10 +481,6 @@ class GamepadControllerHID(InputController):
"""Return True if save button was pressed."""
return self.save_requested
def should_intervene(self):
"""Return True if intervention flag was set."""
return self.intervention_flag
def test_forward_kinematics(robot, fps=10):
logging.info("Testing Forward Kinematics")

View File

@ -552,6 +552,8 @@ class ImageCropResizeWrapper(gym.Wrapper):
obs[k] = F.crop(obs[k], *self.crop_params_dict[k])
obs[k] = F.resize(obs[k], self.resize_size)
# TODO(michel-aractingi): Bug in resize, it returns values outside [0, 1]
obs[k] = obs[k].clamp(0.0, 1.0)
# Check for NaNs after processing
if torch.isnan(obs[k]).any():
@ -569,6 +571,7 @@ class ImageCropResizeWrapper(gym.Wrapper):
obs[k] = obs[k].cpu()
obs[k] = F.crop(obs[k], *self.crop_params_dict[k])
obs[k] = F.resize(obs[k], self.resize_size)
obs[k] = obs[k].clamp(0.0, 1.0)
obs[k] = obs[k].to(device)
return obs, info
@ -588,7 +591,7 @@ class ConvertToLeRobotObservation(gym.ObservationWrapper):
key: observation[key].to(self.device, non_blocking=self.device.type == "cuda")
for key in observation
}
observation = {k: torch.tensor(v, device=self.device) for k, v in observation.items()}
return observation
@ -817,7 +820,14 @@ class EEActionWrapper(gym.ActionWrapper):
fk_func=self.fk_function,
)
if self.use_gripper:
gripper_command = gripper_command * MAX_GRIPPER_COMMAND
# Quantize gripper command to -1, 0 or 1
if gripper_command < -0.2:
gripper_command = -1.0
elif gripper_command > 0.2:
gripper_command = 1.0
else:
gripper_command = 0.0
gripper_state = self.unwrapped.robot.follower_arms["main"].read("Present_Position")[-1]
gripper_action = np.clip(gripper_state + gripper_command, 0, MAX_GRIPPER_COMMAND)
target_joint_pos[-1] = gripper_action
@ -1248,7 +1258,6 @@ def record_dataset(env, policy, cfg):
# Process observation for dataset
obs = {k: v.cpu().squeeze(0).float() for k, v in obs.items()}
obs["observation.images.side"] = torch.clamp(obs["observation.images.side"], 0, 1)
# Add frame to dataset
frame = {**obs, **recorded_action}