Environments

Base Environment

BaseEnv

Bases: ABC

get_joint_states() -> np.ndarray abstractmethod

Get current joint positions.

Source code in drone_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 drone_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, z, yaw].

Source code in drone_planning/envs/base_env.py

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

get_obs() -> Any abstractmethod

Get observations.

Source code in drone_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)

Bases: BaseEnv

PyBullet scene with the X500.

The URDF carries the body pose as joints (base_x_joint … base_yaw_joint from world), so a full 6-DOF configuration is applied by setting joint positions.

Source code in drone_planning/envs/pybullet_env.py

def __init__(self, config: RobotConfig, visualize: bool = True):
    self.config = config
    urdf = _prepare_urdf_for_pybullet(config.urdf_path)
    self.sim = vpb.PyBulletSimulator(urdf, config.joint_names, visualize=visualize)
    self.joint_names = config.joint_names
    self.camera_link_idx = self.sim.link_map[config.camera.link_name]

    from drone_planning.x500 import HOME_JOINTS

    self.set_joint_states(HOME_JOINTS)

set_configuration(config: np.ndarray)

Apply a full 6-DOF config (4 body + 2 gimbal) to the visualization.

Source code in drone_planning/envs/pybullet_env.py

def set_configuration(self, config: np.ndarray):
    """Apply a full 6-DOF config (4 body + 2 gimbal) to the visualization."""
    self.set_joint_states(config)

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 drone_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)
    while client.isConnected():
        keys = client.getKeyboardEvents()
        if key_code in keys and keys[key_code] & pb.KEY_WAS_TRIGGERED:
            break
        time.sleep(0.01)

wait_for_close() -> None

Block until the user closes the PyBullet GUI window.

Source code in drone_planning/envs/pybullet_env.py

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

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 drone_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.
    """
    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))

    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
        ):
            return True
        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)
    return False

draw_box(center, half_extents, color: tuple[float, float, float, float] = (0.75, 0.75, 0.78, 1.0)) -> int

Visual-only axis-aligned box — e.g. a wall of the scene.

Source code in drone_planning/envs/pybullet_env.py

def draw_box(
    self,
    center,
    half_extents,
    color: tuple[float, float, float, float] = (0.75, 0.75, 0.78, 1.0),
) -> int:
    """Visual-only axis-aligned box — e.g. a wall of the scene."""
    client = self.sim.client
    vid = client.createVisualShape(
        shapeType=pb.GEOM_BOX,
        halfExtents=[float(h) for h in half_extents],
        rgbaColor=list(color),
    )
    return client.createMultiBody(
        baseVisualShapeIndex=vid,
        basePosition=list(np.asarray(center, dtype=float)),
    )