yapf

Author: yichao-liang

predicators/envs/pybullet_domino/__init__.py

@@ -11,10 +11,8 @@
     env = PyBulletDominoFanEnv(use_gui=True)
 """
 
-from predicators.envs.pybullet_domino.composed_env import (
-    PyBulletDominoEnvNew,
-    PyBulletDominoFanEnvNew,
-)
+from predicators.envs.pybullet_domino.composed_env import \
+    PyBulletDominoEnvNew, PyBulletDominoFanEnvNew
 
 # Backward-compatible aliases
 PyBulletDominoEnv = PyBulletDominoEnvNew

predicators/envs/pybullet_domino/components/__init__.py

@@ -1,18 +1,18 @@
 """Domino environment components.
 
 Each component encapsulates a specific aspect of the domino environment
-(e.g., dominoes, fans, balls, ramps) and can be composed together to create
-different environment variants.
+(e.g., dominoes, fans, balls, ramps) and can be composed together to
+create different environment variants.
 """
 
+from predicators.envs.pybullet_domino.components.ball_component import \
+    BallComponent
 from predicators.envs.pybullet_domino.components.base_component import \
     DominoEnvComponent
 from predicators.envs.pybullet_domino.components.domino_component import \
     DominoComponent
 from predicators.envs.pybullet_domino.components.fan_component import \
     FanComponent
-from predicators.envs.pybullet_domino.components.ball_component import \
-    BallComponent
 
 __all__ = [
     "DominoEnvComponent",

predicators/envs/pybullet_domino/components/ball_component.py

@@ -39,17 +39,17 @@ class BallComponent(DominoEnvComponent):
     ball_restitution: ClassVar[float] = 0.3
     ball_linear_damping: ClassVar[float] = 0.5
     ball_angular_damping: ClassVar[float] = 0.3
-    ball_color: ClassVar[Tuple[float, float, float, float]] = (0.0, 0.0, 1.0,
-                                                               1.0)
+    ball_color: ClassVar[Tuple[float, float, float,
+                               float]] = (0.0, 0.0, 1.0, 1.0)
 
     # =========================================================================
     # TARGET CONFIGURATION
     # =========================================================================
     target_thickness: ClassVar[float] = 0.00001
     target_mass: ClassVar[float] = 0.0
     target_friction: ClassVar[float] = 0.04
-    target_color: ClassVar[Tuple[float, float, float, float]] = (0.0, 1.0, 0.0,
-                                                                 1.0)
+    target_color: ClassVar[Tuple[float, float, float,
+                                 float]] = (0.0, 1.0, 0.0, 1.0)
 
     def __init__(self,
                  workspace_bounds: Optional[Dict[str, float]] = None,
@@ -267,7 +267,8 @@ def get_init_dict_entries(
             min_distance: float = 0.15
             target_x = rng.uniform(self.x_lb + 0.05, self.x_ub - 0.05)
             target_y = rng.uniform(self.y_lb + 0.05, self.y_ub - 0.05)
-            while np.sqrt((target_x - ball_x)**2 + (target_y - ball_y)**2) < min_distance:
+            while np.sqrt((target_x - ball_x)**2 +
+                          (target_y - ball_y)**2) < min_distance:
                 target_x = rng.uniform(self.x_lb + 0.05, self.x_ub - 0.05)
                 target_y = rng.uniform(self.y_lb + 0.05, self.y_ub - 0.05)
         else:

predicators/envs/pybullet_domino/components/base_component.py

@@ -19,14 +19,16 @@
 class DominoEnvComponent(ABC):
     """Abstract base class for all domino environment components.
 
-    Components encapsulate specific functionality (e.g., dominoes, fans, balls)
-    and can be composed together to create different environment configurations.
+    Components encapsulate specific functionality (e.g., dominoes, fans,
+    balls) and can be composed together to create different environment
+    configurations.
     """
 
     def __init__(self) -> None:
         """Initialize the component.
 
-        Subclasses should create their types, predicates, and objects here.
+        Subclasses should create their types, predicates, and objects
+        here.
         """
         self._physics_client_id: Optional[int] = None
 
@@ -46,7 +48,8 @@ def get_types(self) -> Set[Type]:
     def get_predicates(self) -> Set[Predicate]:
         """Return the predicates introduced by this component.
 
-        These predicates will be added to the environment's predicate set.
+        These predicates will be added to the environment's predicate
+        set.
         """
         raise NotImplementedError
 
@@ -123,15 +126,16 @@ def extract_feature(self, obj: Object, feature: str) -> Optional[float]:
     def set_physics_client_id(self, physics_client_id: int) -> None:
         """Set the physics client ID for this component.
 
-        Called by the composed environment after PyBullet initialization.
+        Called by the composed environment after PyBullet
+        initialization.
         """
         self._physics_client_id = physics_client_id
 
     def step(self) -> None:
         """Called each simulation step.
 
-        Override this method to add per-step physics updates (e.g., wind forces
-        from fans). By default, does nothing.
+        Override this method to add per-step physics updates (e.g., wind
+        forces from fans). By default, does nothing.
         """
         pass
 

predicators/envs/pybullet_domino/components/domino_component.py

@@ -410,10 +410,10 @@ def _StartBlock_holds(self, state: State,
         """Check if domino is the start block (light green)."""
         domino, = objects
         eps = 1e-3
-        return (abs(state.get(domino, "r") - self.start_domino_color[0]) < eps
-                and abs(state.get(domino, "g") - self.start_domino_color[1]) <
-                eps and abs(state.get(domino, "b") -
-                            self.start_domino_color[2]) < eps)
+        return (
+            abs(state.get(domino, "r") - self.start_domino_color[0]) < eps
+            and abs(state.get(domino, "g") - self.start_domino_color[1]) < eps
+            and abs(state.get(domino, "b") - self.start_domino_color[2]) < eps)
 
     def _MovableBlock_holds(self, state: State,
                             objects: Sequence[Object]) -> bool:
@@ -425,14 +425,14 @@ def _MovableBlock_holds(self, state: State,
                 and abs(state.get(domino, "b") - self.domino_color[2]) < eps)
 
     def _TargetDomino_holds(self, state: State,
-                           objects: Sequence[Object]) -> bool:
+                            objects: Sequence[Object]) -> bool:
         """Check if domino is a target (pink or glued red)."""
         domino, = objects
         eps = 1e-3
         return (self._DominoGlued_holds(state, objects)) or (
-            abs(state.get(domino, "r") - self.target_domino_color[0]) < eps
-            and abs(state.get(domino, "g") - self.target_domino_color[1]) < eps
-            and abs(state.get(domino, "b") - self.target_domino_color[2]) < eps)
+            abs(state.get(domino, "r") - self.target_domino_color[0]) < eps and
+            abs(state.get(domino, "g") - self.target_domino_color[1]) < eps and
+            abs(state.get(domino, "b") - self.target_domino_color[2]) < eps)
 
     def _DominoNotGlued_holds(self, state: State,
                               objects: Sequence[Object]) -> bool:
@@ -513,7 +513,9 @@ def place_domino(self,
             "is_held": 0.0,
         }
 
-    def place_pivot_or_target(self, x: float, y: float,
+    def place_pivot_or_target(self,
+                              x: float,
+                              y: float,
                               rot: float = 0.0) -> Dict:
         """Create a dictionary with placement parameters for a pivot/target."""
         return {
@@ -549,23 +551,23 @@ def Toppled(self) -> Predicate:
 
 
 def create_domino_block(
-        color: Tuple[float, float, float, float],
-        half_extents: Tuple[float, float, float],
-        mass: float,
-        friction: float,
-        position: Pose3D = (0.0, 0.0, 0.0),
-        orientation: Quaternion = (0.0, 0.0, 0.0, 1.0),
-        physics_client_id: int = 0,
-        add_top_triangle: bool = False,
-        *,
-        restitution: float = 0.02,
-        rolling_friction: float = 0.006,
-        spinning_friction: Optional[float] = None,
-        linear_damping: float = 0.0,
-        angular_damping: float = 0.03,
-        friction_anchor: bool = True,
-        ccd: bool = True,
-        ccd_swept_radius: Optional[float] = None,
+    color: Tuple[float, float, float, float],
+    half_extents: Tuple[float, float, float],
+    mass: float,
+    friction: float,
+    position: Pose3D = (0.0, 0.0, 0.0),
+    orientation: Quaternion = (0.0, 0.0, 0.0, 1.0),
+    physics_client_id: int = 0,
+    add_top_triangle: bool = False,
+    *,
+    restitution: float = 0.02,
+    rolling_friction: float = 0.006,
+    spinning_friction: Optional[float] = None,
+    linear_damping: float = 0.0,
+    angular_damping: float = 0.03,
+    friction_anchor: bool = True,
+    ccd: bool = True,
+    ccd_swept_radius: Optional[float] = None,
 ) -> int:
     """Create a domino-tuned block with appropriate physics settings."""
     block_id = create_pybullet_block(

predicators/envs/pybullet_domino/components/fan_component.py

@@ -239,10 +239,8 @@ def initialize_pybullet(self, physics_client_id: int) -> Dict[str, Any]:
     def store_pybullet_bodies(self, pybullet_bodies: Dict[str, Any]) -> None:
         """Store PyBullet body IDs on fan and switch objects."""
         fan_ids_by_side = [
-            pybullet_bodies["fan_ids_left"],
-            pybullet_bodies["fan_ids_right"],
-            pybullet_bodies["fan_ids_back"],
-            pybullet_bodies["fan_ids_front"]
+            pybullet_bodies["fan_ids_left"], pybullet_bodies["fan_ids_right"],
+            pybullet_bodies["fan_ids_back"], pybullet_bodies["fan_ids_front"]
         ]
 
         for side_idx, fan_obj in enumerate(self._fans):
@@ -487,7 +485,8 @@ def _Controls_holds(self, state: State, objects: Sequence[Object]) -> bool:
     # -------------------------------------------------------------------------
 
     def get_init_dict_entries(
-            self, rng: "np.random.Generator",
+            self,
+            rng: "np.random.Generator",
             all_off: bool = True) -> Dict[Object, Dict[str, Any]]:
         """Return initial state dict entries for fans, switches, and sides."""
         init_dict: Dict[Object, Dict[str, Any]] = {}

predicators/envs/pybullet_domino/composed_env.py

@@ -10,16 +10,16 @@
 import pybullet as p
 
 from predicators import utils
-from predicators.envs.pybullet_domino.components.base_component import \
-    DominoEnvComponent
 from predicators.envs.pybullet_domino.components.ball_component import \
     BallComponent
+from predicators.envs.pybullet_domino.components.base_component import \
+    DominoEnvComponent
 from predicators.envs.pybullet_domino.components.domino_component import \
     DominoComponent
 from predicators.envs.pybullet_domino.components.fan_component import \
     FanComponent
-from predicators.envs.pybullet_domino.task_generators.domino_task_generator \
-    import DominoTaskGenerator
+from predicators.envs.pybullet_domino.task_generators.domino_task_generator import \
+    DominoTaskGenerator
 from predicators.envs.pybullet_env import PyBulletEnv
 from predicators.pybullet_helpers.camera import create_gui_connection
 from predicators.pybullet_helpers.geometry import Pose, Pose3D, Quaternion
@@ -201,9 +201,8 @@ def initialize_pybullet(
 
         # Load robot
         robot_ee_orn = cls.get_robot_ee_home_orn()
-        ee_home = Pose(
-            (cls.robot_init_x, cls.robot_init_y, cls.robot_init_z),
-            robot_ee_orn)
+        ee_home = Pose((cls.robot_init_x, cls.robot_init_y, cls.robot_init_z),
+                       robot_ee_orn)
         base_pose = Pose(cls.robot_base_pos,
                          cls.robot_base_orn) if cls.robot_base_pos else None
         pybullet_robot = create_single_arm_pybullet_robot(
@@ -221,14 +220,14 @@ def initialize_pybullet(
                                  physics_client_id=physics_client_id)
 
         # Add second table for more space
-        create_object(
-            asset_path="urdf/table.urdf",
-            position=(cls.table_pos[0], cls.table_pos[1] + cls.table_width / 2,
-                      cls.table_pos[2]),
-            orientation=cls.table_orn,
-            scale=1.0,
-            use_fixed_base=True,
-            physics_client_id=physics_client_id)
+        create_object(asset_path="urdf/table.urdf",
+                      position=(cls.table_pos[0],
+                                cls.table_pos[1] + cls.table_width / 2,
+                                cls.table_pos[2]),
+                      orientation=cls.table_orn,
+                      scale=1.0,
+                      use_fixed_base=True,
+                      physics_client_id=physics_client_id)
 
         bodies = {"table_id": table_id}
         return physics_client_id, pybullet_robot, bodies

predicators/envs/pybullet_domino/task_generators/__init__.py

@@ -1,13 +1,14 @@
 """Task generators for the domino environment.
 
-Task generators create EnvironmentTask instances with initial states and goals.
-Different generators can be composed with different component combinations.
+Task generators create EnvironmentTask instances with initial states and
+goals. Different generators can be composed with different component
+combinations.
 """
 
 from predicators.envs.pybullet_domino.task_generators.base_generator import \
     TaskGenerator
-from predicators.envs.pybullet_domino.task_generators.domino_task_generator \
-    import DominoTaskGenerator
+from predicators.envs.pybullet_domino.task_generators.domino_task_generator import \
+    DominoTaskGenerator
 
 __all__ = [
     "TaskGenerator",