From 9381551df71adf9f86d8cb3be486da0790aebe6c Mon Sep 17 00:00:00 2001 From: yichao-liang Date: Tue, 14 Jan 2025 10:00:57 +0000 Subject: [PATCH] balance can generate 2 blocks demo --- predicators/envs/pybullet_balance.py | 70 +++++-------------- .../ground_truth_models/balance/options.py | 23 +++--- predicators/settings.py | 2 +- 3 files changed, 35 insertions(+), 60 deletions(-) diff --git a/predicators/envs/pybullet_balance.py b/predicators/envs/pybullet_balance.py index 0fd0107c8..c286157c2 100644 --- a/predicators/envs/pybullet_balance.py +++ b/predicators/envs/pybullet_balance.py @@ -4,8 +4,8 @@ --make_failure_videos --video_fps 20 \ --pybullet_camera_height 900 --pybullet_camera_width 900 --make_test_videos \ --sesame_task_planning_heuristic "goal_count" \ ---excluded_predicates "Balanced,OnPlate" --sesame_max_skeletons_optimized 100 \ ---sesame_check_expected_atoms False +--excluded_predicates "Balanced,OnPlate" --sesame_max_skeletons_optimized 1 \ +--sesame_check_expected_atoms False --pybullet_ik_validate False """ import logging from pathlib import Path @@ -54,7 +54,7 @@ class PyBulletBalanceEnv(PyBulletEnv): _beam2_pose: ClassVar[Pose3D] = (_table_x, (_plate3_pose[1] + _table2_pose[1]) / 2, _plate_z - 4 * _plate_height) - _beam_half_extents = [0.01, 0.15, _plate_height] + _beam_half_extents = [0.01, 0.15, _plate_height/2] # Button on table _button_radius = 0.04 @@ -222,7 +222,7 @@ def initialize_pybullet( plate3_id = create_pybullet_block( (.9, .9, .9, 1), cls._plate_half_extents, - 1.0, + 0.0, 1.0, cls._plate3_pose, cls._table_orientation, @@ -232,7 +232,7 @@ def initialize_pybullet( plate1_id = create_pybullet_block( (.9, .9, .9, 1), cls._plate_half_extents, - 1.0, + 0.0, 1.0, cls._plate1_pose, cls._table_orientation, @@ -243,7 +243,7 @@ def initialize_pybullet( beam1_id = create_pybullet_block( (0.9, 0.9, 0.9, 1), cls._beam_half_extents, - 1.0, + 0.0, 1.0, cls._beam1_pose, cls._table_orientation, @@ -252,21 +252,18 @@ def initialize_pybullet( beam2_id = create_pybullet_block( (0.9, 0.9, 0.9, 1), cls._beam_half_extents, - 1.0, + 0.0, 1.0, cls._beam2_pose, cls._table_orientation, physics_client_id, ) bodies["beam_ids"] = [beam1_id, beam2_id] - cls.fix_plates_and_beams_in_place(physics_client_id, table2_id, plate1_id, - plate3_id, beam1_id, beam2_id) - button_id = create_pybullet_block( cls._button_color_off, - [cls._button_radius] * 3, - 1.0, + [cls._button_radius, cls._button_radius, cls._button_radius/2], + 0.0, 1.0, (cls.button_x, cls.button_y, cls.button_z), cls._table_orientation, @@ -295,35 +292,6 @@ def initialize_pybullet( return physics_client_id, pybullet_robot, bodies - @staticmethod - def fix_plates_and_beams_in_place(physics_client_id, table_id, plate1_id, - plate3_id, beam1_id, beam2_id): - # Doesn't work for some reason - for child_id in [plate1_id, plate3_id, beam1_id, beam2_id]: - parent_pos, parent_orn = p.getBasePositionAndOrientation(table_id, - physicsClientId=physics_client_id) - child_pos, child_orn = p.getBasePositionAndOrientation(child_id, - physicsClientId=physics_client_id) - rel_pos, rel_orn = p.multiplyTransforms( - p.invertTransform(parent_pos, parent_orn)[0], - p.invertTransform(parent_pos, parent_orn)[1], - child_pos, - child_orn - ) - p.createConstraint( - parentBodyUniqueId=table_id, - parentLinkIndex=-1, - childBodyUniqueId=child_id, - childLinkIndex=-1, - jointType=p.JOINT_FIXED, - jointAxis=(0, 0, 0), - parentFramePosition=rel_pos, - parentFrameOrientation=rel_orn, - childFramePosition=(0, 0, 0), - childFrameOrientation=(0, 0, 0), - physicsClientId=physics_client_id - ) - def _store_pybullet_bodies(self, pybullet_bodies: Dict[str, Any]) -> None: self._plate1.id = pybullet_bodies["table_ids"][0] self._plate3.id = pybullet_bodies["table_ids"][1] @@ -402,9 +370,6 @@ def step(self, action: Action, render_obs: bool = False) -> State: state = super().step(action, render_obs=render_obs) self._update_balance_beam(state) - self.fix_plates_and_beams_in_place(self._physics_client_id, self._table_id, - self._plate1.id, self._plate3.id, - self._beam_ids[0], self._beam_ids[1]) # Turn machine on if self._PressingButton_holds(state, [self._robot, self._machine]): @@ -486,9 +451,6 @@ def _reset_state(self, state: State) -> None: self._prev_diff = 0 # Also do one beam update to make sure the initial positions match self._update_balance_beam(state) - self.fix_plates_and_beams_in_place(self._physics_client_id, self._table_id, - self._plate1.id, self._plate3.id, - self._beam_ids[0], self._beam_ids[1]) # Update the button color if self._MachineOn_holds(state, [self._machine, self._robot]): @@ -519,7 +481,7 @@ def _update_balance_beam(self, state: State) -> None: if diff == self._prev_diff: return - shift_per_block = 0.01 + shift_per_block = 0.007 shift_amount = abs(diff) * shift_per_block block_objs = state.get_objects(self._block_type) left_dropping = diff > 0 @@ -586,8 +548,6 @@ def shift_plate(is_left: bool, dropping: bool): shift_plate(False, True) shift_blocks(False, True) - # Right side update - self._prev_diff = diff @@ -907,7 +867,11 @@ def _make_tasks(self, num_tasks: int, possible_num_blocks: List[int], init_state = self._sample_state_from_piles(piles, rng) goal = { GroundAtom(self._MachineOn, [self._machine, self._robot]), - # GroundAtom(self._DirectlyOn, [piles[1][3], piles[0][1]]), + # GroundAtom(self._DirectlyOn, [piles[1][4], piles[0][0]]), + # GroundAtom(self._DirectlyOn, [piles[1][3], piles[1][4]]), + # GroundAtom(self._DirectlyOn, [piles[0][4], piles[0][5]]), + # GroundAtom(self._DirectlyOn, [piles[0][3], piles[0][4]]), + # GroundAtom(self._DirectlyOn, [piles[0][2], piles[0][3]]), } # } # while True: # repeat until goal is not satisfied @@ -940,6 +904,10 @@ def _sample_initial_piles(self, num_blocks: int, if (block_num == 0 or block_num == 1): n_piles += 1 piles.append([]) + # For generating a 0:6 pile + # if (block_num == 0): + # n_piles += 1 + # piles.append([]) # Add block to pile piles[-1].append(block) return piles diff --git a/predicators/ground_truth_models/balance/options.py b/predicators/ground_truth_models/balance/options.py index baf63a043..cf0945d85 100644 --- a/predicators/ground_truth_models/balance/options.py +++ b/predicators/ground_truth_models/balance/options.py @@ -31,6 +31,7 @@ class PyBulletBalanceGroundTruthOptionFactory(GroundTruthOptionFactory): _move_to_pose_tol: ClassVar[float] = 1e-4 _finger_action_nudge_magnitude: ClassVar[float] = 1e-3 _offset_z: ClassVar[float] = 0.01 + _transport_z: ClassVar[float] = env_cls.z_ub - 0.2 @classmethod def get_env_names(cls) -> Set[str]: @@ -82,7 +83,7 @@ def close_fingers_func(state: State, objects: Sequence[Object], # Move to far above the block which we will grasp. cls._create_blocks_move_to_above_block_option( name="MoveEndEffectorToPreGrasp", - z_func=lambda _: cls.env_cls.z_ub, + z_func=lambda _: cls._transport_z, finger_status="open", pybullet_robot=pybullet_robot, option_types=option_types, @@ -103,11 +104,13 @@ def close_fingers_func(state: State, objects: Sequence[Object], # Move back up. cls._create_blocks_move_to_above_block_option( name="MoveEndEffectorBackUp", - z_func=lambda _: cls.env_cls.z_ub, + z_func=lambda _: cls._transport_z, finger_status="closed", pybullet_robot=pybullet_robot, option_types=option_types, - params_space=params_space), + params_space=params_space, + move_to_pose_tol=cls._move_to_pose_tol*100 + ), ], # "Pick up block ?block" ) @@ -121,7 +124,7 @@ def close_fingers_func(state: State, objects: Sequence[Object], # Move to above the block on which we will stack. cls._create_blocks_move_to_above_block_option( name="MoveEndEffectorToPreStack", - z_func=lambda _: cls.env_cls.z_ub, + z_func=lambda _: cls._transport_z, finger_status="closed", pybullet_robot=pybullet_robot, option_types=option_types, @@ -142,11 +145,13 @@ def close_fingers_func(state: State, objects: Sequence[Object], # Move back up. cls._create_blocks_move_to_above_block_option( name="MoveEndEffectorBackUp", - z_func=lambda _: cls.env_cls.z_ub, + z_func=lambda _: cls._transport_z, finger_status="open", pybullet_robot=pybullet_robot, option_types=option_types, - params_space=params_space), + params_space=params_space, + move_to_pose_tol=cls._move_to_pose_tol*100 + ), ], # annotation="Stack the block in hand onto block ?otherblock" ) @@ -228,7 +233,9 @@ def _create_blocks_move_to_above_block_option( cls, name: str, z_func: Callable[[float], float], finger_status: str, pybullet_robot: SingleArmPyBulletRobot, option_types: List[Type], - params_space: Box) -> ParameterizedOption: + params_space: Box, + move_to_pose_tol: float = _move_to_pose_tol, + ) -> ParameterizedOption: """Creates a ParameterizedOption for moving to a pose above that of the block argument. @@ -255,7 +262,7 @@ def _get_current_and_target_pose_and_finger_status( return create_move_end_effector_to_pose_option( pybullet_robot, name, option_types, params_space, _get_current_and_target_pose_and_finger_status, - cls._move_to_pose_tol, CFG.pybullet_max_vel_norm, + move_to_pose_tol, CFG.pybullet_max_vel_norm, cls._finger_action_nudge_magnitude, validate=CFG.pybullet_ik_validate) diff --git a/predicators/settings.py b/predicators/settings.py index 9da6f55c0..1d32b090d 100644 --- a/predicators/settings.py +++ b/predicators/settings.py @@ -764,7 +764,7 @@ def get_arg_specific_settings(cls, args: Dict[str, Any]) -> Dict[str, Any]: "pybullet_cover": 1000, "pybullet_blocks": 1000, "pybullet_coffee": 2000, - "pybullet_coffee_pixel": 2000, + "pybullet_balance": 2000, "pybullet_grow": 2000, "pybullet_circuit": 2000, "pybullet_float": 2000,