|
| 1 | +// Copyright (c) 2026 Robotics and AI Institute LLC dba RAI Institute. All rights reserved. |
| 2 | +#pragma once |
| 3 | + |
| 4 | +#include <condition_variable> |
| 5 | +#include <cstdint> |
| 6 | +#include <functional> |
| 7 | +#include <mutex> |
| 8 | +#include <thread> |
| 9 | + |
| 10 | +namespace exploy::control { |
| 11 | + |
| 12 | +/** |
| 13 | + * @brief Abstract base class for controller execution strategies. |
| 14 | + * |
| 15 | + * A Worker owns the read → work → write pipeline. Concrete subclasses decide |
| 16 | + * *when* and *how* the three callbacks are invoked: |
| 17 | + * |
| 18 | + * - `read_fn` — called on the main thread to snapshot robot state. |
| 19 | + * - `work_fn` — runs the ONNX inference. |
| 20 | + * - `write_fn` — called on the main thread to dispatch joint targets. |
| 21 | + * |
| 22 | + * Use `setCallbacks()` to register all three before calling `update()`. |
| 23 | + */ |
| 24 | +class Worker { |
| 25 | + public: |
| 26 | + virtual ~Worker() = default; |
| 27 | + |
| 28 | + /** @brief Reset the worker to its initial state. */ |
| 29 | + virtual void reset() {} |
| 30 | + |
| 31 | + /** |
| 32 | + * @brief Advance the control pipeline by one tick. |
| 33 | + * |
| 34 | + * @param time_us Current timestamp in microseconds. |
| 35 | + * @return `true` if the update succeeded or was safely skipped, `false` on error. |
| 36 | + */ |
| 37 | + virtual bool update(uint64_t time_us) = 0; |
| 38 | + |
| 39 | + /** |
| 40 | + * @brief Register the read / work / write callbacks. |
| 41 | + * |
| 42 | + * Must be called exactly once before the first `update()`. All three |
| 43 | + * arguments must be non-null; the function returns `false` otherwise. |
| 44 | + * |
| 45 | + * @param read_fn Reads observations from the robot state (main thread). |
| 46 | + * @param work_fn Runs ONNX inference (may execute on a background thread). |
| 47 | + * @param write_fn Writes joint targets back to the robot (main thread). |
| 48 | + * @return `true` on success. |
| 49 | + */ |
| 50 | + bool setCallbacks(std::function<bool()> read_fn, std::function<bool()> work_fn, |
| 51 | + std::function<bool()> write_fn); |
| 52 | + |
| 53 | + protected: |
| 54 | + std::function<bool()> read_fn_; |
| 55 | + std::function<bool()> work_fn_; |
| 56 | + std::function<bool()> write_fn_; |
| 57 | +}; |
| 58 | + |
| 59 | +/** |
| 60 | + * @brief Synchronous worker — runs read → work → write on the calling thread. |
| 61 | + * |
| 62 | + * All three callbacks execute inline inside `update()`. The call blocks until |
| 63 | + * inference is complete, so the caller's thread must afford the full inference |
| 64 | + * latency every control cycle. |
| 65 | + * |
| 66 | + * Phase is maintained across updates: the first call initialises the phase |
| 67 | + * reference and subsequent calls fire whenever the elapsed time exceeds the |
| 68 | + * configured period. |
| 69 | + */ |
| 70 | +class SyncWorker : public Worker { |
| 71 | + public: |
| 72 | + /** |
| 73 | + * @param update_rate_hz Desired control frequency in Hz. |
| 74 | + */ |
| 75 | + explicit SyncWorker(double update_rate_hz); |
| 76 | + |
| 77 | + bool update(uint64_t time_us) override; |
| 78 | + void reset() override; |
| 79 | + |
| 80 | + private: |
| 81 | + uint64_t period_ms_; |
| 82 | + uint64_t last_scheduled_update_us_ = 0; |
| 83 | + bool first_run_ = true; |
| 84 | +}; |
| 85 | + |
| 86 | +/** |
| 87 | + * @brief Asynchronous worker — offloads ONNX inference to a background thread. |
| 88 | + * |
| 89 | + * The pipeline is split across two consecutive `update()` calls: |
| 90 | + * |
| 91 | + * 1. **First call at cycle boundary** — `read_fn` executes on the main thread, |
| 92 | + * then `work_fn` is dispatched to a dedicated background thread. |
| 93 | + * 2. **Subsequent calls** — if inference has finished, `write_fn` runs on the |
| 94 | + * main thread to commit joint targets. If the worker is still busy |
| 95 | + * (overrun), the cycle is skipped and `update()` returns `true`. |
| 96 | + * |
| 97 | + * This decouples the main control loop from the inference latency, allowing |
| 98 | + * the robot to keep receiving state updates while the GPU or CPU is busy. |
| 99 | + * |
| 100 | + * Thread safety: the internal mutex guards all shared state between the main |
| 101 | + * thread and the worker thread. |
| 102 | + * |
| 103 | + * `reset()` stops the background thread and clears all state. The thread is |
| 104 | + * re-started on the next `update()` call. |
| 105 | + * |
| 106 | + * **Error handling**: when `work_fn` returns false the worker latches into a |
| 107 | + * faulted state. All subsequent `update()` calls return `false` immediately |
| 108 | + * without dispatching new work. Call `reset()` to clear the fault and resume. |
| 109 | + */ |
| 110 | +class AsyncWorker : public Worker { |
| 111 | + public: |
| 112 | + /** |
| 113 | + * @param update_rate_hz Desired control frequency in Hz. |
| 114 | + */ |
| 115 | + explicit AsyncWorker(double update_rate_hz); |
| 116 | + ~AsyncWorker() override; |
| 117 | + |
| 118 | + void reset() override; |
| 119 | + bool update(uint64_t time_us) override; |
| 120 | + |
| 121 | + private: |
| 122 | + void startWorker(); |
| 123 | + void stopWorker(); |
| 124 | + void threadLoop(); |
| 125 | + |
| 126 | + // Main-thread-only — no synchronization needed. |
| 127 | + uint64_t period_ms_; |
| 128 | + uint64_t last_scheduled_update_us_ = 0; |
| 129 | + bool first_run_ = true; |
| 130 | + uint64_t consecutive_overruns_ = 0; |
| 131 | + std::thread thread_; |
| 132 | + |
| 133 | + // Shared between the main thread and the worker thread — all guarded by mutex_. |
| 134 | + std::mutex mutex_; |
| 135 | + std::condition_variable cv_; |
| 136 | + bool stop_ = false; |
| 137 | + bool working_ = false; |
| 138 | + bool work_requested_ = false; |
| 139 | + bool work_finished_ = false; |
| 140 | + bool work_successful_ = true; |
| 141 | + bool faulted_ = false; |
| 142 | +}; |
| 143 | + |
| 144 | +} // namespace exploy::control |
0 commit comments