ros.node

This file contains the estimator node for ROS 2.

Essentially this file provides a wrapper for the estimators to be used with ROS data. The estimates get published to ROS as well.

TODO Subscribed and published topics...

MPEstimator(settings, stop_event=None)

This class starts a batched estimator for all drones in settings.drone_names and the necessary subscribers and publishers.

Note: Since all results are published individually, one might want to synchronize on the receiving end. For that see https://docs.ros.org/en/rolling/p/flex_sync/ https://github.com/ros2/message_filters

In the future, a custom message type might be advantageous. However, to keep it simple, we use standard message types.

Start the subscriber and publisher processes.

Parameters:

Name	Type	Description	Default
settings	Munch	The estimator settings, including the names of all drones in the batch.	required
stop_event	Event | None	Event that stops the estimator and its processes once set. If None, a new event is created.	None

Source code in drone_estimators/ros/node.py

def __init__(self, settings: Munch, stop_event: Event | None = None):
    """Start the subscriber and publisher processes.

    Args:
        settings: The estimator settings, including the names of all drones in the batch.
        stop_event: Event that stops the estimator and its processes once set. If None, a new
            event is created.
    """
    self.settings = settings
    self.drone_names = settings.drone_names
    n_drones = len(self.drone_names)

    ctx = mp.get_context("spawn")
    self._stop_event = ctx.Event() if stop_event is None else stop_event
    self._publish_update = ctx.Event()
    startup = ctx.Barrier(3)  # Main process, _subscriber_loop, _publisher_loop

    # Logger setup
    self.logger = logging.getLogger("ESTIMATOR" + "_" + "_".join(self.drone_names))
    self.logger.setLevel(logging.INFO)
    # Create console handler
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    # Create formatter with [name][LEVEL] prefix
    formatter = logging.Formatter("[%(name)s][%(levelname)s] %(message)s")
    # Add formatter to handler
    console_handler.setFormatter(formatter)
    # Add handler to logger (only once)
    if not self.logger.hasHandlers():
        self.logger.addHandler(console_handler)
    # All buffers contain one row per drone, see _view.
    # We allocate a buffer that is used to store the most recent tf message from the tf
    # subscriber. The message contains four fields: The first element is the number of messages
    # since the array was last cleared. This allows us to throw warnings in case the estimator
    # loop cannot keep up. The second field is the timestamp in s, the third is the position
    # (3d) and the fourth is the quaternion (4d)
    self._tf_msg_buffer = ctx.Array("d", [0.0] * (1 + 1 + 3 + 4) * n_drones)
    # Allocate the command subscriber buffer
    cmd_dim = 4
    self._cmd_msg_buffer = ctx.Array("d", [0.0] * (1 + 1 + cmd_dim) * n_drones)
    # Estimated state publisher buffers
    self._pose_buffer = ctx.Array("d", [0.0] * 7 * n_drones)
    self._twist_buffer = ctx.Array("d", [0.0] * 6 * n_drones)
    self._rotor_vel_buffer = ctx.Array("d", [0.0] * 4 * n_drones)  # Rotor velocities, optional
    self._wrench_buffer = ctx.Array("d", [0.0] * 6 * n_drones)  # External wrench, optional

    args = (
        self.drone_names,
        self._tf_msg_buffer,
        self._cmd_msg_buffer,
        startup,
        self._stop_event,
    )
    self._sub_process = ctx.Process(target=self._subscriber_loop, args=args)

    args = (
        self.drone_names,
        self._pose_buffer,
        self._twist_buffer,
        self._rotor_vel_buffer,
        self._wrench_buffer,
        self._publish_update,
        startup,
        self._stop_event,
    )
    self._pub_process = ctx.Process(target=self._publisher_loop, args=args)

    self._sub_process.start()
    self._pub_process.start()
    startup.wait(10.0)
    self.logger.info("Subscriber and Publisher process started.")

close()

Stop the subscriber and publisher processes and save the data if enabled.

Source code in drone_estimators/ros/node.py

def close(self):
    """Stop the subscriber and publisher processes and save the data if enabled."""
    self.logger.info(f"Estimator {self.drone_names} shutdown")
    self._stop_event.set()
    self._sub_process.join()
    self._pub_process.join()

    if self.settings.save_data:
        self.logger.info("Saving data...")
        for i, drone_name in enumerate(self.drone_names):
            filename = f"data_{drone_name}_"
            info = f"{self.settings.estimator_type}"
            if self.settings.estimator_type != "diff":
                info = info + f"_{self.settings.dynamics_model}"
            with open(filename + info + ".pkl", "wb") as f:
                pickle.dump(self.data_est[i], f)
            with open(filename + "measurement" + ".pkl", "wb") as f:
                pickle.dump(self.data_meas[i], f)

run()

Main estimator loop.

Source code in drone_estimators/ros/node.py

def run(self):
    """Main estimator loop."""
    self._init_estimator()  # done here such that errors can be raised properly

    names = np.array(self.drone_names)
    n_drones = len(self.drone_names)
    tf_msg_buffer = _view(self._tf_msg_buffer, n_drones)
    cmd_msg_buffer = _view(self._cmd_msg_buffer, n_drones)
    pose_buffer = _view(self._pose_buffer, n_drones)
    twist_buffer = _view(self._twist_buffer, n_drones)
    rotor_vel_buffer = _view(self._rotor_vel_buffer, n_drones)
    wrench_buffer = _view(self._wrench_buffer, n_drones)

    k = 0
    global_time = time.perf_counter()
    # The loop sleeps until an absolute deadline that advances by one period each cycle, so a
    # slow cycle is followed by a shorter sleep instead of delaying all subsequent cycles.
    period = 1 / self.frequency
    next_deadline = time.time() + period
    # time.sleep overshoots the requested duration. The overshoot of the previous sleep is
    # measured at the start of the next cycle and compensated, low pass filtered to smooth it
    sleep_overshoot = 0.0
    sleep_end = None  # Expected end of the previous sleep, None if the cycle did not sleep

    try:
        # self.logger.info(f"{self.estimator.data.pos=}")
        # Estimation loop
        while not self._stop_event.is_set():
            loop_start_time = time.time()

            with self._tf_msg_buffer.get_lock():
                data = tf_msg_buffer.copy()
                tf_msg_buffer[:, 0] = 0
            n_tf_msg, tf_timestamp, pos, quat = (
                data[:, 0],
                data[:, 1],
                data[:, 2:5],
                data[:, 5:],
            )
            with self._cmd_msg_buffer.get_lock():
                data = cmd_msg_buffer.copy()
                cmd_msg_buffer[:, 0] = 0
            n_cmd_messages, cmd_timestep, cmd = data[:, 0], data[:, 1], data[:, 2:]
            has_tf, has_cmd = n_tf_msg >= 1, n_cmd_messages >= 1

            outdated = (cmd_timestep < tf_timestamp - 1) & (cmd_timestep > 0)
            if np.any(outdated) and self.input_needed:
                self.logger.warning(
                    f"Last command of {names[outdated]} is older than 1s. Assuming zeros as input."
                )
                u = np.array(self.estimator.data.u)
                u[outdated] = 0.0
                self.estimator.set_input(u)

            if np.any(has_cmd) and self.input_needed:
                # The command is as it is sent to the drone, meaning for attitude interface:
                # roll (deg), pitch (deg), yaw (deg), thrust (PWM)
                # All the models run with rad and N, so we need to convert the RPYT command
                cmd[..., -1] = pwm2force(
                    cmd[..., -1], self.params["thrust_max"] * 4, self.params["pwm_max"]
                )
                cmd[..., :-1] = np.deg2rad(cmd[..., :-1])
                u = np.array(self.estimator.data.u)
                u[has_cmd] = cmd[has_cmd]
                self.estimator.set_input(u)  # TODO # compare times?

            if np.any(n_tf_msg > 2):
                self.logger.warning(
                    f"Dropping tf messages of {names[n_tf_msg > 2]} because estimator loop can't keep up"
                )
                # TODO check for frequencies. If Vicon is running at higher frequency, of course estimator cant keep up

            if np.any(has_tf):
                # Drones without a measurement have dt = 0 and are neither predicted nor corrected
                dt = np.where(has_tf, tf_timestamp - self.time_stamp_last_prediction, 0.0)
                self.time_stamp_last_prediction = np.where(
                    has_tf, tf_timestamp, self.time_stamp_last_prediction
                )

                self.estimator.predict(dt)
                self.estimator.correct(pos, quat, mask=has_tf)

                # estimator_data = self.estimator.step(pos, quat, dt)

            time_stamp_now = time.time()
            dt = time_stamp_now - self.time_stamp_last_prediction
            self.time_stamp_last_prediction = np.full(n_drones, time_stamp_now)
            estimator_data = self.estimator.predict(dt)

            # Giving new estimate to publisher
            with self._pose_buffer.get_lock():
                pose_buffer[:, :3] = estimator_data.pos
                pose_buffer[:, 3:] = estimator_data.quat
            with self._twist_buffer.get_lock():
                twist_buffer[:, :3] = estimator_data.vel
                twist_buffer[:, 3:] = estimator_data.ang_vel
            if estimator_data.rotor_vel is not None:
                with self._rotor_vel_buffer.get_lock():
                    rotor_vel_buffer[:] = estimator_data.rotor_vel
            if estimator_data.dist_f is not None:
                with self._wrench_buffer.get_lock():
                    wrench_buffer[:, :3] = estimator_data.dist_f
                    if estimator_data.dist_t is not None:
                        wrench_buffer[:, 3:] = estimator_data.dist_t
            self._publish_update.set()

            if self.settings.save_data:
                estimator_data = jax.tree.map(np.asarray, estimator_data)
                for i in range(n_drones):
                    estimate_i = jax.tree.map(lambda x: x[i], estimator_data)
                    append_state(self.data_est[i], time_stamp_now, estimate_i)
                    if has_tf[i]:
                        if has_cmd[i] and self.input_needed:
                            append_measurement(
                                self.data_meas[i], tf_timestamp[i], pos[i], quat[i], cmd[i]
                            )
                        if not self.input_needed:
                            append_measurement(
                                self.data_meas[i], tf_timestamp[i], pos[i], quat[i], None
                            )

            k += 1
            if k == 1000:
                self.logger.info(f"Freq: {k / (time.perf_counter() - global_time)}")
                k, global_time = 0, time.perf_counter()

            # Sleep until next cycle
            if sleep_end is not None:  # Overshoot of the last sleep until this cycle started
                overshoot = np.clip(loop_start_time - sleep_end, 0.0, 1e-3)
                sleep_overshoot += 0.1 * (overshoot - sleep_overshoot)
            sleep_end = None
            now = time.time()
            if now > next_deadline + 0.1:  # Fell far behind (e.g. jit compile), do not catch up
                next_deadline = now
            remaining = next_deadline - now - sleep_overshoot
            next_deadline += period
            if remaining > 0:
                sleep_end = now + remaining
                time.sleep(remaining)
    except KeyboardInterrupt:
        self._stop_event.set()

launch_estimators(estimators)

Launch one process per batch of estimators and block until a keyboard interrupt.

Parameters:

Name	Type	Description	Default
estimators	dict	The settings of each estimator, keyed by name. Estimators with identical settings (except for the drone name) are run as one batch.	required

Source code in drone_estimators/ros/node.py

def launch_estimators(estimators: dict):
    """Launch one process per batch of estimators and block until a keyboard interrupt.

    Args:
        estimators: The settings of each estimator, keyed by name. Estimators with identical
            settings (except for the drone name) are run as one batch.
    """
    processes = []
    seen_drone_names = []
    ctx = mp.get_context("spawn")
    stop_event = ctx.Event()

    # Estimators with identical settings (except for the drone name) are run as one batch
    batches = {}
    for k, settings in estimators.items():
        name = settings.drone_name
        if name in seen_drone_names:
            print(f"[ESTIMATOR_{name}]  Estimator for {name} already existing. Check settings")
            continue
        seen_drone_names.append(name)
        batch_settings = {k: v for k, v in settings.items() if k != "drone_name"}
        key = repr(sorted(batch_settings.items()))
        batches.setdefault(key, munchify(batch_settings | {"drone_names": []}))
        batches[key].drone_names.append(name)

    try:
        for settings in batches.values():
            # not sure if daemon should be True or False
            p = ctx.Process(target=launch_node, args=(settings, stop_event))
            processes.append(p)
            p.start()

        try:
            while True:
                time.sleep(10.0)
        except KeyboardInterrupt:
            print("\nKeyboard interrupt received. Terminating nodes...")

    finally:
        stop_event.set()

        for p in processes:
            p.join(timeout=2)

        for p in processes:
            if p.is_alive():
                print(f"Force terminating process {p.pid}")
                p.terminate()
                p.join()

        if rclpy.ok():
            rclpy.shutdown()
        print("All nodes terminated.")

launch_node(settings, stop_event)

Run an estimator node until the stop event is set.

Parameters:

Name	Type	Description	Default
settings	Munch	The estimator settings, including the names of all drones in the batch.	required
stop_event	Event	Event that stops the estimator and its processes once set.	required

Source code in drone_estimators/ros/node.py

def launch_node(settings: Munch, stop_event: Event):
    """Run an estimator node until the stop event is set.

    Args:
        settings: The estimator settings, including the names of all drones in the batch.
        stop_event: Event that stops the estimator and its processes once set.
    """
    rclpy.init()
    estimator = MPEstimator(settings, stop_event=stop_event)
    try:
        estimator.run()
    finally:
        estimator.close()