src/robomme/robomme_env/utils/object_generation.py

import numpy as np
import torch
import sapien
from typing import Optional, Tuple, Sequence, Union
from mani_skill.utils.structs.pose import Pose
from mani_skill.utils.geometry.rotation_conversions import (
    euler_angles_to_matrix,
    matrix_to_quaternion,
)
import mani_skill.envs.utils.randomization as randomization  # Only used if needed elsewhere
from mani_skill.examples.motionplanning.base_motionplanner.utils import (
    compute_grasp_info_by_obb,
    get_actor_obb,
)
from mani_skill.utils.building import actors
from mani_skill.utils.geometry.rotation_conversions import (
    euler_angles_to_matrix,
    matrix_to_quaternion,
)
from transforms3d.euler import euler2quat
from mani_skill.utils import sapien_utils
from mani_skill.envs.scene import ManiSkillScene
from mani_skill.utils.building.actor_builder import ActorBuilder
from mani_skill.utils.structs.pose import Pose
from mani_skill.utils.structs.types import Array
from typing import Optional, Union

def _color_to_rgba(color: Union[str, Sequence[float]]) -> Tuple[float, float, float, float]:
    """Convert a hex string or RGB/RGBA tuple to an RGBA tuple accepted by SAPIEN."""
    if isinstance(color, str):
        return sapien_utils.hex2rgba(color)
    if len(color) == 3:
        return (float(color[0]), float(color[1]), float(color[2]), 1.0)
    if len(color) == 4:
        return tuple(float(c) for c in color)
    raise ValueError("color must be a hex string or a sequence of 3/4 floats")


def build_peg(
    env_or_scene,
    length: float,
    radius: float,
    *,
    initial_pose: Optional["sapien.Pose"] = None,
    head_color: str = "#EC7357",
    tail_color: str = "#F5F5F5",
    density: float = 1200.0,
    name: str = "peg",
) -> Tuple["sapien.Articulation", "sapien.Link", "sapien.Link"]:
    """Construct a peg articulation with head and tail links tied by a fixed joint.

    Args:
        env_or_scene: Environment or scene providing `create_articulation_builder`.
        length: Total length of the peg (meters).
        radius: Half-width of the rectangular cross section (meters).
        initial_pose: Optional pose for the articulation root; defaults to placing
            the head centered at positive x.
        head_color: Hex color for the head visual.
        tail_color: Hex color for the tail visual.
        density: Collision density (kg/m^3) shared by both links.
        name: Name assigned to the articulation.

    Returns:
        The articulation along with the head and tail links.
    """

    scene = getattr(env_or_scene, "scene", env_or_scene)
    if initial_pose is None:
        initial_pose = sapien.Pose(p=[length / 2, 0.0, radius], q=[1, 0, 0, 0])

    builder = scene.create_articulation_builder()
    builder.initial_pose = initial_pose

    head_builder = builder.create_link_builder()
    head_builder.set_name("peg_head")
    head_builder.add_box_collision(
        half_size=[length / 2 * 0.9, radius, radius], density=density
    )
    head_material = sapien.render.RenderMaterial(
        base_color=_color_to_rgba(head_color),
        roughness=0.5,
        specular=0.5,
    )
    head_builder.add_box_visual(
        half_size=[length / 2, radius, radius],
        material=head_material,
    )

    tail_builder = builder.create_link_builder(head_builder)
    tail_builder.set_name("peg_tail")
    tail_builder.set_joint_name("peg_fixed_joint")
    tail_builder.set_joint_properties(
        type="fixed",
        limits=[[0.0, 0.0]],
        pose_in_parent=sapien.Pose(p=[-length, 0.0, 0.0], q=[1, 0, 0, 0]),
        pose_in_child=sapien.Pose(p=[0.0, 0.0, 0.0], q=[1, 0, 0, 0]),
        friction=0.0,
        damping=0.0,
    )
    tail_builder.add_box_collision(
        half_size=[length / 2 * 0.9, radius, radius], density=density
    )
    tail_material = sapien.render.RenderMaterial(
        base_color=_color_to_rgba(tail_color),
        roughness=0.5,
        specular=0.5,
    )
    tail_builder.add_box_visual(
        half_size=[length / 2, radius, radius],
        material=tail_material,
    )

    peg = builder.build(name=name, fix_root_link=False)
    link_map = {link.get_name(): link for link in peg.get_links()}
    peg_head = link_map["peg_head"]
    peg_tail = link_map["peg_tail"]
    return peg, peg_head, peg_tail


def build_box_with_hole(self, inner_radius, outer_radius, depth, center=(0, 0)):
    builder = self.scene.create_actor_builder()
    thickness = (outer_radius - inner_radius) * 0.5
    # x-axis is hole direction
    half_center = [x * 0.5 for x in center]
    half_sizes = [
        [depth, thickness - half_center[0], outer_radius],
        [depth, thickness + half_center[0], outer_radius],
        [depth, outer_radius, thickness - half_center[1]],
        [depth, outer_radius, thickness + half_center[1]],
    ]
    offset = thickness + inner_radius
    poses = [
        sapien.Pose([0, offset + half_center[0], 0]),
        sapien.Pose([0, -offset + half_center[0], 0]),
        sapien.Pose([0, 0, offset + half_center[1]]),
        sapien.Pose([0, 0, -offset + half_center[1]]),
    ]

    mat = sapien.render.RenderMaterial(
        base_color=sapien_utils.hex2rgba("#FFD289"), roughness=0.5, specular=0.5
    )

    for half_size, pose in zip(half_sizes, poses):
        builder.add_box_collision(pose, half_size)
        builder.add_box_visual(pose, half_size, material=mat)
    box=builder.build_kinematic(f"box_with_hole")
    return box
def _safe_unit(v, eps=1e-12):
    n = np.linalg.norm(v)
    if n < eps:
        return v
    return v / n

def _trimesh_box_to_obb2d(obb_box, extra_pad=0.0):
    """
    Convert trimesh.primitives.Box (world frame) to 2D OBB representation: center c(2,), axes A(2x2), half-extents h(2,)
    extra_pad: Margins to expand outward on XY plane (meters)
    """
    # Compatible with obb potentially wrapped in .primitive
    b = getattr(obb_box, "primitive", obb_box)
    T = np.asarray(b.transform, dtype=np.float64)  # 4x4
    ex = np.asarray(b.extents, dtype=np.float64)   # 3

    R = T[:3, :3]
    t = T[:3, 3]

    c = t[:2].copy()

    # Take projection of X, Y axes on plane as two axes of 2D OBB
    u = _safe_unit(R[:2, 0])  # x-axis projection
    v = _safe_unit(R[:2, 1])  # y-axis projection
    A = np.stack([u, v], axis=1)  # 2x2, each column is an axis

    h = 0.5 * ex[:2].astype(np.float64)
    if extra_pad > 0:
        h = h + float(extra_pad)
    return c, A, h

def _obb2d_intersect(c1, A1, h1, c2, A2, h2):
    """
    2D OBB SAT detection. c*: (2,), A*: (2x2) columns are axes, h*: (2,)
    Returns True indicating intersection (including contact), False indicating separation
    """
    d = c2 - c1
    axes = [A1[:, 0], A1[:, 1], A2[:, 0], A2[:, 1]]

    for a in axes:
        a = _safe_unit(a)
        # Projected radius
        r1 = abs(np.dot(A1[:, 0], a)) * h1[0] + abs(np.dot(A1[:, 1], a)) * h1[1]
        r2 = abs(np.dot(A2[:, 0], a)) * h2[0] + abs(np.dot(A2[:, 1], a)) * h2[1]
        dist = abs(np.dot(d, a))
        if dist > (r1 + r2):
            return False  # Separating axis exists -> No intersection
    return True  # All axes overlap -> Intersection/Contact

def _yaw_to_quat_tensor(yaw: float, device):
    """
    Get quaternion consistent with ManiSkill/your conversion tools using z-axis Euler angle (shape [1,4], float32, device aligned)
    """
    # euler_angles_to_matrix accepts [roll, pitch, yaw] (radians), returns Nx3x3
    angles = torch.tensor([[0.0, 0.0, float(yaw)]], dtype=torch.float32, device=device)
    R = euler_angles_to_matrix(angles,convention="XYZ")            # (1, 3, 3)
    q = matrix_to_quaternion(R)                   # (1, 4) Convention same as ManiSkill
    return q

def _build_new_cube_obb2d(x, y, half_size_xy, yaw, pad_xy=0.0):
    """
    Construct 2D OBB for "cube ready to be placed": center/axes/half-extents
    half_size_xy: float, half length of cube on XY
    yaw: rotation around z-axis (radians)
    pad_xy: extra padding on half length on XY (for minimum gap)
    """
    c = np.array([x, y], dtype=np.float64)
    cos_y = np.cos(yaw)
    sin_y = np.sin(yaw)
    A = np.array([[cos_y, -sin_y],
                  [sin_y,  cos_y]], dtype=np.float64)  # Columns are axes
    h = np.array([half_size_xy + pad_xy, half_size_xy + pad_xy], dtype=np.float64)
    return c, A, h

def spawn_random_cube(
        self,
        region_center=[0, 0],
        region_half_size=0.1,
        half_size=0.01,
        color=(1, 0, 0, 1),
        name_prefix="cube_extra",
        min_gap=0.005,
        max_trials=256,
        avoid=None,
        random_yaw=True,
        include_existing=True,
        include_goal=True,
        generator=None
    ):
    """
    Drop a cube (onto table) in rectangular region using rejection sampling, and return the cube actor.
    - Uses OBB precise collision (2D projection + SAT), places only if min_gap is satisfied.
    - avoid: Input a list of objects. Can be [actor, ...] or [(actor, pad), ...] (pad in meters).
    - generator: Must pass torch.Generator for randomization.
    """
    # Cache
    if not hasattr(self, "_spawned_cubes"):
        self._spawned_cubes = []
        self._spawned_count = 0

    center = np.array(region_center if region_center is not None else self.cube_spawn_center, dtype=np.float64)

    # Support two types of input: scalar or 2D array
    if region_half_size is None:
        region_half_size = self.cube_spawn_half_size

    # Compatible with two input formats
    if isinstance(region_half_size, (list, tuple, np.ndarray)):
        # 2D array input: independent control for xy
        area_half = np.array(region_half_size, dtype=np.float64)
        if area_half.shape == ():  # Handle 0-dim array
            area_half = np.array([float(area_half), float(area_half)], dtype=np.float64)
        elif len(area_half) == 1:
            area_half = np.array([float(area_half[0]), float(area_half[0])], dtype=np.float64)
        elif len(area_half) != 2:
            raise ValueError("region_half_size array must contain 1 or 2 elements [x_half, y_half]")
    else:
        # Scalar input: xy remain consistent
        area_half = np.array([float(region_half_size), float(region_half_size)], dtype=np.float64)

    hs_new = float(half_size if half_size is not None else self.cube_half_size)

    # Let cube fall completely inside region (independent control for xy)
    x_low = center[0] - area_half[0] + hs_new
    x_high = center[0] + area_half[0] - hs_new
    y_low = center[1] - area_half[1] + hs_new
    y_high = center[1] + area_half[1] - hs_new
    if x_low > x_high or y_low > y_high:
        raise ValueError("spawn_random_cube: Sampling region too small, cannot fit cube of this size.")

    # === Assemble Obstacle OBB (2D) List ===
    obb2d_list = []  # [(c, A, h), ...]

    def _push_actor_as_obb2d(actor, pad=0.0):
        try:
            # Special handling for board_with_hole
            if hasattr(actor, '_board_side') and hasattr(actor, '_hole_side'):
                # This is our board with hole, manually add its OBB
                board_side = actor._board_side
                hole_side = actor._hole_side

                # Get board world position
                actor_pos = actor.pose.p
                if isinstance(actor_pos, torch.Tensor):
                    actor_pos = actor_pos[0].detach().cpu().numpy()

                board_center = np.array(actor_pos[:2], dtype=np.float64)
                board_half = board_side / 2
                hole_half = hole_side / 2

                # Add OBBs for four rectangular strips
                # Top strip
                if board_half > hole_half:  # Ensure enough space
                    top_height = board_half - hole_half
                    top_center = board_center + np.array([0, hole_half + top_height / 2])
                    A_top = np.eye(2)  # No rotation
                    h_top = np.array([board_half + pad, top_height / 2 + pad])
                    obb2d_list.append((top_center, A_top, h_top))

                    # Bottom strip
                    bottom_center = board_center + np.array([0, -(hole_half + top_height / 2)])
                    obb2d_list.append((bottom_center, A_top, h_top))

                    # Left strip
                    left_width = board_half - hole_half
                    left_center = board_center + np.array([-(hole_half + left_width / 2), 0])
                    h_left = np.array([left_width / 2 + pad, hole_half + pad])
                    obb2d_list.append((left_center, A_top, h_left))

                    # Right strip
                    right_center = board_center + np.array([hole_half + left_width / 2, 0])
                    obb2d_list.append((right_center, A_top, h_left))
                return

            obb = get_actor_obb(actor, to_world_frame=True, vis=False)
            obb2d = _trimesh_box_to_obb2d(obb, extra_pad=float(pad))
            obb2d_list.append(obb2d)
        except Exception:
            # Some objects (like site/marker) do not have physical mesh, ignore or use circle approximation below
            pass

    if include_existing:
        # Main cube
        if hasattr(self, "cube") and self.cube is not None:
            _push_actor_as_obb2d(self.cube, pad=0.0)

        # Historically spawned cubes
        for ac in self._spawned_cubes:
            _push_actor_as_obb2d(ac, pad=0.0)

    # User specified extra avoidance
    if avoid:
        for it in avoid:
            if isinstance(it, tuple):
                # Check if it's a pre-made OBB tuple (c, A, h) or (actor, pad)
                if len(it) == 3 and isinstance(it[0], np.ndarray) and isinstance(it[1], np.ndarray):
                    # Pre-made OBB: (center, axes, half_sizes)
                    obb2d_list.append(it)
                else:
                    # Actor with padding
                    act_i, pad_i = it
                    _push_actor_as_obb2d(act_i, pad=float(pad_i))
            else:
                _push_actor_as_obb2d(it, pad=0.0)

    # Target point (if no mesh), supplement with "circle + circumscribed circle" conservative approximation (optional)
    circle_list = []  # [(xy(2,), R)], for objects without mesh
    def _actor_xy(actor):
        p = actor.pose.p
        if isinstance(p, torch.Tensor):
            p = p[0].detach().cpu().numpy()
        return np.array(p[:2], dtype=np.float64)

    if include_goal and hasattr(self, "goal_site") and self.goal_site is not None:
        try:
            # If goal_site has mesh, it will be covered in _push_actor_as_obb2d, here only as a fallback
            _push_actor_as_obb2d(self.goal_site, pad=0.0)
        except Exception:
            # Degrade to circle approximation: goal radius + new cube circumscribed circle radius
            R_goal = float(getattr(self, "goal_thresh", 0.03))
            R_new_ext = np.sqrt(2.0) * hs_new
            circle_list.append((_actor_xy(self.goal_site), R_goal + R_new_ext + min_gap))

    # === Sampling Iteration ===
    if generator is None:
        raise ValueError("spawn_random_cube: generator argument must be explicitly passed for randomization")

    device = self.device

    for trial in range(int(max_trials)):
        # Use simple uniform sampling to ensure good spatial coverage
        # Complex sampling strategies often reduce coverage

        u1 = torch.rand(1, generator=generator).item()
        u2 = torch.rand(1, generator=generator).item()

        # Map directly to sampling region - Uniform distribution provides best spatial coverage
        x = float(x_low + u1 * (x_high - x_low))
        y = float(y_low + u2 * (y_high - y_low))

        if random_yaw:
            # Use more random yaw generation, sampling from full [0, 2π] range
            yaw_sample = torch.rand(1, generator=generator).item()
            yaw = float(yaw_sample * 2 * np.pi)
        else:
            yaw = 0.0

        # New cube's 2D OBB (reflect min_gap in "new object half length expansion", avoid adding to both sides causing double)
        c_new, A_new, h_new = _build_new_cube_obb2d(x, y, hs_new, yaw, pad_xy=float(min_gap))

        # Check collision one by one with OBB obstacles
        hit = False
        for (c_obs, A_obs, h_obs) in obb2d_list:
            if _obb2d_intersect(c_obs, A_obs, h_obs, c_new, A_new, h_new):
                hit = True
                break
        if hit:
            continue

        # Check circular conservative obstacles (if present)
        for (xy_c, R_c) in circle_list:
            if np.linalg.norm(np.asarray([x, y], dtype=np.float64) - xy_c) < R_c:
                hit = True
                break
        if hit:
            continue

        # Passing detection, create cube (pose and collision detection use same yaw to ensure consistency)
        q = _yaw_to_quat_tensor(yaw, device=device)

        cube = actors.build_cube(
            self.scene,
            half_size=hs_new,
            color=color,
            name=name_prefix,  # Use name_prefix directly, do not add counter
            initial_pose=Pose.create_from_pq(
                torch.tensor([[x, y, hs_new]], device=device, dtype=torch.float32),
                q,
            ),
        )
        cube._cube_half_size = hs_new
        self._spawned_cubes.append(cube)
        self._spawned_count += 1
        return cube

    raise RuntimeError("spawn_random_cube: Region crowded or constraints too tight, no feasible position found. Try: increase region/decrease cube/decrease min_gap.")

def _build_new_target_obb2d(x, y, half_size_xy, yaw, pad_xy=0.0):
    """
    Construct 2D OBB for "target ready to be placed": center/axes/half-extents
    half_size_xy: float, half length of target on XY
    yaw: rotation around z-axis (radians)
    pad_xy: extra padding on half length on XY (for minimum gap)
    """
    c = np.array([x, y], dtype=np.float64)
    cos_y = np.cos(yaw)
    sin_y = np.sin(yaw)
    A = np.array([[cos_y, -sin_y],
                  [sin_y,  cos_y]], dtype=np.float64)  # Each column is an axis
    h = np.array([half_size_xy + pad_xy, half_size_xy + pad_xy], dtype=np.float64)
    return c, A, h

def spawn_random_target(
        self,
        region_center=[0, 0],
        region_half_size=0.1,
        radius=0.01,
        thickness=0.005,
        name_prefix="target_extra",
        min_gap=0.005,
        max_trials=256,
        avoid=None,          # Supports [actor, ...] or [(actor, pad), ...]
        include_existing=True,   # Whether to automatically avoid existing main target and generated extra targets
        include_goal=True,       # Whether to treat goal_site as obstacle (approximate with circle, conservative)
        generator=None,
        randomize=True,      # Control whether to randomize position
        target_style="purple",  # Choose which color scheme target to create
    ):
    """
    Drop a target (onto table) in rectangular region using rejection sampling, and return the target actor.
    - Uses OBB precise collision (2D projection + SAT), places only if min_gap is satisfied.
    - avoid: Input a list of objects. Can be [actor, ...] or [(actor, pad), ...] (pad in meters).
    - generator: Must pass torch.Generator for randomization (when randomize=True).
    - randomize: Control whether to randomize position. If False, generate directly at region_center.
    """
    # Cache
    random_yaw=False
    if not hasattr(self, "_spawned_targets"):
        self._spawned_targets = []
        self._spawned_target_count = 0

    center = np.array(region_center if region_center is not None else getattr(self, 'target_spawn_center', [0, 0]), dtype=np.float64)
    area_half = float(region_half_size if region_half_size is not None else getattr(self, 'target_spawn_half_size', 0.1))
    target_radius = float(radius if radius is not None else getattr(self, 'target_radius', 0.01))
    target_thickness = float(thickness if thickness is not None else getattr(self, 'target_thickness', 0.005))

    # Let target fall completely inside region
    x_low = center[0] - area_half + target_radius
    x_high = center[0] + area_half - target_radius
    y_low = center[1] - area_half + target_radius
    y_high = center[1] + area_half - target_radius
    if x_low > x_high or y_low > y_high:
        raise ValueError("spawn_random_target: Sampling region too small, cannot fit target of this size.")

    # === Assemble Obstacle OBB (2D) List ===
    obb2d_list = []  # [(c, A, h), ...]

    def _push_actor_as_obb2d(actor, pad=0.0):
        try:
            # Special handling for board_with_hole
            if hasattr(actor, '_board_side') and hasattr(actor, '_hole_side'):
                # This is our board with hole, manually add its OBB
                board_side = actor._board_side
                hole_side = actor._hole_side

                # Get board world position
                actor_pos = actor.pose.p
                if isinstance(actor_pos, torch.Tensor):
                    actor_pos = actor_pos[0].detach().cpu().numpy()

                board_center = np.array(actor_pos[:2], dtype=np.float64)
                board_half = board_side / 2
                hole_half = hole_side / 2

                # Add OBBs for four rectangular strips
                # Top strip
                if board_half > hole_half:  # Ensure enough space
                    top_height = board_half - hole_half
                    top_center = board_center + np.array([0, hole_half + top_height / 2])
                    A_top = np.eye(2)  # No rotation
                    h_top = np.array([board_half + pad, top_height / 2 + pad])
                    obb2d_list.append((top_center, A_top, h_top))

                    # Bottom strip
                    bottom_center = board_center + np.array([0, -(hole_half + top_height / 2)])
                    obb2d_list.append((bottom_center, A_top, h_top))

                    # Left strip
                    left_width = board_half - hole_half
                    left_center = board_center + np.array([-(hole_half + left_width / 2), 0])
                    h_left = np.array([left_width / 2 + pad, hole_half + pad])
                    obb2d_list.append((left_center, A_top, h_left))

                    # Right strip
                    right_center = board_center + np.array([hole_half + left_width / 2, 0])
                    obb2d_list.append((right_center, A_top, h_left))
                return

            obb = get_actor_obb(actor, to_world_frame=True, vis=False)
            obb2d = _trimesh_box_to_obb2d(obb, extra_pad=float(pad))
            obb2d_list.append(obb2d)
        except Exception:
            # Some objects (like site/marker) do not have physical mesh, ignore or use circle approximation below
            pass

    if include_existing:
        # Main cube
        if hasattr(self, "cube") and self.cube is not None:
            _push_actor_as_obb2d(self.cube, pad=0.0)

        # Main target
        if hasattr(self, "target") and self.target is not None:
            _push_actor_as_obb2d(self.target, pad=0.0)

        # Historically spawned cubes
        if hasattr(self, "_spawned_cubes"):
            for ac in self._spawned_cubes:
                _push_actor_as_obb2d(ac, pad=0.0)

    # Target point (if no mesh), supplement with "circle + circumscribed circle" conservative approximation (optional)
    circle_list = []  # [(xy(2,), R)], for objects without mesh
    def _actor_xy(actor):
        p = actor.pose.p
        if isinstance(p, torch.Tensor):
            p = p[0].detach().cpu().numpy()
        return np.array(p[:2], dtype=np.float64)

    # Historically spawned targets - Treat as circular obstacles
    if include_existing:
        for ac in self._spawned_targets:
            target_r = getattr(ac, "_target_radius", target_radius)
            circle_list.append((_actor_xy(ac), target_r))

    # User specified extra avoidance
    if avoid:
        for it in avoid:
            if isinstance(it, tuple):
                # Check if it's a pre-made OBB tuple (c, A, h) or (actor, pad)
                if len(it) == 3 and isinstance(it[0], np.ndarray) and isinstance(it[1], np.ndarray):
                    # Pre-made OBB: (center, axes, half_sizes)
                    obb2d_list.append(it)
                else:
                    # Actor with padding
                    act_i, pad_i = it
                    # Check if it is a target (circular)
                    if hasattr(act_i, "_target_radius"):
                        target_r = getattr(act_i, "_target_radius", target_radius)
                        circle_list.append((_actor_xy(act_i), target_r + float(pad_i)))
                    else:
                        _push_actor_as_obb2d(act_i, pad=float(pad_i))
            else:
                # Check if it is a target (circular)
                if hasattr(it, "_target_radius"):
                    target_r = getattr(it, "_target_radius", target_radius)
                    circle_list.append((_actor_xy(it), target_r))
                else:
                    _push_actor_as_obb2d(it, pad=0.0)

    if include_goal and hasattr(self, "goal_site") and self.goal_site is not None:
        try:
            # If goal_site has mesh, it will be covered in _push_actor_as_obb2d, here only as a fallback
            _push_actor_as_obb2d(self.goal_site, pad=0.0)
        except Exception:
            # Degrade to circle approximation: goal radius + new target circumscribed circle radius
            R_goal = float(getattr(self, "goal_thresh", 0.03))
            R_new_ext = target_radius
            circle_list.append((_actor_xy(self.goal_site), R_goal + R_new_ext + min_gap))

    # === Sampling Iteration ===
    if generator is None:
        raise ValueError("spawn_random_target: generator argument must be explicitly passed for randomization")

    device = self.device

    target_builders = {
        "purple": build_purple_white_target,
        "gray": build_gray_white_target,
        "green": build_green_white_target,
        "red": build_red_white_target,
    }
    if isinstance(target_style, str):
        builder_key = target_style.lower()
        if builder_key not in target_builders:
            raise ValueError(f"spawn_random_target: Unknown target_style '{target_style}'. Supported: {list(target_builders.keys())}")
        target_builder = target_builders[builder_key]
    elif callable(target_style):
        target_builder = target_style
    else:
        raise ValueError("spawn_random_target: target_style must be a string or callable builder function")

    for _ in range(int(max_trials)):
        x = float(torch.rand(1, generator=generator).item() * (x_high - x_low) + x_low)
        y = float(torch.rand(1, generator=generator).item() * (y_high - y_low) + y_low)

        if random_yaw:
            yaw = float(torch.rand(1, generator=generator).item() * 2 * np.pi - np.pi)
        else:
            yaw = 0.0

        # New target's circular collision detection (target is circular, circular detection is more accurate)
        target_pos = np.array([x, y], dtype=np.float64)
        target_collision_radius = target_radius + min_gap

        # Check collision with OBB obstacles (check circular target against square obstacles)
        hit = False
        for (c_obs, A_obs, h_obs) in obb2d_list:
            # Calculate minimum distance from circle center to OBB
            # Convert circle center to OBB local coordinate system
            local_pos = A_obs.T @ (target_pos - c_obs)
            # Calculate closest point from circle center to OBB
            closest_point = np.clip(local_pos, -h_obs, h_obs)
            # Convert back to world coordinate system
            closest_world = c_obs + A_obs @ closest_point
            # Calculate distance
            dist = np.linalg.norm(target_pos - closest_world)
            if dist < target_collision_radius:
                hit = True
                break
        if hit:
            continue

        # Check collision with circular obstacles (circle vs circle)
        for (xy_c, R_c) in circle_list:
            if np.linalg.norm(target_pos - xy_c) < (target_collision_radius + R_c):
                hit = True
                break
        if hit:
            continue

        # Passed detection, create target (pose and collision detection use same yaw to ensure consistency)
        rotate = np.array([np.cos(yaw/2), 0, 0, np.sin(yaw/2)])  # Quaternion for z-axis rotation
        angles = torch.deg2rad(torch.tensor([0.0, 90.0, 0.0], dtype=torch.float32))  # (3,)
        rotate = matrix_to_quaternion(
            euler_angles_to_matrix(angles, convention="XYZ")
        )
        target = target_builder(
            scene=self.scene,
            radius=target_radius,
            thickness=target_thickness,
            name=name_prefix,  # Use name_prefix directly, do not add counter
            body_type="kinematic",  # Visualization only
            add_collision=False,  # Disable collision
            initial_pose=sapien.Pose(p=[x, y, target_thickness], q=rotate),
        )
        target._target_radius = target_radius
        self._spawned_targets.append(target)
        self._spawned_target_count += 1
        return target

    raise RuntimeError("spawn_random_target: Region crowded or constraints too tight, no feasible position found. Try: increase region/decrease target/decrease min_gap.")


def create_button_obb(center_xy=(-0.3, 0), half_size=0.05):
    """
    Create a manual OBB for button collision avoidance.

    Args:
        center_xy: Button center position (x, y)
        half_size: Safe zone half-size around button (default 0.05m)

    Returns:
        Tuple (center, axes, half_sizes) for use in avoid lists
    """
    return (
        np.array(center_xy, dtype=np.float64),  # center
        np.eye(2, dtype=np.float64),  # axes (identity for axis-aligned)
        np.array([half_size, half_size], dtype=np.float64)  # half-sizes
    )

def build_button(
            self,
            center_xy=(0.15, 0.10),  # Button (x,y) on table
            base_half=[0.025, 0.025, 0.005],  # Base half-size [x,y,z]
            cap_radius=0.015,  # Button cap radius
            cap_half_len=0.006,  # Button cap half-length
            travel=None,  # Press travel
            stiffness=800.0,
            damping=40.0,
            scale: float = None,  # ⭐ New: scaling factor
            generator=None,
            name: str = "button",  # ⭐ New: button name
            randomize: bool = True,  # ⭐ New: whether to randomize position
            randomize_range=(0.1, 0.4),  # ⭐ New: randomization range, (range_x, range_y)
    ):
        # ------- Scaling and Travel -------
        if scale is None:
            # If not passed, use default scaling from environment
            scale = getattr(self, "button_scale", 1.0)
        scale = float(scale)

        # Travel priority: argument > environment base
        if travel is None:
            # Scale proportionally using base travel
            base_travel = getattr(self, "_button_travel_base", 0.1)
            travel = base_travel * scale
        else:
            # If travel explicitly passed, also follow scale (to keep absolute value, change next line to pass)
            travel = float(travel) * scale

        # Size scaling
        base_half = [bh * scale for bh in base_half]
        cap_radius = float(cap_radius) * scale
        cap_half_len = float(cap_half_len) * scale

        # Record current button travel for other functions
        self.button_travel = float(travel)

        # ------- Position Randomization -------
        cx, cy = float(center_xy[0]), float(center_xy[1])

        if randomize:
            if not isinstance(randomize_range, (tuple, list, np.ndarray)):
                raise TypeError("randomize_range must be a sequence of length 2.")
            if len(randomize_range) != 2:
                raise ValueError("randomize_range must contain exactly two elements.")
            range_x, range_y = float(randomize_range[0]), float(randomize_range[1])
            offset = torch.rand(2, generator=generator) - 0.5
            cx += float(offset[0]) * range_x
            cy += float(offset[1]) * range_y
        center_xy = (cx, cy)

        scene = self.scene
        builder = scene.create_articulation_builder()

        # Initial pose: lift base center to z=base_half[2]
        builder.initial_pose = sapien.Pose(p=[cx, cy, base_half[2]])

        # Root: Base
        base = builder.create_link_builder()
        base.set_name("button_base")
        base.add_box_collision(half_size=base_half, density=200000)
        base.add_box_visual(half_size=base_half)

        # Child: Button cap (vertical sliding)
        cap = builder.create_link_builder(base)
        cap.set_name("button_cap")
        cap.set_joint_name("button_joint")

        R_up = euler2quat(0, -np.pi / 2, 0)  # Align joint x-axis with world z

        cap.set_joint_properties(
            type="prismatic",
            limits=[[-travel, 0.0]],  # Negative direction is pressed
            pose_in_parent=sapien.Pose(p=[0, 0, base_half[2]], q=R_up),
            pose_in_child=sapien.Pose(p=[0, 0, 0.0], q=R_up),
            friction=0.0,
            damping=0.0,
        )

        cap.add_cylinder_collision(
            half_length=cap_half_len, radius=cap_radius,
            pose=sapien.Pose(p=[0, 0, cap_half_len], q=R_up), density=1500
        )
        material = sapien.render.RenderMaterial()
        material.set_base_color([0.5, 0.5, 0.5, 1.0])
        cap.add_cylinder_visual(
            half_length=cap_half_len, radius=cap_radius,
            pose=sapien.Pose(p=[0, 0, cap_half_len], q=R_up), material=material
        )



        button = builder.build(name=name, fix_root_link=True)

        j = {j.name: j for j in button.get_joints()}["button_joint"]
        j.set_drive_properties(stiffness=stiffness, damping=damping)
        j.set_drive_target(0.0)

        self.button = button
        self.button_joint = j

        cap_link = next(
            link for link in button.get_links()
            if link.get_name() == "button_cap"
        )
        cap_link = next(link for link in button.get_links()
                if link.get_name() == "button_cap")
        if not hasattr(self, "cap_links"):
            self.cap_links = {}
        self.cap_links[name] = [cap_link]   # name is "button_left", "button_right", etc.
        self.cap_link = self.cap_links[name]  # Compatible with old logic

        # Provide an OBB for downstream placement logic using the scaled button footprint
        button_obb = create_button_obb(
            center_xy=center_xy,
            half_size=max(base_half[0], base_half[1]) * 1.5,
        )
        return button_obb
def build_bin(
        self,
        *,
        inner_side: float = 0.04,  # Inner opening side length (full length, meters), originally 2*inner_side_half_len = 0.04
        wall_thickness: float = 0.005,  # Wall thickness (full thickness, meters)
        wall_height: float = 0.05,  # Wall height (full height, meters)
        floor_thickness: float = 0.004,  # Floor thickness (full thickness, meters)
        callsign=None,
        position=None,  # Add position argument
        z_rotation_deg=0.0  # Add z-axis rotation angle argument (degrees)
):
    """
    Assemble an "open box" using 1 floor + 4 wall strips.
    All dimensions use "full size (meters)", automatically converted to half-size internally.
    Refer to cube generation method, let bin bottom sit on table (z=0).
    """
    inner_side = self.cube_half_size * 2.5
    wall_height = self.cube_half_size * 2.5

    # ---- Convert full size to half size (consistent with add_box_* interface) ----
    inner_half = inner_side * 0.5
    t = wall_thickness * 0.5  # Half wall thickness
    h = wall_height * 0.5  # Half wall height
    tf = floor_thickness * 0.5  # Half floor thickness

    # ---- Component half sizes (in world coordinates [x, y, z]) ----
    # Floor: covers inner opening + two side wall thicknesses
    bottom_half = [inner_half + t, inner_half + t, tf]
    # Left/Right Wall: thickness along x, height along z, length along y
    lr_wall_half = [t, inner_half + t, h]
    # Front/Back Wall: thickness along y, height along z, length along x
    fb_wall_half = [inner_half + t, t, h]

    # ---- Determine bin position (refer to cube way) ----
    if position is None:
        base_pos = [0.0, 0.0, 0.0]
    else:
        base_pos = list(position)

    # Build geometry as "opening up" in local coordinate system, then flip to "opening down" globally
    # Floor on table, walls extend up from floor top (flipped becomes down)
    base_z = tf  # Floor center height (half of floor thickness)

    # ---- Component placement positions (relative to bin builder origin) ----
    # Wall center horizontal offset = inner half + half wall thickness
    offset = inner_half + t
    # Wall center vertical position = floor thickness + half wall height
    z_wall = tf + h

    poses = [
        sapien.Pose([0.0, 0.0, 0]),
        # Floor: on table, half thickness height
        sapien.Pose([0.0, 0.0, base_z]),
        # Left/Right Wall (+/- x direction): extend up from floor top
        sapien.Pose([-offset, 0.0, z_wall]),
        sapien.Pose([+offset, 0.0, z_wall]),
        # Front/Back Wall (+/- y direction): extend up from floor top
        sapien.Pose([0.0, -offset, z_wall]),
        sapien.Pose([0.0, +offset, z_wall]),
    ]
    half_sizes = [
        [self.cube_half_size,self.cube_half_size,self.cube_half_size],
        bottom_half,
        lr_wall_half,  # Left
        lr_wall_half,  # Right
        fb_wall_half,  # Front
        fb_wall_half,  # Back
    ]

    builder = self.scene.create_actor_builder()

    # Let bin "clasp" on table: flip 180 degrees around x-axis, opening down, then rotate around z-axis
    angles = torch.deg2rad(torch.tensor([180.0, 0.0, z_rotation_deg], dtype=torch.float32))  # (3,)
    rotate = matrix_to_quaternion(
        euler_angles_to_matrix(angles, convention="XYZ")
    )
    # Lowest point after rotation is -(tf + 2h), translate to z=0 to sit on table
    builder.set_initial_pose(
        sapien.Pose(
            p=[base_pos[0], base_pos[1], tf + 2 * h],
            q=rotate,
        )
    )

    for pose, half_size in zip(poses, half_sizes):
        builder.add_box_collision(pose, half_size)
        builder.add_box_visual(pose, half_size)

    bin_actor = builder.build_dynamic(name=callsign)

    return bin_actor

def spawn_random_bin(
        self,
        avoid=None,
        region_center=[-0.1, 0],
        region_half_size=0.3,
        min_gap=0.05,
        name_prefix="bin",
        max_trials=256,
        generator=None
):
    """
    Drop a bin in rectangular region using rejection sampling, and return the bin actor.
    Use OBB precise collision detection, place only if min_gap is satisfied.
    """
    if avoid is None:
        avoid = []

    center = np.array(region_center, dtype=np.float64)
    area_half = float(region_half_size)

    # Calculate bin size (for collision detection)
    inner_side = self.cube_half_size * 2.5
    wall_thickness = 0.005
    bin_half_size = (inner_side + wall_thickness) * 0.5  # Half of bin total size

    # Let bin fall completely inside region
    x_low = center[0] - area_half + bin_half_size
    x_high = center[0] + area_half - bin_half_size
    y_low = center[1] - area_half + bin_half_size
    y_high = center[1] + area_half - bin_half_size

    if x_low > x_high or y_low > y_high:
        raise ValueError("_spawn_random_bin: Sampling region too small, cannot fit bin of this size.")

    # === Assemble Obstacle OBB (2D) List ===
    obb2d_list = []  # [(c, A, h), ...]

    def _push_actor_as_obb2d(actor, pad=0.0):
        try:
            obb = get_actor_obb(actor, to_world_frame=True, vis=False)
            obb2d = _trimesh_box_to_obb2d(obb, extra_pad=float(pad))
            obb2d_list.append(obb2d)
        except Exception:
            # Some objects (like site/marker) do not have physical mesh, ignore
            pass

    # Collect avoidance object OBBs
    for item in avoid:
        if isinstance(item, tuple):
            # Check if it's a pre-made OBB tuple (c, A, h) or (actor, pad)
            if len(item) == 3 and isinstance(item[0], np.ndarray) and isinstance(item[1], np.ndarray):
                # Pre-made OBB: (center, axes, half_sizes)
                obb2d_list.append(item)
            else:
                # Actor with padding
                actor, pad = item
                _push_actor_as_obb2d(actor, pad)
        else:
            _push_actor_as_obb2d(item, min_gap)

    for trial in range(int(max_trials)):
        x = float(torch.rand(1, generator=generator).item() * (x_high - x_low) + x_low)
        y = float(torch.rand(1, generator=generator).item() * (y_high - y_low) + y_low)

        # New bin square collision detection
        bin_pos = np.array([x, y], dtype=np.float64)
        bin_collision_half_size = bin_half_size + min_gap

        # Detect collision with other OBB obstacles
        hit = False
        for (c_obs, A_obs, h_obs) in obb2d_list:
            # Simplify: treat bin as square, detect collision with OBB
            # Calculate bin center to OBB closest distance
            local_pos = A_obs.T @ (bin_pos - c_obs)
            closest_point = np.clip(local_pos, -h_obs, h_obs)
            closest_world = c_obs + A_obs @ closest_point
            dist = np.linalg.norm(bin_pos - closest_world)
            if dist < bin_collision_half_size:
                hit = True
                break

        if hit:
            continue

        # Passing detection, create bin (at specified position), with random z-axis rotation
        z_rotation = float(torch.rand(1, generator=generator).item() * 90.0)  # 0-360 degrees
        bin_actor = build_bin(self, callsign=name_prefix, position=[x, y, 0.002], z_rotation_deg=z_rotation)

        return bin_actor

    raise RuntimeError("_spawn_random_bin: Region crowded or constraints too tight, no feasible position found. Try: increase region/decrease bin/decrease min_gap.")

def spawn_fixed_cube(
        self,
        position,  # [x, y, z] fixed position
        half_size=None,
        color=(1, 0, 0, 1),
        name_prefix="fixed_cube",
        yaw=0.0,  # rotation around z-axis (radians)
        dynamic=False,
    ):
    """
    Generate a cube at fixed position, no collision detection.
    Use builder pattern to create dynamic object, refer to build_bin implementation.
    """
    hs = float(half_size if half_size is not None else self.cube_half_size)

    # Ensure position is array format
    pos = np.array(position, dtype=np.float64)
    if len(pos) == 2:
        # If only x,y provided, set z to cube half height (let cube bottom sit on table)
        pos = np.append(pos, hs)

    # Create actor builder
    builder = self.scene.create_actor_builder()

    # Generate rotation quaternion (rotate yaw angle around z-axis)
    if yaw != 0.0:
        angles = torch.tensor([0.0, 0.0, float(yaw)], dtype=torch.float32)
        R = euler_angles_to_matrix(angles.unsqueeze(0), convention="XYZ")[0]
        q = matrix_to_quaternion(R.unsqueeze(0))[0]
        rotate = q
    else:
        rotate = torch.tensor([1.0, 0.0, 0.0, 0.0])  # Identity quaternion

    # Set initial position and rotation
    builder.set_initial_pose(
        sapien.Pose(
            p=[pos[0], pos[1], pos[2]],
            q=rotate.numpy() if isinstance(rotate, torch.Tensor) else rotate
        )
    )

    # Add box geometry (collision and visual)
    half_size_list = [hs, hs, hs]
    if  dynamic==True:
        # Collision geometry stays at builder origin; initial pose already positions the actor
        builder.add_box_collision(sapien.Pose([0, 0, 0]), half_size_list)

    # Create material
    material = sapien.render.RenderMaterial()
    material.set_base_color(color)
    builder.add_box_visual(sapien.Pose([0, 0, 0]), half_size_list, material=material)

    # Choose build method based on dynamic argument
    if dynamic==True:
        cube = builder.build_dynamic(name=name_prefix)
    else:
        cube = builder.build_kinematic(name=name_prefix)

    # Set cube attribute
    cube._cube_half_size = hs

    return cube

def build_board_with_hole(
        self,
        *,
        board_side=0.01,  # Square board side length
        hole_side=0.06,  # Square hole side length
        thickness=0.02,  # Board thickness
        position=None,  # Board position [x, y] or [x, y, z]
        rotation_quat=None,  # Rotation quaternion [w, x, y, z]
        name="board_with_hole"
):
    """
    Create a square board with a square hole
    Combine four rectangular strips: top, bottom, left, right

    Args:
        height: If provided, overwrite z coordinate in position
    """
    if position is None:
        position = [0.3, 0, 0]  # Default position, bottom on table


    # Board and hole half lengths
    board_half = board_side / 2
    hole_half = hole_side / 2
    thickness_half = thickness / 2

    # Use input position as board bottom, calculate board center position
    # Input position is bottom position, need to add thickness_half to get center position
    center_position = [position[0], position[1], position[2] + thickness_half]

    # Create actor builder
    builder = self.scene.create_actor_builder()

    # Set board initial position (use center position)
    if rotation_quat is None:
        rotation_quat = [1.0, 0.0, 0.0, 0.0]  # No rotation
    builder.set_initial_pose(
        sapien.Pose(
            p=center_position,
            q=rotation_quat
        )
    )

    # Create material - brown board
    material = sapien.render.RenderMaterial()
    material.set_base_color([0.8, 0.6, 0.4, 1.0])  # Light brown

    # Four rectangular strips dimensions and positions
    # Top strip
    top_width = board_side  # Full board width
    top_height = board_half - hole_half  # From hole top to board top
    top_center_y = hole_half + top_height / 2
    builder.add_box_collision(
        sapien.Pose([0, top_center_y, 0]),
        [top_width / 2, top_height / 2, thickness_half]
    )
    builder.add_box_visual(
        sapien.Pose([0, top_center_y, 0]),
        [top_width / 2, top_height / 2, thickness_half],
        material=material
    )

    # Bottom strip
    bottom_width = board_side  # Full board width
    bottom_height = board_half - hole_half  # From board bottom to hole bottom
    bottom_center_y = -(hole_half + bottom_height / 2)
    builder.add_box_collision(
        sapien.Pose([0, bottom_center_y, 0]),
        [bottom_width / 2, bottom_height / 2, thickness_half]
    )
    builder.add_box_visual(
        sapien.Pose([0, bottom_center_y, 0]),
        [bottom_width / 2, bottom_height / 2, thickness_half],
        material=material
    )

    # Left strip - only within hole height range
    left_width = board_half - hole_half  # From board left to hole left
    left_height = hole_side  # Hole height
    left_center_x = -(hole_half + left_width / 2)
    builder.add_box_collision(
        sapien.Pose([left_center_x, 0, 0]),
        [left_width / 2, left_height / 2, thickness_half]
    )
    builder.add_box_visual(
        sapien.Pose([left_center_x, 0, 0]),
        [left_width / 2, left_height / 2, thickness_half],
        material=material
    )

    # Right strip - only within hole height range
    right_width = board_half - hole_half  # From hole right to board right
    right_height = hole_side  # Hole height
    right_center_x = hole_half + right_width / 2
    builder.add_box_collision(
        sapien.Pose([right_center_x, 0, 0]),
        [right_width / 2, right_height / 2, thickness_half]
    )
    builder.add_box_visual(
        sapien.Pose([right_center_x, 0, 0]),
        [right_width / 2, right_height / 2, thickness_half],
        material=material
    )

    # Add a black cube at hole center with same size as hole but half height (visual only, no collision)
    hole_cube_half_size_xy = hole_half  # cube half size same as hole
    hole_cube_half_height = thickness_half / 2  # cube height is half of board thickness

    # Create black material
    black_material = sapien.render.RenderMaterial()
    black_material.set_base_color([0.0, 0.0, 0.0, 1.0])  # Black

    # Add black cube (visual only, no collision)
    # Position: cube bottom at board bottom, so cube center at -thickness_half + hole_cube_half_height
    cube_center_z = -thickness_half + hole_cube_half_height
    builder.add_box_visual(
        sapien.Pose([0, 0, cube_center_z]),  # Black cube bottom at board bottom
        [hole_cube_half_size_xy, hole_cube_half_size_xy, hole_cube_half_height],
        material=black_material
    )

    # Build actor
    board_actor = builder.build_kinematic(name=name)

    # Store board attributes
    board_actor._board_side = board_side
    board_actor._hole_side = hole_side
    board_actor._thickness = thickness

    return board_actor


def build_purple_white_target(
    scene: ManiSkillScene,
    radius: float,
    thickness: float,
    name: str,
    body_type: str = "dynamic",
    add_collision: bool = True,
    scene_idxs: Optional[Array] = None,
    initial_pose: Optional[Union[Pose, sapien.Pose]] = None,
):
    TARGET_PURPLE = (np.array([160, 32, 240, 255]) / 255).tolist()
    builder = scene.create_actor_builder()
    builder.add_cylinder_visual(
        radius=radius,
        half_length=thickness / 2,
        material=sapien.render.RenderMaterial(base_color=TARGET_PURPLE),
    )
    builder.add_cylinder_visual(
        radius=radius * 4 / 5,
        half_length=thickness / 2 + 1e-5,
        material=sapien.render.RenderMaterial(base_color=[1, 1, 1, 1]),
    )
    builder.add_cylinder_visual(
        radius=radius * 3 / 5,
        half_length=thickness / 2 + 2e-5,
        material=sapien.render.RenderMaterial(base_color=TARGET_PURPLE),
    )
    builder.add_cylinder_visual(
        radius=radius * 2 / 5,
        half_length=thickness / 2 + 3e-5,
        material=sapien.render.RenderMaterial(base_color=[1, 1, 1, 1]),
    )
    builder.add_cylinder_visual(
        radius=radius * 1 / 5,
        half_length=thickness / 2 + 4e-5,
        material=sapien.render.RenderMaterial(base_color=TARGET_PURPLE),
    )
    if add_collision:
        builder.add_cylinder_collision(
            radius=radius,
            half_length=thickness / 2,
        )
        builder.add_cylinder_collision(
            radius=radius * 4 / 5,
            half_length=thickness / 2 + 1e-5,
        )
        builder.add_cylinder_collision(
            radius=radius * 3 / 5,
            half_length=thickness / 2 + 2e-5,
        )
        builder.add_cylinder_collision(
            radius=radius * 2 / 5,
            half_length=thickness / 2 + 3e-5,
        )
        builder.add_cylinder_collision(
            radius=radius * 1 / 5,
            half_length=thickness / 2 + 4e-5,
        )
    return _build_by_type(builder, name, body_type, scene_idxs, initial_pose)

def build_gray_white_target(
    scene: ManiSkillScene,
    radius: float,
    thickness: float,
    name: str,
    body_type: str = "dynamic",
    add_collision: bool = True,
    scene_idxs: Optional[Array] = None,
    initial_pose: Optional[Union[Pose, sapien.Pose]] = None,
):
    TARGET_GRAY = (np.array([128, 128, 128, 255]) / 255).tolist()
    builder = scene.create_actor_builder()
    builder.add_cylinder_visual(
        radius=radius,
        half_length=thickness / 2,
        material=sapien.render.RenderMaterial(base_color=TARGET_GRAY),
    )
    builder.add_cylinder_visual(
        radius=radius * 4 / 5,
        half_length=thickness / 2 + 1e-5,
        material=sapien.render.RenderMaterial(base_color=[1, 1, 1, 1]),
    )
    builder.add_cylinder_visual(
        radius=radius * 3 / 5,
        half_length=thickness / 2 + 2e-5,
        material=sapien.render.RenderMaterial(base_color=TARGET_GRAY),
    )
    builder.add_cylinder_visual(
        radius=radius * 2 / 5,
        half_length=thickness / 2 + 3e-5,
        material=sapien.render.RenderMaterial(base_color=[1, 1, 1, 1]),
    )
    builder.add_cylinder_visual(
        radius=radius * 1 / 5,
        half_length=thickness / 2 + 4e-5,
        material=sapien.render.RenderMaterial(base_color=TARGET_GRAY),
    )
    if add_collision:
        builder.add_cylinder_collision(
            radius=radius,
            half_length=thickness / 2,
        )
        builder.add_cylinder_collision(
            radius=radius * 4 / 5,
            half_length=thickness / 2 + 1e-5,
        )
        builder.add_cylinder_collision(
            radius=radius * 3 / 5,
            half_length=thickness / 2 + 2e-5,
        )
        builder.add_cylinder_collision(
            radius=radius * 2 / 5,
            half_length=thickness / 2 + 3e-5,
        )
        builder.add_cylinder_collision(
            radius=radius * 1 / 5,
            half_length=thickness / 2 + 4e-5,
        )
    return _build_by_type(builder, name, body_type, scene_idxs, initial_pose)

def build_green_white_target(
    scene: ManiSkillScene,
    radius: float,
    thickness: float,
    name: str,
    body_type: str = "dynamic",
    add_collision: bool = True,
    scene_idxs: Optional[Array] = None,
    initial_pose: Optional[Union[Pose, sapien.Pose]] = None,
):
    TARGET_GREEN = (np.array([34, 139, 34, 255]) / 255).tolist()
    builder = scene.create_actor_builder()
    builder.add_cylinder_visual(
        radius=radius,
        half_length=thickness / 2,
        material=sapien.render.RenderMaterial(base_color=TARGET_GREEN),
    )
    builder.add_cylinder_visual(
        radius=radius * 4 / 5,
        half_length=thickness / 2 + 1e-5,
        material=sapien.render.RenderMaterial(base_color=[1, 1, 1, 1]),
    )
    builder.add_cylinder_visual(
        radius=radius * 3 / 5,
        half_length=thickness / 2 + 2e-5,
        material=sapien.render.RenderMaterial(base_color=TARGET_GREEN),
    )
    builder.add_cylinder_visual(
        radius=radius * 2 / 5,
        half_length=thickness / 2 + 3e-5,
        material=sapien.render.RenderMaterial(base_color=[1, 1, 1, 1]),
    )
    builder.add_cylinder_visual(
        radius=radius * 1 / 5,
        half_length=thickness / 2 + 4e-5,
        material=sapien.render.RenderMaterial(base_color=TARGET_GREEN),
    )
    if add_collision:
        builder.add_cylinder_collision(
            radius=radius,
            half_length=thickness / 2,
        )
        builder.add_cylinder_collision(
            radius=radius * 4 / 5,
            half_length=thickness / 2 + 1e-5,
        )
        builder.add_cylinder_collision(
            radius=radius * 3 / 5,
            half_length=thickness / 2 + 2e-5,
        )
        builder.add_cylinder_collision(
            radius=radius * 2 / 5,
            half_length=thickness / 2 + 3e-5,
        )
        builder.add_cylinder_collision(
            radius=radius * 1 / 5,
            half_length=thickness / 2 + 4e-5,
        )
    return _build_by_type(builder, name, body_type, scene_idxs, initial_pose)

def build_red_white_target(
    scene: ManiSkillScene,
    radius: float,
    thickness: float,
    name: str,
    body_type: str = "dynamic",
    add_collision: bool = True,
    scene_idxs: Optional[Array] = None,
    initial_pose: Optional[Union[Pose, sapien.Pose]] = None,
):
    TARGET_RED = (np.array([200, 33, 33, 255]) / 255).tolist()
    builder = scene.create_actor_builder()
    builder.add_cylinder_visual(
        radius=radius,
        half_length=thickness / 2,
        material=sapien.render.RenderMaterial(base_color=TARGET_RED),
    )
    builder.add_cylinder_visual(
        radius=radius * 4 / 5,
        half_length=thickness / 2 + 1e-5,
        material=sapien.render.RenderMaterial(base_color=[1, 1, 1, 1]),
    )
    builder.add_cylinder_visual(
        radius=radius * 3 / 5,
        half_length=thickness / 2 + 2e-5,
        material=sapien.render.RenderMaterial(base_color=TARGET_RED),
    )
    builder.add_cylinder_visual(
        radius=radius * 2 / 5,
        half_length=thickness / 2 + 3e-5,
        material=sapien.render.RenderMaterial(base_color=[1, 1, 1, 1]),
    )
    builder.add_cylinder_visual(
        radius=radius * 1 / 5,
        half_length=thickness / 2 + 4e-5,
        material=sapien.render.RenderMaterial(base_color=TARGET_RED),
    )
    if add_collision:
        builder.add_cylinder_collision(
            radius=radius,
            half_length=thickness / 2,
        )
        builder.add_cylinder_collision(
            radius=radius * 4 / 5,
            half_length=thickness / 2 + 1e-5,
        )
        builder.add_cylinder_collision(
            radius=radius * 3 / 5,
            half_length=thickness / 2 + 2e-5,
        )
        builder.add_cylinder_collision(
            radius=radius * 2 / 5,
            half_length=thickness / 2 + 3e-5,
        )
        builder.add_cylinder_collision(
            radius=radius * 1 / 5,
            half_length=thickness / 2 + 4e-5,
        )
    return _build_by_type(builder, name, body_type, scene_idxs, initial_pose)

def _build_by_type(
    builder: ActorBuilder,
    name,
    body_type,
    scene_idxs: Optional[Array] = None,
    initial_pose: Optional[Union[Pose, sapien.Pose]] = None,
):
    if scene_idxs is not None:
        builder.set_scene_idxs(scene_idxs)
    if initial_pose is not None:
        builder.set_initial_pose(initial_pose)
    if body_type == "dynamic":
        actor = builder.build(name=name)
    elif body_type == "static":
        actor = builder.build_static(name=name)
    elif body_type == "kinematic":
        actor = builder.build_kinematic(name=name)
    else:
        raise ValueError(f"Unknown body type {body_type}")
    return actor