balance can generate 2 blocks demo

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

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)
 

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,