Environments

Base Environment

BaseEnv

Bases: ABC

get_joint_states() -> np.ndarray abstractmethod

Get current joint positions.

Source code in autolife_planning/envs/base_env.py

@abstractmethod
def get_joint_states(self) -> np.ndarray:
    """Get current joint positions."""
    raise NotImplementedError

set_joint_states(config: np.ndarray) abstractmethod

Set joint positions.

Source code in autolife_planning/envs/base_env.py

@abstractmethod
def set_joint_states(self, config: np.ndarray):
    """Set joint positions."""

get_localization() -> np.ndarray abstractmethod

Get robot localization as [x, y, theta].

Source code in autolife_planning/envs/base_env.py

@abstractmethod
def get_localization(self) -> np.ndarray:
    """Get robot localization as [x, y, theta]."""
    raise NotImplementedError

get_obs() -> Any abstractmethod

Get observations.

Source code in autolife_planning/envs/base_env.py

@abstractmethod
def get_obs(self) -> Any:
    """Get observations."""
    raise NotImplementedError

PyBullet Environment

pybullet_env

PyBulletEnv(config: RobotConfig, visualize: bool = True, viz_urdf_path: str | None = None)

Bases: BaseEnv

Source code in autolife_planning/envs/pybullet_env.py

def __init__(
    self,
    config: RobotConfig,
    visualize: bool = True,
    viz_urdf_path: str | None = None,
):
    self.config = config
    urdf = viz_urdf_path if viz_urdf_path else config.urdf_path
    urdf = _prepare_urdf_for_pybullet(urdf)
    self.sim = vpb.PyBulletSimulator(urdf, config.joint_names, visualize=visualize)
    self.joint_names = config.joint_names

    # Find camera link index
    self.camera_link_idx = -1
    # Use skel_id (body id) to query joints
    num_joints = self.sim.client.getNumJoints(self.sim.skel_id)
    for i in range(num_joints):
        info = self.sim.client.getJointInfo(self.sim.skel_id, i)
        # info[12] is the child link name (bytes)
        if info[12].decode("utf-8") == config.camera.link_name:
            self.camera_link_idx = i
            break

    # Set initial pose
    from autolife_planning.autolife import HOME_JOINTS

    self.set_joint_states(HOME_JOINTS[3:])

set_base_position(x: float, y: float, theta: float)

Move the robot base to (x, y) with yaw=theta in the world frame.

Source code in autolife_planning/envs/pybullet_env.py

def set_base_position(self, x: float, y: float, theta: float):
    """Move the robot base to (x, y) with yaw=theta in the world frame."""
    quat = self.sim.client.getQuaternionFromEuler([0, 0, theta])
    self.sim.client.resetBasePositionAndOrientation(
        self.sim.skel_id, [x, y, 0], quat
    )

set_configuration(config: np.ndarray)

Apply a full 24-DOF config (3 base + 21 joints) to the visualization.

Source code in autolife_planning/envs/pybullet_env.py

def set_configuration(self, config: np.ndarray):
    """Apply a full 24-DOF config (3 base + 21 joints) to the visualization."""
    self.set_base_position(config[0], config[1], config[2])
    self.set_joint_states(config[3:])

wait_key(key: str | int, message: str = '') -> None

Block until key is pressed in the PyBullet GUI.

Prints message to stdout, then polls keyboard events until the requested key is triggered. Accepts either a single character ("n") or a raw PyBullet key code (e.g. pb.B3G_RIGHT_ARROW). Returns early if the GUI window is closed.

Source code in autolife_planning/envs/pybullet_env.py

def wait_key(self, key: str | int, message: str = "") -> None:
    """Block until *key* is pressed in the PyBullet GUI.

    Prints *message* to stdout, then polls keyboard events until the
    requested key is triggered.  Accepts either a single character
    (``"n"``) or a raw PyBullet key code (e.g. ``pb.B3G_RIGHT_ARROW``).
    Returns early if the GUI window is closed.
    """
    client = self.sim.client
    key_code = ord(key) if isinstance(key, str) else int(key)
    if message:
        print(message)
    try:
        while client.isConnected():
            keys = client.getKeyboardEvents()
            if key_code in keys and keys[key_code] & pb.KEY_WAS_TRIGGERED:
                break
            time.sleep(0.01)
    except pb.error:
        pass

wait_for_close() -> None

Block until the user closes the PyBullet GUI window.

Source code in autolife_planning/envs/pybullet_env.py

def wait_for_close(self) -> None:
    """Block until the user closes the PyBullet GUI window."""
    client = self.sim.client
    try:
        while client.isConnected():
            client.getKeyboardEvents()  # keeps the event queue alive
            time.sleep(0.05)
    except pb.error:
        pass

animate_path(path: np.ndarray, fps: float = 60.0, next_key: str | int | None = None) -> bool

Interactively play back a full-DOF path, VAMP-style.

Each row of path is handed to :meth:set_configuration. The method blocks on the PyBullet GUI with these controls:

• SPACE — toggle auto-play (loops end → start)
• ← / → — step one waypoint back / forward (while paused)
• next_key (opt.) — exit and return True — useful when the caller wants to drive a sequence of demos

Exits when the user closes the GUI window (returns False) or presses next_key if one is provided (returns True). The animation starts paused so the caller can read the banner before anything moves.

Parameters:

Name	Type	Description	Default
path	ndarray	(N, dof) array of full-DOF configurations.	required
fps	float	playback frame rate while auto-playing.	60.0
next_key	str | int | None	optional single-character string or raw PyBullet key code; pressing it exits the viewer.	None

Returns:

Type	Description
bool	True if next_key was pressed, False otherwise.

Source code in autolife_planning/envs/pybullet_env.py

def animate_path(
    self,
    path: np.ndarray,
    fps: float = 60.0,
    next_key: str | int | None = None,
) -> bool:
    """Interactively play back a full-DOF path, VAMP-style.

    Each row of *path* is handed to :meth:`set_configuration`.  The
    method blocks on the PyBullet GUI with these controls:

    * ``SPACE``             — toggle auto-play (loops end → start)
    * ``← / →``             — step one waypoint back / forward (while paused)
    * ``next_key`` (opt.)   — exit and return ``True`` — useful when
      the caller wants to drive a sequence of demos

    Exits when the user closes the GUI window (returns ``False``) or
    presses *next_key* if one is provided (returns ``True``).  The
    animation starts paused so the caller can read the banner before
    anything moves.

    Args:
      path: ``(N, dof)`` array of full-DOF configurations.
      fps:  playback frame rate while auto-playing.
      next_key: optional single-character string or raw PyBullet key
        code; pressing it exits the viewer.

    Returns:
      ``True`` if *next_key* was pressed, ``False`` otherwise.
    """
    if path is None or len(path) == 0:
        return False

    client = self.sim.client
    dt = 1.0 / fps
    n = int(len(path))
    idx = 0
    playing = False

    left = pb.B3G_LEFT_ARROW
    right = pb.B3G_RIGHT_ARROW
    space = ord(" ")
    next_code: int | None = None
    if next_key is not None:
        next_code = ord(next_key) if isinstance(next_key, str) else int(next_key)

    parts = ["SPACE play/pause", "←/→ step"]
    if next_code is not None:
        parts.append(f"'{next_key}' next")
    parts.append("close window to exit")
    print("  " + "  |  ".join(parts))

    advanced_to_next = False
    try:
        while client.isConnected():
            self.set_configuration(path[idx])

            keys = client.getKeyboardEvents()
            if space in keys and keys[space] & pb.KEY_WAS_TRIGGERED:
                playing = not playing
            elif (
                next_code is not None
                and next_code in keys
                and keys[next_code] & pb.KEY_WAS_TRIGGERED
            ):
                advanced_to_next = True
                break
            elif not playing and left in keys and keys[left] & pb.KEY_WAS_TRIGGERED:
                idx = (idx - 1) % n
            elif (
                not playing and right in keys and keys[right] & pb.KEY_WAS_TRIGGERED
            ):
                idx = (idx + 1) % n
            elif playing:
                idx = (idx + 1) % n

            time.sleep(dt)
    except pb.error:
        pass

    return advanced_to_next

add_mesh(mesh_file: str, position: np.ndarray = np.zeros(3), orientation: np.ndarray = np.array([0, 0, 0, 1]), scale: np.ndarray = np.ones(3), mass: float = 0.0, name: str | None = None)

Add a mesh to the simulation environment directly using the raw PyBullet client, bypassing the wrapper to avoid modifying third_party code.

Source code in autolife_planning/envs/pybullet_env.py

def add_mesh(
    self,
    mesh_file: str,
    position: np.ndarray = np.zeros(3),
    orientation: np.ndarray = np.array([0, 0, 0, 1]),
    scale: np.ndarray = np.ones(3),
    mass: float = 0.0,
    name: str | None = None,
):
    """
    Add a mesh to the simulation environment directly using the raw PyBullet client,
    bypassing the wrapper to avoid modifying third_party code.
    """
    # Ensure the simulator isn't rendering while we load to speed it up
    # We can't easily use the DisableRendering context manager from vamp here
    # without importing it, but we can access the client directly.
    self.sim.client.configureDebugVisualizer(pb.COV_ENABLE_RENDERING, 0)

    vis_shape_id = self.sim.client.createVisualShape(
        shapeType=pb.GEOM_MESH,
        fileName=mesh_file,
        meshScale=scale.tolist(),
    )
    col_shape_id = self.sim.client.createCollisionShape(
        shapeType=pb.GEOM_MESH,
        fileName=mesh_file,
        meshScale=scale.tolist(),
    )

    multibody_id = self.sim.client.createMultiBody(
        baseVisualShapeIndex=vis_shape_id,
        baseCollisionShapeIndex=col_shape_id,
        basePosition=position.tolist(),
        baseOrientation=orientation.tolist(),
        baseMass=mass,
    )

    if name:
        # Add debug text
        self.sim.client.addUserDebugText(
            text=name,
            textPosition=position.tolist(),
            textColorRGB=[0.0, 0.0, 0.0],
        )

    # Try to load texture if it exists
    base_path = os.path.splitext(mesh_file)[0]
    for ext in [".png", ".jpg", ".jpeg", ".tga"]:
        tex_path = base_path + ext
        if os.path.exists(tex_path):
            tex_id = self.sim.client.loadTexture(tex_path)
            self.sim.client.changeVisualShape(
                multibody_id, -1, textureUniqueId=tex_id
            )
            break

    self.sim.client.configureDebugVisualizer(pb.COV_ENABLE_RENDERING, 1)
    return multibody_id

draw_plane(center, half_sizes: tuple[float, float] = (0.35, 0.35), normal=(0.0, 0.0, 1.0), color: tuple[float, float, float, float] = (0.15, 0.55, 0.95, 0.35)) -> int

Visual-only translucent plate at center.

The plate is a thin GEOM_BOX (1 mm thick along its normal), oriented so its thin axis lines up with normal. Useful for showing a plane manifold — e.g. z = z0 — in the scene.

Source code in autolife_planning/envs/pybullet_env.py

def draw_plane(
    self,
    center,
    half_sizes: tuple[float, float] = (0.35, 0.35),
    normal=(0.0, 0.0, 1.0),
    color: tuple[float, float, float, float] = (0.15, 0.55, 0.95, 0.35),
) -> int:
    """Visual-only translucent plate at *center*.

    The plate is a thin ``GEOM_BOX`` (1 mm thick along its normal),
    oriented so its thin axis lines up with *normal*.  Useful for
    showing a plane manifold — e.g. ``z = z0`` — in the scene.
    """
    client = self.sim.client
    vid = client.createVisualShape(
        shapeType=pb.GEOM_BOX,
        halfExtents=[float(half_sizes[0]), float(half_sizes[1]), 0.001],
        rgbaColor=list(color),
    )
    return client.createMultiBody(
        baseVisualShapeIndex=vid,
        basePosition=list(np.asarray(center, dtype=float)),
        baseOrientation=_quat_from_z_axis(normal),
    )

draw_rod(p1, p2, radius: float = 0.008, color: tuple[float, float, float, float] = (0.2, 0.9, 0.3, 1.0)) -> int | None

Visual-only solid rod (GEOM_CYLINDER) from p1 to p2.

Uses a real 3D cylinder rather than addUserDebugLine so the segment has a proper thickness, catches lighting, and shows up correctly in screenshots. Returns None if p1 and p2 coincide.

Source code in autolife_planning/envs/pybullet_env.py

def draw_rod(
    self,
    p1,
    p2,
    radius: float = 0.008,
    color: tuple[float, float, float, float] = (0.20, 0.90, 0.30, 1.0),
) -> int | None:
    """Visual-only solid rod (``GEOM_CYLINDER``) from *p1* to *p2*.

    Uses a real 3D cylinder rather than ``addUserDebugLine`` so the
    segment has a proper thickness, catches lighting, and shows up
    correctly in screenshots.  Returns ``None`` if *p1* and *p2*
    coincide.
    """
    p1 = np.asarray(p1, dtype=float)
    p2 = np.asarray(p2, dtype=float)
    delta = p2 - p1
    length = float(np.linalg.norm(delta))
    if length < 1e-9:
        return None
    client = self.sim.client
    vid = client.createVisualShape(
        shapeType=pb.GEOM_CYLINDER,
        radius=float(radius),
        length=length,
        rgbaColor=list(color),
    )
    return client.createMultiBody(
        baseVisualShapeIndex=vid,
        basePosition=(0.5 * (p1 + p2)).tolist(),
        baseOrientation=_quat_from_z_axis(delta / length),
    )

draw_sphere(center, radius: float, color: tuple[float, float, float, float] = (0.95, 0.3, 0.3, 0.55)) -> int

Visual-only sphere — handy for marking obstacles or targets.

Source code in autolife_planning/envs/pybullet_env.py

def draw_sphere(
    self,
    center,
    radius: float,
    color: tuple[float, float, float, float] = (0.95, 0.30, 0.30, 0.55),
) -> int:
    """Visual-only sphere — handy for marking obstacles or targets."""
    client = self.sim.client
    vid = client.createVisualShape(
        shapeType=pb.GEOM_SPHERE,
        radius=float(radius),
        rgbaColor=list(color),
    )
    return client.createMultiBody(
        baseVisualShapeIndex=vid,
        basePosition=list(np.asarray(center, dtype=float)),
    )

draw_frame(position, rotation, size: float = 0.12, radius: float = 0.006) -> list[int]

RGB coordinate axes at position with a 3x3 rotation matrix.

Red / green / blue rods along the rotation's first / second / third columns — useful for visualising an orientation-lock constraint (the frame at the start and goal look identical).

Source code in autolife_planning/envs/pybullet_env.py

def draw_frame(
    self,
    position,
    rotation,
    size: float = 0.12,
    radius: float = 0.006,
) -> list[int]:
    """RGB coordinate axes at *position* with a 3x3 *rotation* matrix.

    Red / green / blue rods along the rotation's first / second /
    third columns — useful for visualising an orientation-lock
    constraint (the frame at the start and goal look identical).
    """
    position = np.asarray(position, dtype=float)
    rotation = np.asarray(rotation, dtype=float)
    colors = [
        (1.0, 0.15, 0.15, 1.0),
        (0.15, 1.0, 0.15, 1.0),
        (0.15, 0.40, 1.0, 1.0),
    ]
    ids: list[int] = []
    for i, color in enumerate(colors):
        tip = position + rotation[:, i] * size
        body_id = self.draw_rod(position, tip, radius=radius, color=color)
        if body_id is not None:
            ids.append(body_id)
    return ids