diff --git a/examples/inference/TAM.md b/examples/inference/TAM.md new file mode 100644 index 00000000..d7c2e982 --- /dev/null +++ b/examples/inference/TAM.md @@ -0,0 +1,81 @@ +# TAM (Torque Adaptation Module) integration + +[TAM](https://github.com/Dongwon-Son/TAM) adds a learned residual to the +commanded joint torque at 1 kHz inside the RCS controller, conditioned on a +latent computed from the recent control history. One process, two rates: + +- **1 kHz (C++ control thread)**: `Franka::tam_forward` runs a small MLP on + q, dq, the last 8 commanded torques (gravity-included space) and the + current latent, and adds the clipped residual to the controller torque. +- **~5 Hz (Python thread in `franka_tam.py`)**: `run_history_encoder` pulls + the controller's history buffer, streams it through the TAM transformer + encoder (JAX), and pushes each latent back via `set_tam_latent`. + +## Setup + +```shell +pip install -r requirements.txt # includes torque-adaptation-module (JAX) +``` + +A CUDA-capable JAX is strongly recommended on the control machine — the +encoder is ~20x slower than real time on CPU: + +```shell +pip install "jax[cuda12]" +``` + +## Run + +Set `tam = True` on `InferenceConfig` (top of `franka_tam.py`; the file does +not read `franka.json`). Everything else resolves automatically on first run: + +- the default checkpoint (DAgger-finetuned, applied-torque) downloads once + from the TAM GitHub releases into `~/.cache/simadaptor` (override with + `SIMADAPTOR_CACHE_DIR`), verified by SHA-256; +- the ideal-model MJCF is installed with the `torque-adaptation-module` + package. + +To use a different checkpoint or MJCF, pass them to +`RealTimeHistoryAdaptor.from_checkpoint(...)` in `_init_tam`. +The per-joint residual is clipped at 10/10/10/10/2/2/2 Nm +(`FrankaConfig.tam_residual_clip`, a C++-side default); the residual also ramps in over 1 s whenever it +(re)activates. + +## Operational notes + +- The controller applies zero residual until the MLP weights and the first + latent arrive. The first latent needs only ~0.5 s of control history (one + 400 ms encoder patch plus a poll); the estimate then keeps refining as + context grows toward the 4 s attention window. Startup also pays a + one-time JAX JIT warm-up of ~10-15 s before the encoder loop begins. +- Switching controller gains (`pd_mode`) restarts the control thread, but + TAM adaptation continues across it: history timestamps come from a + robot-lifetime monotonic clock, the encoder bridges the short restart + gap with masked padding rows (so its context window is not cut), and the + latent is kept throughout (it encodes plant properties, which a gain + change does not alter). Only the residual re-ramps over 1 s after the + switch, since the gains it interacts with changed. +- Only applied-torque checkpoints are supported; `base_tam_fusion` + checkpoints are rejected at startup. +- The residual MLP adds ~0.1-0.3 ms to every 1 kHz control tick, which + leaves no deadline slack on a stock (non-PREEMPT_RT) kernel. The control + thread therefore elevates itself to a real-time scheduling class by + default (`FrankaConfig.rt_priority`, default 80; 0 disables), trying in + order: + + 1. `SCHED_FIFO` at the configured priority — needs an rtprio rlimit; + 2. RealtimeKit (`SCHED_RR`, the mechanism the desktop audio stack uses) + — **no host configuration needed** on a normal desktop session. + + So on a desktop workstation this works out of the box. Only on headless + / SSH-only setups (where RealtimeKit's policy denies the request) do the + one-time rlimit setup: + + ```shell + echo "$USER - rtprio 99" | sudo tee -a /etc/security/limits.conf + # then open a fresh login session and check: ulimit -r -> 99 + ``` + + If both mechanisms are unavailable the controller prints a warning and + stays on the normal scheduler; expect `communication_constraints_violation` + aborts on a loaded machine in that state. diff --git a/examples/inference/franka_tam.py b/examples/inference/franka_tam.py index cc593b62..adc3a13a 100644 --- a/examples/inference/franka_tam.py +++ b/examples/inference/franka_tam.py @@ -1,6 +1,13 @@ import copy import logging +import os import threading +import time + +# The TAM history encoder (JAX) shares the GPU with everything else on this +# machine; never let JAX preallocate ~75% of device memory. Must be set +# before JAX initializes its backend (first use inside _init_tam). +os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" from dataclasses import dataclass, field from typing import Any @@ -65,6 +72,11 @@ class InferenceConfig: n_action_steps: int | None = None max_rel_mov_joints: float = MAX_REL_MOV_JOINTS max_rel_mov_cart: tuple[float, float] = MAX_REL_MOV_CART + # TAM (Torque Adaptation Module): master switch. When enabled, everything + # else resolves automatically: the default checkpoint is downloaded once + # into ~/.cache/simadaptor and the ideal-model MJCF is installed with the + # torque-adaptation-module package. + tam: bool = False def build_vlagent_obs( @@ -129,29 +141,59 @@ def __init__(self, env: gym.Env, cfg: InferenceConfig): self.frame_rate = SimpleFrameRate(self._cfg.fps) self._action_buffer = [] self._prev_pd_mode = 1.0 - # TODO: load history encoder - self.history_encoder = None + self.tam_runtime = None + if cfg.tam: + self._init_tam() + + def _init_tam(self) -> None: + """Load the TAM checkpoint: the streaming history encoder runs in this + process (JAX; a GPU is strongly recommended) and the adaptor MLP is + exported once into the C++ controller.""" + from simadaptor.deploy.history_runtime import RealTimeHistoryAdaptor + + # from_checkpoint enforces an applied-torque checkpoint: this + # integration records a single commanded-torque stream. + self.tam_runtime = RealTimeHistoryAdaptor.from_checkpoint() + logger.info( + "TAM ready: adaptor binary %d bytes, applied-torque mode", + len(self.tam_runtime.adaptor_weight_bytes()), + ) def run_history_encoder(self): + """Feed the 1 kHz controller history through the TAM history encoder + (~5 Hz) and push the resulting latent back into the C++ controller. + + The C++ side applies zero residual until both the MLP weights and the + first latent have arrived, then ramps the residual in over 1 s.""" + if self.tam_runtime is None: + logger.info("TAM disabled; history encoder not started") + return robot: Franka = self.env.get_wrapper_attr("envs")["right"].get_wrapper_attr("robot")() - # TODO: load TAM weights from disk - weights = None - - # flatten array to send to cpp - # attention: numpy vs eigen has col vs row major (numpy) how values are stored in ram - # and its easier to fix this in python with indexing - robot.set_tam_mlp_weight(weights.reshape((-1,))) - hist_encoder_framerate = SimpleFrameRate(5) # 5hz + + # The weight vector carries the packed adaptor binary, one byte per + # float64 element (parsed and validated on the C++ side). + robot.set_tam_mlp_weight( + np.frombuffer(self.tam_runtime.adaptor_weight_bytes(), dtype=np.uint8).astype(np.float64) + ) + + hist_encoder_framerate = SimpleFrameRate(5) # 5 Hz latent updates + latents_sent = 0 while True: - hist = robot.get_tam_history() - # TODO convert in TAM suitable dataformat - latent = self.history_encoder(hist) - - # flatten array to send to cpp - # attention: numpy vs eigen has col vs row major (numpy) how values are stored in ram - # and its easier to fix this in python with indexing - latent = latent.numpy().reshape((-1,)) - robot.set_tam_latent(latent) + try: + # The runtime handles everything stream-related internally: + # overlapping polls are deduplicated by timestamp, short holes + # (e.g. a controller restart during a gain switch) are bridged + # with masked padding on its dense grid, and a backwards or + # over-long gap restarts the stream. + latent = self.tam_runtime.push_history_samples(robot.get_tam_history()) + if latent is not None: + robot.set_tam_latent(np.asarray(latent, dtype=np.float64).reshape((-1,))) + latents_sent += 1 + if latents_sent == 1: + logger.info("TAM: first latent sent, residual ramps in on the controller") + except Exception: + logger.exception("TAM history encoder step failed; retrying") + time.sleep(0.5) hist_encoder_framerate() def obs_rcs2agents(self, obs: dict, info: dict | None = None) -> Obs: @@ -320,6 +362,7 @@ def get_env(cfg: InferenceConfig) -> gym.Env: "right": rcs.common.Pose(translation=np.array([0, 0, 0]), rpy_vector=np.array([0, 0, 0])), } hw_cfg.robot_cfgs["right"].ignore_realtime = True + hw_cfg.robot_cfgs["right"].tam_enabled = cfg.tam hw_cfg.robot_cfgs["right"].speed_factor = 0.4 hw_cfg.robot_cfgs["right"].policy_rate = cfg.fps env_rel = env_creator.create_env(hw_cfg) @@ -363,10 +406,11 @@ def main() -> None: env_rel = get_env(cfg) controller = ModelInference(env_rel, cfg) - history_encoder_thread = threading.Thread( - target=controller.run_history_encoder, name="history_encoder", daemon=True - ) - history_encoder_thread.start() + if cfg.tam: + history_encoder_thread = threading.Thread( + target=controller.run_history_encoder, name="history_encoder", daemon=True + ) + history_encoder_thread.start() with env_rel: controller.loop() diff --git a/examples/inference/requirements.txt b/examples/inference/requirements.txt index 2bc1996e..851c93a4 100644 --- a/examples/inference/requirements.txt +++ b/examples/inference/requirements.txt @@ -1 +1,3 @@ -vlagents==0.3.0 \ No newline at end of file +vlagents==0.3.0 +# TAM history encoder + adaptor export (installs the `simadaptor` package) +torque-adaptation-module @ git+https://github.com/Dongwon-Son/TAM diff --git a/extensions/rcs_fr3/src/hw/Franka.cpp b/extensions/rcs_fr3/src/hw/Franka.cpp index 171f2b0f..d3a74c61 100644 --- a/extensions/rcs_fr3/src/hw/Franka.cpp +++ b/extensions/rcs_fr3/src/hw/Franka.cpp @@ -19,6 +19,7 @@ namespace rcs { namespace hw { + common::Pose GetFlangeInBaseFrame(const franka::RobotState& robot_state) { return common::Pose(robot_state.O_T_EE) * common::Pose(robot_state.F_T_EE).inverse(); @@ -93,22 +94,86 @@ FrankaState* Franka::get_state() { return state; } -common::Vector7d Franka::tam_forward(const std::array& tau) { - // access weight matrix thread safe - const std::optional weight = this->tam_mlp_weight.load(); - - // no weights set yet, so TAM does not contribute any torque - if (!weight.has_value()) { - return common::Vector7d::Zero(); +void Franka::set_tam_mlp_weight(const Eigen::VectorXd& weight) { + this->tam_mlp_weight.store(weight); + // The vector carries the packed TAM adaptor binary, one byte per element + // (the format written by SimAdaptorInference.export_simadaptor_weights_cpp; + // it includes the MLP architecture and the input normalization statistics). + // Parse here, off the 1 kHz control thread, and publish the parsed model. + std::string bytes(static_cast(weight.size()), '\0'); + for (Eigen::Index i = 0; i < weight.size(); ++i) { + const double v = weight(i); + if (!(v >= 0.0 && v <= 255.0)) { + throw std::runtime_error( + "set_tam_mlp_weight: element " + std::to_string(i) + + " is not a byte value; pass the .bin file bytes as float64"); + } + bytes[static_cast(i)] = static_cast(static_cast(v)); } + std::shared_ptr model = + adaptor::SimAdaptor::LoadFromMemory(bytes.data(), bytes.size()); + if (!model) { + throw std::runtime_error( + "set_tam_mlp_weight: could not parse the TAM adaptor binary"); + } + if (model->dof != 7) { + throw std::runtime_error("set_tam_mlp_weight: expected dof=7, got " + + std::to_string(model->dof)); + } + this->tam_model.store(model); +} - // access latent thread safe - const Eigen::VectorXd latent = this->tam_latent.load(); +common::Vector7d Franka::tam_forward(const std::array& tau, + const franka::RobotState& robot_state, + const std::array& gravity) { + const std::shared_ptr model = + this->tam_model.load(); + common::Vector7d delta = common::Vector7d::Zero(); + std::vector rows; + if (model) { + // The MLP conditions on the last history_steps samples in the + // ideal-model (gravity-included) torque space: history_steps-1 recorded + // ticks plus the current one from the arguments. The newest recorded row + // must be contiguous with now — right after a controller restart the + // buffer still ends before the restart gap (see tam_now()). + const int T = model->history_steps; + const std::deque& past = this->tam_recent; + if (past.size() + 1 >= static_cast(T) && + this->tam_now() - past.back().t <= 0.05) { + rows.reserve(static_cast(T)); + for (auto it = past.end() - static_cast(T - 1); it != past.end(); + ++it) { + const TAMHistorySample& s = *it; + adaptor::SimAdaptor::StreamRow r; + for (int j = 0; j < 7; ++j) { + r.q[j] = s.q[j]; + r.dq[j] = s.dq[j]; + r.tau_model[j] = s.tau_cmd[j] + s.gravity[j]; + } + rows.push_back(r); + } + adaptor::SimAdaptor::StreamRow now; + for (int j = 0; j < 7; ++j) { + now.q[j] = robot_state.q[j]; + now.dq[j] = robot_state.dq[j]; + now.tau_model[j] = tau[j] + gravity[j]; + } + rows.push_back(now); + } + } - // TODO reshape weight and latent to match nn - // TODO forward nn - common::Vector7d tam_tau = common::Vector7d::Zero(); - return tam_tau; + const bool active = + !rows.empty() && model->forward_stream(rows, this->tam_latent.load(), + this->m_cfg.tam_residual_clip, + delta); + if (!active) { + this->tam_active_ticks = 0; + return common::Vector7d::Zero(); + } + // Ramp the residual in over ~1 s whenever it (re)activates so enabling + // TAM (or the first latent) never steps the torque. + this->tam_active_ticks = std::min(this->tam_active_ticks + 1, 1000); + return (static_cast(this->tam_active_ticks) / 1000.0) * delta; } void Franka::set_default_robot_behavior() { @@ -413,6 +478,7 @@ void Franka::stop_control_thread() { } void Franka::osc() { + adaptor::TryElevateControlThreadPriority(this->m_cfg.rt_priority); franka::Model model = this->robot.loadModel(); const Eigen::Vector3d kp_p_cfg = this->m_cfg.kp_p; const double kp_r_cfg = this->m_cfg.kp_r; @@ -420,6 +486,11 @@ void Franka::osc() { const bool allow_high_collision = this->m_cfg.allow_high_collision; this->controller_time = 0.0; + // The TAM history buffer intentionally survives controller restarts: + // timestamps come from a robot-lifetime monotonic clock, so a restart is + // just a short gap in a continuous stream. Only the residual ramp starts + // over (the gains may have changed). + this->tam_active_ticks = 0; // conservative collision and impedance behavior this->set_default_robot_behavior(); @@ -629,7 +700,8 @@ void Franka::osc() { std::chrono::duration_cast(t2 - t1); if (this->m_cfg.tam_enabled) { - Eigen::VectorXd::Map(&tau_d_array[0], 7) += tam_forward(tau_d_array); + Eigen::VectorXd::Map(&tau_d_array[0], 7) += + tam_forward(tau_d_array, robot_state, gravity_array); } std::array tau_d_rate_limited = franka::limitRate( @@ -639,12 +711,16 @@ void Franka::osc() { if (this->m_cfg.tam_enabled) { // safe q, q_dot and tau - this->tam_history.push_back( - TAMHistorySample{.t = period.toSec(), - .q = robot_state.q, - .dq = robot_state.dq, - .tau_cmd = tau_d_rate_limited, - .gravity = gravity_array}); + const TAMHistorySample sample{.t = this->tam_now(), + .q = robot_state.q, + .dq = robot_state.dq, + .tau_cmd = tau_d_rate_limited, + .gravity = gravity_array}; + this->tam_history.push_back(sample); + this->tam_recent.push_back(sample); + if (this->tam_recent.size() > 32) { + this->tam_recent.pop_front(); + } } return tau_d_rate_limited; @@ -658,12 +734,18 @@ void Franka::osc() { } void Franka::joint_controller() { + adaptor::TryElevateControlThreadPriority(this->m_cfg.rt_priority); franka::Model model = this->robot.loadModel(); const common::Vector7d Kp = this->m_cfg.kp; const common::Vector7d Kd = this->m_cfg.kd; const common::Vector7d torque_limit = this->m_cfg.torque_limit; const bool allow_high_collision = this->m_cfg.allow_high_collision; this->controller_time = 0.0; + // The TAM history buffer intentionally survives controller restarts: + // timestamps come from a robot-lifetime monotonic clock, so a restart is + // just a short gap in a continuous stream. Only the residual ramp starts + // over (the gains may have changed). + this->tam_active_ticks = 0; // conservative collision and impedance behavior this->set_default_robot_behavior(); @@ -738,8 +820,10 @@ void Franka::joint_controller() { auto time = std::chrono::duration_cast(t2 - t1); + std::array gravity_array = model.gravity(robot_state); if (this->m_cfg.tam_enabled) { - Eigen::VectorXd::Map(&tau_d_array[0], 7) += tam_forward(tau_d_array); + Eigen::VectorXd::Map(&tau_d_array[0], 7) += + tam_forward(tau_d_array, robot_state, gravity_array); } std::array tau_d_rate_limited = franka::limitRate( @@ -749,12 +833,16 @@ void Franka::joint_controller() { if (this->m_cfg.tam_enabled) { // safe q, q_dot and tau - this->tam_history.push_back( - TAMHistorySample{.t = period.toSec(), - .q = robot_state.q, - .dq = robot_state.dq, - .tau_cmd = tau_d_rate_limited, - .gravity = model.gravity(robot_state)}); + const TAMHistorySample sample{.t = this->tam_now(), + .q = robot_state.q, + .dq = robot_state.dq, + .tau_cmd = tau_d_rate_limited, + .gravity = gravity_array}; + this->tam_history.push_back(sample); + this->tam_recent.push_back(sample); + if (this->tam_recent.size() > 32) { + this->tam_recent.pop_front(); + } } return tau_d_rate_limited; @@ -778,6 +866,7 @@ void Franka::zero_torque_guiding() { } void Franka::zero_torque_controller() { + adaptor::TryElevateControlThreadPriority(this->m_cfg.rt_priority); this->set_default_robot_behavior(); if (this->m_cfg.allow_high_collision) { // High collision threshold values for high impedance. diff --git a/extensions/rcs_fr3/src/hw/Franka.h b/extensions/rcs_fr3/src/hw/Franka.h index ca3e440c..5193bd54 100644 --- a/extensions/rcs_fr3/src/hw/Franka.h +++ b/extensions/rcs_fr3/src/hw/Franka.h @@ -5,6 +5,8 @@ #include #include +#include +#include #include #include #include @@ -13,6 +15,7 @@ #include #include "rcs/Kinematics.h" +#include "simadaptor.h" #include "rcs/LinearPoseTrajInterpolator.h" #include "rcs/Pose.h" #include "rcs/Robot.h" @@ -30,7 +33,10 @@ struct TAMHistorySample { }; const double DEFAULT_SPEED_FACTOR = 0.2; -const size_t TAM_HISTORY_SIZE = 200; +// 4 s of 1 kHz samples: the history-encoder thread polls at ~5 Hz and +// tolerates late polls without losing stream continuity (200 rows gave +// the Python side a zero-margin 0.2 s window). +const size_t TAM_HISTORY_SIZE = 4000; struct FrankaLoad { double load_mass; @@ -85,7 +91,19 @@ struct FrankaConfig : common::RobotConfig { double approach_cartesian_speed = 0.1; double approach_rotation_speed = 0.5; bool tam_enabled = false; + // Per-joint |clip| of the TAM residual torque in Nm before it is added to + // the controller torque (wrist joints have a 12 Nm limit; keep headroom). + common::Vector7d tam_residual_clip = + (common::Vector7d() << 10., 10., 10., 10., 2., 2., 2.).finished(); bool ignore_realtime = false; + // Best-effort real-time scheduling for the async control thread at this + // priority (0 disables): SCHED_FIFO when the rtprio rlimit allows it, + // otherwise SCHED_RR via RealtimeKit, otherwise a warning and the normal + // scheduler. Works on stock kernels (unlike ignore_realtime=false, which + // requires PREEMPT_RT). On by default: the 1 kHz torque loop always + // benefits, and the TAM residual's ~0.1-0.3 ms per tick misses deadlines + // on a loaded non-RT machine without it. + int rt_priority = 80; size_t dof = 7; Eigen::Matrix joint_limits = (Eigen::Matrix(2, 7) << @@ -137,6 +155,23 @@ class Franka : public common::Robot { common::ThreadSafeValue> tam_mlp_weight{ std::nullopt}; common::ThreadSafeFixedBuffer tam_history{TAM_HISTORY_SIZE}; + // Parsed TAM MLP (set_tam_mlp_weight parses off the control thread). + common::ThreadSafeValue> tam_model; + // Robot-lifetime monotonic epoch for TAM history timestamps: controller + // restarts (e.g. gain changes) must not restart the encoder's timeline — + // they only leave a short, maskable gap in an otherwise continuous stream. + const std::chrono::steady_clock::time_point tam_epoch = + std::chrono::steady_clock::now(); + double tam_now() const { + return std::chrono::duration(std::chrono::steady_clock::now() - + tam_epoch).count(); + } + // Control-thread-only: ticks since the residual became active (1 s ramp). + int tam_active_ticks = 0; + // Control-thread-only ring of the newest history samples so tam_forward + // reads its MLP window without touching the shared buffer's mutex (the + // control thread is the only writer of both). + std::deque tam_recent; void osc(); void joint_controller(); void zero_torque_controller(); @@ -195,9 +230,10 @@ class Franka : public common::Robot { common::Pose get_base_pose_in_world_coordinates() override; - void set_tam_mlp_weight(const Eigen::VectorXd& weight) { - this->tam_mlp_weight.store(weight); - } + // Parses the packed TAM adaptor binary (one byte per vector element) and + // publishes it to the control thread; throws std::runtime_error on a + // malformed buffer. Parsing happens here, off the 1 kHz thread. + void set_tam_mlp_weight(const Eigen::VectorXd& weight); void set_tam_latent(const Eigen::VectorXd& latent) { this->tam_latent.store(latent); } @@ -206,7 +242,9 @@ class Franka : public common::Robot { return this->tam_history.to_vector(); } - common::Vector7d tam_forward(const std::array& tau); + common::Vector7d tam_forward(const std::array& tau, + const franka::RobotState& robot_state, + const std::array& gravity); void reset() override; void close() override {}; diff --git a/extensions/rcs_fr3/src/hw/simadaptor.h b/extensions/rcs_fr3/src/hw/simadaptor.h new file mode 100644 index 00000000..a4742bcc --- /dev/null +++ b/extensions/rcs_fr3/src/hw/simadaptor.h @@ -0,0 +1,1245 @@ +#pragma once + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace adaptor { + +// Best-effort real-time scheduling for the calling thread (a 1 kHz robot +// control loop). Two mechanisms, tried in order: +// 1. SCHED_FIFO at the given priority — needs an rtprio rlimit +// (one line in /etc/security/limits.conf + a fresh login). +// 2. RealtimeKit (the D-Bus service the desktop audio stack uses), which +// can grant SCHED_RR without any host configuration. rtkit caps the +// priority (typically 20) and requires a finite RLIMIT_RTTIME on the +// process; both are fine here — SCHED_RR at any priority preempts all +// normal threads, and a control thread blocks every millisecond so it +// can never approach the runtime limit. +// If both fail the thread stays on the normal scheduler and a warning is +// printed; expect communication-constraint violations on a loaded machine +// in that state. priority <= 0 disables. +inline void TryElevateControlThreadPriority(int priority) { + if (priority <= 0) { + return; + } + sched_param param{}; + param.sched_priority = priority; + if (pthread_setschedparam(pthread_self(), SCHED_FIFO, ¶m) == 0) { + std::cerr << "[rcs] control thread on SCHED_FIFO priority " << priority + << std::endl; + return; + } + + rlimit rttime{}; + rttime.rlim_cur = 200000; // 200 ms of uninterrupted RT CPU (rtkit requires a finite cap) + rttime.rlim_max = 200000; + setrlimit(RLIMIT_RTTIME, &rttime); + const int rtkit_priority = std::min(priority, 20); + const long pid = static_cast(getpid()); + const long tid = static_cast(syscall(SYS_gettid)); + const std::string cmd = + "busctl --system --timeout=2 call org.freedesktop.RealtimeKit1 " + "/org/freedesktop/RealtimeKit1 org.freedesktop.RealtimeKit1 " + "MakeThreadRealtimeWithPID ttu " + + std::to_string(pid) + " " + std::to_string(tid) + " " + + std::to_string(rtkit_priority) + " >/dev/null 2>&1"; + const int sys_rc = std::system(cmd.c_str()); + (void)sys_rc; + const int policy = sched_getscheduler(0); + if (policy == SCHED_RR || policy == SCHED_FIFO) { + std::cerr << "[rcs] control thread on SCHED_RR priority " << rtkit_priority + << " via RealtimeKit" << std::endl; + return; + } + + std::cerr << "[rcs] real-time scheduling unavailable (SCHED_FIFO denied, " + "RealtimeKit not reachable); control thread stays on the " + "normal scheduler. For the strong SCHED_FIFO path add " + "\" - rtprio 99\" to /etc/security/limits.conf and open " + "a fresh login session." << std::endl; +} + + +using M = Eigen::Matrix; +using V = Eigen::VectorXf; + +// Binary format flags (optional int32 in weights header). +// Header layout variants: +// - v0: int32 dof, emb_dim, hidden, depth +// - v1: v0 + int32 history_steps +// - v2: v1 + int32 flags (bitfield, see below) +// Flags: +// bit0: monotone command map +// bit1: jointwise conditioning +// bit2: external torque auxiliary head +// bit3: jointwise direct-residual no-command-map head +// bit4: command-conditioned head that consumes the current desired torque +static constexpr uint32_t kFlagUseMonotoneMap = 1u << 0; +static constexpr uint32_t kFlagJointwiseConditioning = 1u << 1; +static constexpr uint32_t kFlagPredictExternalTorque = 1u << 2; +static constexpr uint32_t kFlagJointwiseDirectResidualHead = 1u << 3; +static constexpr uint32_t kFlagCommandConditionedHead = 1u << 4; +static constexpr int kJointDirectProjDim = 16; + +inline float gelu_scalar(float x) { + const float k = std::sqrt(2.0f / static_cast(M_PI)); + return 0.5f * x * (1.0f + std::tanh(k * (x + 0.044715f * x * x * x))); +} + +inline void gelu_inplace(M& X) { + X = X.unaryExpr([](float v) { return gelu_scalar(v); }); +} + +inline M softplus(const M& Z) { + const Eigen::ArrayXXf x = Z.array(); + const Eigen::ArrayXXf y = x.max(0.0f) + ((-x.abs()).exp() + 1.0f).log(); + return y.matrix(); +} + +inline M softplus_k(const M& Z, float k) { + const Eigen::ArrayXXf x = (k * Z.array()); + const Eigen::ArrayXXf y = x.max(0.0f) + ((-x.abs()).exp() + 1.0f).log(); + return (y / k).matrix(); +} + +inline M three_regime_map(const M& tau, + const M& xl, + const M& xh, + const M& s_neg, + const M& s_mid, + const M& s_pos, + float k = 30.0f) { + const M tau_anchor = 0.5f * (xl + xh); + const M h1 = softplus_k(tau - xl, k) - softplus_k(tau_anchor - xl, k); + const M h2 = softplus_k(tau - xh, k) - softplus_k(tau_anchor - xh, k); + return (s_neg.array() * (tau - tau_anchor).array() + + (s_mid - s_neg).array() * h1.array() + + (s_pos - s_mid).array() * h2.array()) + .matrix(); +} + +struct Dense { + M W; + V b; + bool use_bias{true}; + + Dense() = default; + Dense(int Fin, int Fout, bool use_bias_ = true) + : W(Fin, Fout), b(Fout), use_bias(use_bias_) {} + + M forward(const M& X) const { + M Y = X * W; + if (use_bias) { + Y.rowwise() += b.transpose(); + } + return Y; + } +}; + +struct LayerNorm { + float eps = 1e-6f; + V gamma; + V beta; + + LayerNorm() = default; + LayerNorm(int dim) : eps(1e-6f), gamma(dim), beta(dim) {} + + M forward(const M& X) const { + const int B = static_cast(X.rows()); + const int D = static_cast(X.cols()); + M Y(B, D); + for (int i = 0; i < B; ++i) { + float mean = X.row(i).mean(); + float var = 0.0f; + for (int j = 0; j < D; ++j) { + const float d = X(i, j) - mean; + var += d * d; + } + var /= static_cast(D); + const float denom = 1.0f / std::sqrt(var + eps); + for (int j = 0; j < D; ++j) { + const float normed = (X(i, j) - mean) * denom; + Y(i, j) = normed * gamma(j) + beta(j); + } + } + return Y; + } +}; + +struct NormStats { + V mean_q; + V mean_qd; + V mean_tau; + V inv_std_q; + V inv_std_qd; + V inv_std_tau; + bool enabled{false}; + + void disable() { + enabled = false; + mean_q.resize(0); + mean_qd.resize(0); + mean_tau.resize(0); + inv_std_q.resize(0); + inv_std_qd.resize(0); + inv_std_tau.resize(0); + } + + void apply(M& q, M& qd, M& tau) const { + if (!enabled) return; + const auto apply_vec = [](M& X, const V& mean, const V& inv_std) { + X.rowwise() -= mean.transpose(); + X = (X.array().rowwise() * inv_std.transpose().array()).matrix(); + }; + apply_vec(q, mean_q, inv_std_q); + apply_vec(qd, mean_qd, inv_std_qd); + apply_vec(tau, mean_tau, inv_std_tau); + } +}; + +struct AdaLN { + int emb_dim; + int cond_dim; + Dense proj; + LayerNorm ln; + + AdaLN() = default; + AdaLN(int emb_dim_, int cond_dim_) + : emb_dim(emb_dim_), + cond_dim(cond_dim_), + proj(cond_dim_, 2 * emb_dim_), + ln(emb_dim_) {} + + M forward(const M& x, const M& cond) const { + M x_norm = ln.forward(x); + M ss = proj.forward(cond); + M gamma = ss.leftCols(emb_dim); + M beta = ss.rightCols(emb_dim); + return (x_norm.array() * (1.0f + gamma.array()) + beta.array()).matrix(); + } +}; + +struct SimAdaptorBlock { + int hidden_dim; + int out_dim; + int cond_dim; + AdaLN adaln1; + Dense fc1; + AdaLN adaln2; + Dense fc2; + + SimAdaptorBlock() = default; + SimAdaptorBlock(int hidden_dim_, int out_dim_, int cond_dim_) + : hidden_dim(hidden_dim_), + out_dim(out_dim_), + cond_dim(cond_dim_), + adaln1(hidden_dim_, cond_dim_), + fc1(hidden_dim_, hidden_dim_), + adaln2(hidden_dim_, cond_dim_), + fc2(hidden_dim_, out_dim_) {} + + M forward(const M& x, const M& cond) const { + M h = adaln1.forward(x, cond); + h = fc1.forward(h); + gelu_inplace(h); + h = adaln2.forward(h, cond); + h = fc2.forward(h); + return h; + } +}; + +struct ForwardResult { + M delta_tau; + bool has_external_torque{false}; + M external_torque; +}; + +struct SimAdaptor { + int dof; + int emb_dim; // per-joint embedding width when jointwise conditioning is enabled + int hidden; + int depth; + int history_steps; + bool use_monotone_map{false}; + bool use_jointwise_conditioning{false}; + bool predict_external_torque{false}; + bool use_jointwise_direct_head{false}; + bool use_command_conditioned_head{false}; + uint32_t flags{0}; + + Dense q_stem; + Dense qd_stem; + Dense tau_stem; + LayerNorm ln_q; + LayerNorm ln_qd; + LayerNorm ln_tau; + std::vector blocks; + SimAdaptorBlock out_block; + SimAdaptorBlock external_torque_block; + std::vector joint_global_projs; + Dense direct_tau_des_proj; + Dense joint_direct_tau_hyper; + Dense joint_direct_projected; + Dense joint_direct_projected_out; + NormStats norm_stats; + + SimAdaptor(int dof_, + int emb_dim_, + int hidden_, + int depth_, + int history_steps_ = 1, + bool use_monotone_map_ = false, + bool use_jointwise_conditioning_ = false, + bool predict_external_torque_ = false, + bool use_jointwise_direct_head_ = false, + bool use_command_conditioned_head_ = false) + : dof(dof_), + emb_dim(emb_dim_), + hidden(hidden_), + depth(depth_), + history_steps(history_steps_), + use_monotone_map(use_monotone_map_), + use_jointwise_conditioning(use_jointwise_conditioning_), + predict_external_torque(predict_external_torque_), + use_jointwise_direct_head(use_jointwise_direct_head_), + use_command_conditioned_head(use_command_conditioned_head_), + flags((use_monotone_map_ ? kFlagUseMonotoneMap : 0u) | + (use_jointwise_conditioning_ ? kFlagJointwiseConditioning : 0u) | + (predict_external_torque_ ? kFlagPredictExternalTorque : 0u) | + (use_jointwise_direct_head_ ? kFlagJointwiseDirectResidualHead : 0u) | + (use_command_conditioned_head_ ? kFlagCommandConditionedHead : 0u)), + q_stem(use_jointwise_conditioning_ ? history_steps_ : (dof_ * history_steps_), + hidden_, + false), + qd_stem(use_jointwise_conditioning_ ? history_steps_ : (dof_ * history_steps_), + hidden_, + false), + tau_stem(use_jointwise_conditioning_ ? history_steps_ : (dof_ * history_steps_), + hidden_, + false), + ln_q(hidden_), + ln_qd(hidden_), + ln_tau(hidden_), + out_block(hidden_, + use_jointwise_conditioning_ + ? (use_monotone_map_ ? 6 : 1) + : (use_monotone_map_ ? (6 * dof_) : dof_), + emb_dim_), + external_torque_block(hidden_, + use_jointwise_conditioning_ ? 1 : dof_, + emb_dim_), + direct_tau_des_proj(hidden_ + dof_, hidden_, true), + joint_direct_tau_hyper(hidden_, 2 * kJointDirectProjDim, true), + joint_direct_projected(use_command_conditioned_head_ ? (2 * kJointDirectProjDim) + : kJointDirectProjDim, + kJointDirectProjDim, + true), + joint_direct_projected_out( + kJointDirectProjDim, + (use_jointwise_conditioning_ && use_monotone_map_ && + use_command_conditioned_head_) + ? 6 + : 1, + true) { + const int middle = std::max(0, depth_ - 1); + blocks.reserve(middle); + if (use_jointwise_conditioning_) { + joint_global_projs.reserve(middle); + } + for (int i = 0; i < middle; ++i) { + blocks.emplace_back(hidden_, hidden_, emb_dim_); + if (use_jointwise_conditioning_) { + joint_global_projs.emplace_back(hidden_, hidden_, true); + } + } + } + + int input_cols() const { return dof * history_steps; } + + int expected_history_embedding_cols() const { + return use_jointwise_conditioning ? (dof * emb_dim) : emb_dim; + } + + + // A single sample of a streaming controller history (torque in the + // ideal-model, gravity-included space). + struct StreamRow { + std::array q{}; + std::array dq{}; + std::array tau_model{}; + }; + + // Residual for the newest sample of a 1 kHz stream. + // ``rows`` must hold exactly ``history_steps`` samples, oldest first; + // ``latent`` is the history embedding. Returns false (with ``delta`` zero) + // when the inputs do not match the model or the forward fails. The result + // is clipped per joint to ``|clip|``; ramping is the caller's policy. + bool forward_stream(const std::vector& rows, + const Eigen::VectorXd& latent, + const Eigen::Matrix& clip, + Eigen::Matrix& delta) const { + delta.setZero(); + const int T = history_steps; + const int D = dof; + if (D > 7 || static_cast(rows.size()) != T || + latent.size() != expected_history_embedding_cols()) { + return false; + } + M emb(1, latent.size()); + for (Eigen::Index i = 0; i < latent.size(); ++i) { + emb(0, i) = static_cast(latent(i)); + } + M q_hist(1, T * D); + M dq_hist(1, T * D); + M tau_hist(1, T * D); + for (int t = 0; t < T; ++t) { + const StreamRow& r = rows[static_cast(t)]; + const int offset = t * D; + for (int j = 0; j < D; ++j) { + q_hist(0, offset + j) = static_cast(r.q[j]); + dq_hist(0, offset + j) = static_cast(r.dq[j]); + tau_hist(0, offset + j) = static_cast(r.tau_model[j]); + } + } + M out; + try { + out = forward(q_hist, dq_hist, tau_hist, emb); + } catch (...) { + return false; + } + if (out.size() != D) { + return false; + } + for (int j = 0; j < D; ++j) { + const double v = static_cast(out(0, j)); + const double lim = std::abs(clip(j)); + delta(j) = std::isfinite(v) ? std::clamp(v, -lim, lim) : 0.0; + } + return true; + } + + M forward(const M& q, const M& qd, const M& tau, const M& history_emb) const { + return forward_with_aux(q, qd, tau, history_emb).delta_tau; + } + + ForwardResult forward_with_aux(const M& q, + const M& qd, + const M& tau, + const M& history_emb) const { + ForwardResult out; + out.delta_tau = M::Zero(q.rows(), dof); + const int expected_cols = input_cols(); + if (q.cols() != expected_cols || qd.cols() != expected_cols || tau.cols() != expected_cols) { + std::cerr << "Input dim mismatch: expected " << expected_cols << " cols but got " + << q.cols() << "/" << qd.cols() << "/" << tau.cols() << ".\n"; + return out; + } + if (q.rows() != qd.rows() || q.rows() != tau.rows()) { + std::cerr << "Batch dim mismatch across q/qd/tau: " << q.rows() << "/" << qd.rows() << "/" + << tau.rows() << ".\n"; + return out; + } + if (history_emb.rows() != q.rows()) { + std::cerr << "History embedding batch mismatch: expected " << q.rows() << " rows but got " + << history_emb.rows() << ".\n"; + return out; + } + const int expected_emb_cols = expected_history_embedding_cols(); + if (history_emb.cols() != expected_emb_cols) { + std::cerr << "History embedding dim mismatch: expected " << expected_emb_cols + << " cols but got " << history_emb.cols() << ".\n"; + return out; + } + + if (use_jointwise_conditioning) { + return forward_jointwise(q, qd, tau, history_emb); + } + return forward_global(q, qd, tau, history_emb); + } + + static inline bool read_array(std::istream& in, float* dst, size_t count) { + in.read(reinterpret_cast(dst), static_cast(count * sizeof(float))); + return in.good(); + } + + static bool load_dense(std::istream& in, Dense& d) { + if (!read_array(in, d.W.data(), static_cast(d.W.size()))) return false; + if (d.use_bias) { + if (!read_array(in, d.b.data(), static_cast(d.b.size()))) return false; + } else { + d.b.setZero(); + } + return true; + } + + static bool load_layernorm(std::istream& in, LayerNorm& ln) { + if (!read_array(in, ln.gamma.data(), static_cast(ln.gamma.size()))) return false; + if (!read_array(in, ln.beta.data(), static_cast(ln.beta.size()))) return false; + return true; + } + + static bool load_block(std::istream& in, SimAdaptorBlock& blk) { + if (!load_layernorm(in, blk.adaln1.ln)) return false; + if (!load_dense(in, blk.adaln1.proj)) return false; + if (!load_dense(in, blk.fc1)) return false; + if (!load_layernorm(in, blk.adaln2.ln)) return false; + if (!load_dense(in, blk.adaln2.proj)) return false; + if (!load_dense(in, blk.fc2)) return false; + return true; + } + + static std::unique_ptr LoadFromFile(const std::string& path) { + std::ifstream in(path, std::ios::binary); + if (!in.is_open()) return nullptr; + + in.seekg(0, std::ios::end); + const std::streampos file_size_pos = in.tellg(); + if (file_size_pos <= 0) return nullptr; + const size_t file_size = static_cast(file_size_pos); + in.seekg(0, std::ios::beg); + return LoadFromStream(in, file_size); + } + + // Same binary layout as LoadFromFile, from an in-memory buffer. + static std::unique_ptr LoadFromMemory(const char* data, + size_t size) { + if (data == nullptr || size == 0) return nullptr; + std::istringstream in(std::string(data, size), std::ios::binary); + return LoadFromStream(in, size); + } + + static std::unique_ptr LoadFromStream(std::istream& in, + const size_t file_size) { + + int32_t dof_f = 0; + int32_t emb_f = 0; + int32_t hidden_f = 0; + int32_t depth_f = 0; + in.read(reinterpret_cast(&dof_f), sizeof(int32_t)); + in.read(reinterpret_cast(&emb_f), sizeof(int32_t)); + in.read(reinterpret_cast(&hidden_f), sizeof(int32_t)); + in.read(reinterpret_cast(&depth_f), sizeof(int32_t)); + if (!in.good()) return nullptr; + + auto try_match_size = [&](int header_ints, + int32_t history, + uint32_t flags_candidate, + size_t& stats_floats_out) -> bool { + if (history <= 0) return false; + const size_t header_bytes = static_cast(header_ints) * sizeof(int32_t); + if (file_size < header_bytes) return false; + const size_t payload_bytes = file_size - header_bytes; + if (payload_bytes % sizeof(float) != 0) return false; + const size_t payload_floats = payload_bytes / sizeof(float); + const bool mono = (flags_candidate & kFlagUseMonotoneMap) != 0u; + const bool jointwise = (flags_candidate & kFlagJointwiseConditioning) != 0u; + const bool predict_external = + (flags_candidate & kFlagPredictExternalTorque) != 0u; + const bool jointwise_direct = + (flags_candidate & kFlagJointwiseDirectResidualHead) != 0u; + const bool command_conditioned = + (flags_candidate & kFlagCommandConditionedHead) != 0u; + const uint32_t known_flags = + kFlagUseMonotoneMap | kFlagJointwiseConditioning | + kFlagPredictExternalTorque | kFlagJointwiseDirectResidualHead | + kFlagCommandConditionedHead; + if ((flags_candidate & ~known_flags) != 0u) return false; + if (jointwise_direct && (!jointwise || mono)) return false; + if (command_conditioned && (!jointwise && mono)) return false; + const size_t weight_floats = + ModelParamCount( + dof_f, + history, + emb_f, + hidden_f, + depth_f, + mono, + jointwise, + jointwise_direct, + command_conditioned, + predict_external); + if (payload_floats < weight_floats) return false; + const size_t stats_floats = payload_floats - weight_floats; + const size_t per_dof = 6ull * static_cast(dof_f); + const size_t per_hist = 6ull * static_cast(dof_f) * static_cast(history); + if (stats_floats == 0 || stats_floats == per_dof || stats_floats == per_hist) { + stats_floats_out = stats_floats; + return true; + } + return false; + }; + + int32_t history_f = 1; + uint32_t flags_f = 0; + bool use_mono = false; + bool use_jointwise = false; + bool predict_external = false; + bool use_jointwise_direct = false; + bool use_command_conditioned = false; + int header_ints = 4; + size_t stats_floats = 0; + + const std::streampos after_header4 = in.tellg(); + int32_t history_candidate = 1; + in.read(reinterpret_cast(&history_candidate), sizeof(int32_t)); + if (in.good()) { + const std::streampos after_header5 = in.tellg(); + + int32_t flags_candidate_i32 = 0; + in.read(reinterpret_cast(&flags_candidate_i32), sizeof(int32_t)); + if (in.good()) { + const uint32_t flags_candidate = static_cast(flags_candidate_i32); + size_t stats_candidate = 0; + if (try_match_size(6, history_candidate, flags_candidate, stats_candidate)) { + history_f = history_candidate; + flags_f = flags_candidate; + use_mono = (flags_f & kFlagUseMonotoneMap) != 0u; + use_jointwise = (flags_f & kFlagJointwiseConditioning) != 0u; + predict_external = + (flags_f & kFlagPredictExternalTorque) != 0u; + use_jointwise_direct = + (flags_f & kFlagJointwiseDirectResidualHead) != 0u; + use_command_conditioned = + (flags_f & kFlagCommandConditionedHead) != 0u; + header_ints = 6; + stats_floats = stats_candidate; + } else { + in.clear(); + in.seekg(after_header5); + } + } else { + in.clear(); + in.seekg(after_header5); + } + + if (header_ints == 4) { + size_t stats_std = 0; + size_t stats_mono = 0; + const bool ok_std = try_match_size(5, history_candidate, 0u, stats_std); + const bool ok_mono = try_match_size(5, history_candidate, kFlagUseMonotoneMap, stats_mono); + if (ok_std ^ ok_mono) { + history_f = history_candidate; + use_mono = ok_mono; + use_jointwise = false; + predict_external = false; + use_jointwise_direct = false; + use_command_conditioned = false; + flags_f = use_mono ? kFlagUseMonotoneMap : 0u; + header_ints = 5; + stats_floats = ok_mono ? stats_mono : stats_std; + } else { + in.clear(); + in.seekg(after_header4); + } + } + } else { + in.clear(); + in.seekg(after_header4); + } + + if (header_ints == 4) { + const size_t header_bytes = 4ull * sizeof(int32_t); + if (file_size < header_bytes) return nullptr; + const size_t payload_bytes = file_size - header_bytes; + if (payload_bytes % sizeof(float) != 0) return nullptr; + const size_t payload_floats = payload_bytes / sizeof(float); + + int best_history = -1; + bool best_mono = false; + size_t best_stats = 0; + + const size_t per_step = + 3ull * static_cast(dof_f) * static_cast(hidden_f); + int max_history = + per_step > 0 ? static_cast(payload_floats / per_step + 1) : 1; + max_history = std::min(max_history, 1024); + + for (int h = 1; h <= max_history; ++h) { + for (const bool mono_candidate : {false, true}) { + size_t stats_candidate = 0; + if (try_match_size(4, + h, + mono_candidate ? kFlagUseMonotoneMap : 0u, + stats_candidate)) { + if (best_history != -1) { + return nullptr; + } + best_history = h; + best_mono = mono_candidate; + best_stats = stats_candidate; + } + } + } + if (best_history <= 0) return nullptr; + history_f = best_history; + use_mono = best_mono; + use_jointwise = false; + predict_external = false; + use_jointwise_direct = false; + use_command_conditioned = false; + flags_f = use_mono ? kFlagUseMonotoneMap : 0u; + header_ints = 4; + stats_floats = best_stats; + } + + in.clear(); + in.seekg(static_cast(header_ints * sizeof(int32_t)), std::ios::beg); + + auto model = std::make_unique( + dof_f, + emb_f, + hidden_f, + depth_f, + history_f, + use_mono, + use_jointwise, + predict_external, + use_jointwise_direct, + use_command_conditioned); + model->flags = flags_f; + + if (!load_dense(in, model->q_stem)) return nullptr; + if (!load_dense(in, model->qd_stem)) return nullptr; + if (!load_dense(in, model->tau_stem)) return nullptr; + if (!load_layernorm(in, model->ln_q)) return nullptr; + if (!load_layernorm(in, model->ln_qd)) return nullptr; + if (!load_layernorm(in, model->ln_tau)) return nullptr; + for (size_t i = 0; i < model->blocks.size(); ++i) { + if (!load_block(in, model->blocks[i])) return nullptr; + if (model->use_jointwise_conditioning) { + if (!load_dense(in, model->joint_global_projs[i])) return nullptr; + } + } + if (model->use_jointwise_conditioning && model->use_command_conditioned_head) { + if (!load_dense(in, model->joint_direct_tau_hyper)) return nullptr; + if (!load_dense(in, model->joint_direct_projected)) return nullptr; + if (!load_dense(in, model->joint_direct_projected_out)) return nullptr; + } else if (model->use_jointwise_direct_head) { + if (!load_dense(in, model->joint_direct_tau_hyper)) return nullptr; + if (!load_dense(in, model->joint_direct_projected)) return nullptr; + if (!load_dense(in, model->joint_direct_projected_out)) return nullptr; + } else { + if (!model->use_jointwise_conditioning && model->use_command_conditioned_head) { + if (!load_dense(in, model->direct_tau_des_proj)) return nullptr; + } + if (!load_block(in, model->out_block)) return nullptr; + } + if (model->predict_external_torque) { + if (!load_block(in, model->external_torque_block)) return nullptr; + } + model->norm_stats.disable(); + + if (stats_floats > 0) { + const size_t per_dof = 6ull * static_cast(model->dof); + const size_t per_hist = + 6ull * static_cast(model->dof) * static_cast(model->history_steps); + if (stats_floats == per_dof || stats_floats == per_hist) { + if (!load_norm_stats_stream( + in, stats_floats, model->dof, model->history_steps, model->norm_stats)) { + std::cerr << "[error] Failed to load norm_stats; continuing without normalization.\n"; + model->norm_stats.disable(); + } else { + std::cout << "Loaded norm_stats (" << (stats_floats == per_dof ? "per-DoF" : "per-step") + << ").\n"; + } + } else { + std::cerr << "[error] Unrecognized trailing data (" << stats_floats + << " floats). Expected " << per_dof << " (per-DoF) or " << per_hist + << " (per-step). Ignoring.\n"; + } + } else { + std::cout << "No norm_stats found; continuing without normalization.\n"; + } + + return model; + } + + private: + ForwardResult forward_global(const M& q, + const M& qd, + const M& tau, + const M& history_emb) const { + ForwardResult result; + M q_in = q; + M qd_in = qd; + M tau_in = tau; + norm_stats.apply(q_in, qd_in, tau_in); + + M tau_des; + if (use_monotone_map || use_command_conditioned_head) { + tau_des = tau_in.rightCols(dof); + tau_in.rightCols(dof).setZero(); + } + + M h_q = ln_q.forward(q_stem.forward(q_in)); + M h_qd = ln_qd.forward(qd_stem.forward(qd_in)); + M h_tau = ln_tau.forward(tau_stem.forward(tau_in)); + M h = h_q + h_qd + h_tau; + for (const auto& blk : blocks) { + h = blk.forward(h, history_emb); + gelu_inplace(h); + } + + if (!use_monotone_map) { + if (use_command_conditioned_head) { + M h_tau(h.rows(), hidden + dof); + h_tau.leftCols(hidden) = h; + h_tau.rightCols(dof) = tau_des; + h_tau = direct_tau_des_proj.forward(h_tau); + gelu_inplace(h_tau); + result.delta_tau = out_block.forward(h_tau, history_emb); + } else { + result.delta_tau = out_block.forward(h, history_emb); + } + } else { + M params = out_block.forward(h, history_emb); + if (params.cols() != 6 * dof) { + std::cerr << "Monotone head dim mismatch: expected " << (6 * dof) << " cols but got " + << params.cols() << ".\n"; + result.delta_tau = M::Zero(q.rows(), dof); + } else { + const auto slice = [&](int idx) { return params.middleCols(idx * dof, dof); }; + const M s_neg = slice(0); + const M s_mid_raw = slice(1); + const M s_pos = slice(2); + const M xl_raw = slice(3); + const M dx_raw = slice(4); + const M bias = slice(5); + + const M s_mid = (softplus(s_mid_raw).array() + 1e-4f).matrix(); + const M xl = xl_raw; + const M xh = (xl_raw.array() + softplus(dx_raw).array() + 1e-3f).matrix(); + + result.delta_tau = + three_regime_map(tau_des, xl, xh, s_neg, s_mid, s_pos) + bias; + } + } + + if (predict_external_torque) { + M external = external_torque_block.forward(h, history_emb); + if (external.cols() != dof) { + std::cerr << "External torque head dim mismatch: expected " << dof + << " cols but got " << external.cols() << ".\n"; + } else { + result.has_external_torque = true; + result.external_torque = std::move(external); + } + } + return result; + } + + ForwardResult forward_jointwise(const M& q, + const M& qd, + const M& tau, + const M& history_emb) const { + ForwardResult result; + const int batch = static_cast(q.rows()); + M q_in = q; + M qd_in = qd; + M tau_in = tau; + norm_stats.apply(q_in, qd_in, tau_in); + + M tau_des; + if (use_monotone_map || use_jointwise_direct_head || use_command_conditioned_head) { + const int last_offset = (history_steps - 1) * dof; + tau_des = tau_in.middleCols(last_offset, dof); + tau_in.middleCols(last_offset, dof).setZero(); + } + + const M q_joint = reshape_flat_history_by_joint(q_in, batch); + const M qd_joint = reshape_flat_history_by_joint(qd_in, batch); + const M tau_joint = reshape_flat_history_by_joint(tau_in, batch); + const M cond_joint = reshape_flat_embedding_by_joint(history_emb, batch); + + M h_q = ln_q.forward(q_stem.forward(q_joint)); + M h_qd = ln_qd.forward(qd_stem.forward(qd_joint)); + M h_tau = ln_tau.forward(tau_stem.forward(tau_joint)); + M h = h_q + h_qd + h_tau; + + for (size_t i = 0; i < blocks.size(); ++i) { + h = blocks[i].forward(h, cond_joint); + gelu_inplace(h); + const M h_global = mean_over_joint_rows(h, batch); + const M projected = joint_global_projs[i].forward(h_global); + add_joint_residual_inplace(h, projected, batch); + } + + if (use_command_conditioned_head) { + M linear_bias = joint_direct_tau_hyper.forward(h); + if (linear_bias.cols() != 2 * kJointDirectProjDim) { + std::cerr << "Jointwise command tau hyper dim mismatch: expected " + << (2 * kJointDirectProjDim) << " cols but got " + << linear_bias.cols() << ".\n"; + result.delta_tau = M::Zero(batch, dof); + } else { + const M linear = linear_bias.leftCols(kJointDirectProjDim); + const M bias = linear_bias.rightCols(kJointDirectProjDim); + M tau_projected = project_tau_basis(linear, bias, tau_des, batch); + gelu_inplace(tau_projected); + const M other_joint_tau_context = other_joint_context(tau_projected, batch); + const M head_input = + concat_columns(use_monotone_map ? bias : tau_projected, + other_joint_tau_context); + M features = joint_direct_projected.forward(head_input); + gelu_inplace(features); + const M out = joint_direct_projected_out.forward(features); + if (use_monotone_map) { + if (out.cols() != 6) { + std::cerr << "Jointwise command monotone head dim mismatch: expected 6 cols but got " + << out.cols() << ".\n"; + result.delta_tau = M::Zero(batch, dof); + } else { + const M s_neg = collect_joint_column(out, batch, 0); + const M s_mid_raw = collect_joint_column(out, batch, 1); + const M s_pos = collect_joint_column(out, batch, 2); + const M xl_raw = collect_joint_column(out, batch, 3); + const M dx_raw = collect_joint_column(out, batch, 4); + const M map_bias = collect_joint_column(out, batch, 5); + + const M s_mid = (softplus(s_mid_raw).array() + 1e-4f).matrix(); + const M xl = xl_raw; + const M xh = (xl_raw.array() + softplus(dx_raw).array() + 1e-3f).matrix(); + + result.delta_tau = + three_regime_map(tau_des, xl, xh, s_neg, s_mid, s_pos) + map_bias; + } + } else { + if (out.cols() != 1) { + std::cerr << "Jointwise command direct output dim mismatch: expected 1 col but got " + << out.cols() << ".\n"; + result.delta_tau = M::Zero(batch, dof); + } else { + result.delta_tau = collapse_joint_scalar_output(out, batch); + } + } + } + } else if (use_jointwise_direct_head) { + M linear_bias = joint_direct_tau_hyper.forward(h); + if (linear_bias.cols() != 2 * kJointDirectProjDim) { + std::cerr << "Jointwise direct tau hyper dim mismatch: expected " + << (2 * kJointDirectProjDim) << " cols but got " + << linear_bias.cols() << ".\n"; + result.delta_tau = M::Zero(batch, dof); + } else { + const M linear = linear_bias.leftCols(kJointDirectProjDim); + const M bias = linear_bias.rightCols(kJointDirectProjDim); + M projected = project_tau_basis(linear, bias, tau_des, batch); + gelu_inplace(projected); + projected = joint_direct_projected.forward(projected); + gelu_inplace(projected); + const M out = joint_direct_projected_out.forward(projected); + if (out.cols() != 1) { + std::cerr << "Jointwise direct output dim mismatch: expected 1 col but got " + << out.cols() << ".\n"; + result.delta_tau = M::Zero(batch, dof); + } else { + result.delta_tau = collapse_joint_scalar_output(out, batch); + } + } + } else if (!use_monotone_map) { + const M out = out_block.forward(h, cond_joint); + if (out.cols() != 1) { + std::cerr << "Jointwise head dim mismatch: expected 1 col but got " << out.cols() + << ".\n"; + result.delta_tau = M::Zero(batch, dof); + } else { + result.delta_tau = collapse_joint_scalar_output(out, batch); + } + } else { + const M out = out_block.forward(h, cond_joint); + if (out.cols() != 6) { + std::cerr << "Jointwise monotone head dim mismatch: expected 6 cols but got " + << out.cols() << ".\n"; + result.delta_tau = M::Zero(batch, dof); + } else { + const M s_neg = collect_joint_column(out, batch, 0); + const M s_mid_raw = collect_joint_column(out, batch, 1); + const M s_pos = collect_joint_column(out, batch, 2); + const M xl_raw = collect_joint_column(out, batch, 3); + const M dx_raw = collect_joint_column(out, batch, 4); + const M bias = collect_joint_column(out, batch, 5); + + const M s_mid = (softplus(s_mid_raw).array() + 1e-4f).matrix(); + const M xl = xl_raw; + const M xh = (xl_raw.array() + softplus(dx_raw).array() + 1e-3f).matrix(); + + result.delta_tau = + three_regime_map(tau_des, xl, xh, s_neg, s_mid, s_pos) + bias; + } + } + + if (predict_external_torque) { + const M external = external_torque_block.forward(h, cond_joint); + if (external.cols() != 1) { + std::cerr << "Jointwise external torque head dim mismatch: expected 1 col but got " + << external.cols() << ".\n"; + } else { + result.has_external_torque = true; + result.external_torque = collapse_joint_scalar_output(external, batch); + } + } + return result; + } + + M reshape_flat_history_by_joint(const M& flat, int batch) const { + M out(batch * dof, history_steps); + for (int b = 0; b < batch; ++b) { + for (int t = 0; t < history_steps; ++t) { + const int src_offset = t * dof; + for (int j = 0; j < dof; ++j) { + out(b * dof + j, t) = flat(b, src_offset + j); + } + } + } + return out; + } + + M reshape_flat_embedding_by_joint(const M& flat, int batch) const { + M out(batch * dof, emb_dim); + for (int b = 0; b < batch; ++b) { + for (int j = 0; j < dof; ++j) { + const int src_offset = j * emb_dim; + for (int e = 0; e < emb_dim; ++e) { + out(b * dof + j, e) = flat(b, src_offset + e); + } + } + } + return out; + } + + M mean_over_joint_rows(const M& joint_rows, int batch) const { + M out(batch, joint_rows.cols()); + out.setZero(); + for (int b = 0; b < batch; ++b) { + for (int j = 0; j < dof; ++j) { + out.row(b) += joint_rows.row(b * dof + j); + } + out.row(b) /= static_cast(dof); + } + return out; + } + + void add_joint_residual_inplace(M& joint_rows, const M& residual, int batch) const { + for (int b = 0; b < batch; ++b) { + for (int j = 0; j < dof; ++j) { + joint_rows.row(b * dof + j) += residual.row(b); + } + } + } + + M project_tau_basis(const M& scale, const M& bias, const M& tau_des, int batch) const { + M out(scale.rows(), kJointDirectProjDim); + for (int b = 0; b < batch; ++b) { + for (int j = 0; j < dof; ++j) { + const int row = b * dof + j; + const float tau_value = tau_des(b, j); + for (int p = 0; p < kJointDirectProjDim; ++p) { + out(row, p) = scale(row, p) * tau_value + bias(row, p); + } + } + } + return out; + } + + M other_joint_context(const M& joint_rows, int batch) const { + M out(joint_rows.rows(), joint_rows.cols()); + for (int b = 0; b < batch; ++b) { + Eigen::RowVectorXf sum = Eigen::RowVectorXf::Zero(joint_rows.cols()); + for (int j = 0; j < dof; ++j) { + sum += joint_rows.row(b * dof + j); + } + for (int j = 0; j < dof; ++j) { + const int row = b * dof + j; + out.row(row) = sum - joint_rows.row(row); + } + } + return out; + } + + static M concat_columns(const M& left, const M& right) { + M out(left.rows(), left.cols() + right.cols()); + out.leftCols(left.cols()) = left; + out.rightCols(right.cols()) = right; + return out; + } + + M collapse_joint_scalar_output(const M& joint_rows, int batch) const { + M out(batch, dof); + for (int b = 0; b < batch; ++b) { + for (int j = 0; j < dof; ++j) { + out(b, j) = joint_rows(b * dof + j, 0); + } + } + return out; + } + + M collect_joint_column(const M& joint_rows, int batch, int col) const { + M out(batch, dof); + for (int b = 0; b < batch; ++b) { + for (int j = 0; j < dof; ++j) { + out(b, j) = joint_rows(b * dof + j, col); + } + } + return out; + } + + static size_t DenseParamCount(int in_dim, int out_dim, bool use_bias) { + return static_cast(in_dim) * static_cast(out_dim) + + (use_bias ? static_cast(out_dim) : 0ull); + } + + static size_t BlockParamCount(int hidden_dim, int out_dim, int cond_dim) { + const size_t h = static_cast(hidden_dim); + const size_t od = static_cast(out_dim); + const size_t cond = static_cast(cond_dim); + const size_t adaln_proj = (2 * h * cond) + (2 * h); + const size_t ln = 2 * h; + const size_t fc1 = (h * h) + h; + const size_t fc2 = (h * od) + od; + return ln + adaln_proj + fc1 + ln + adaln_proj + fc2; + } + + static size_t ModelParamCount(int dof, + int history_steps, + int emb_dim, + int hidden, + int depth, + bool use_monotone_map, + bool use_jointwise_conditioning, + bool use_jointwise_direct_head, + bool use_command_conditioned_head, + bool predict_external_torque) { + const size_t h = static_cast(hidden); + const size_t d = static_cast(dof); + const size_t emb = static_cast(emb_dim); + const size_t hist = static_cast(history_steps); + const size_t stems = + use_jointwise_conditioning ? (3ull * hist * h) : (3ull * d * hist * h); + const size_t ln_stems = 6ull * h; + const size_t middle_blocks = + static_cast(std::max(0, depth - 1)) * BlockParamCount(hidden, hidden, emb); + const size_t joint_global = + use_jointwise_conditioning + ? static_cast(std::max(0, depth - 1)) * DenseParamCount(hidden, hidden, true) + : 0ull; + const size_t out_dim = + use_jointwise_conditioning ? (use_monotone_map ? 6ull : 1ull) + : (use_monotone_map ? (6ull * d) : d); + const bool jointwise_command_head = + use_jointwise_conditioning && use_command_conditioned_head; + const bool global_command_head = + (!use_jointwise_conditioning) && (!use_monotone_map) && + use_command_conditioned_head; + const size_t direct_head_params = + jointwise_command_head + ? DenseParamCount(hidden, 2 * kJointDirectProjDim, true) + + DenseParamCount(2 * kJointDirectProjDim, kJointDirectProjDim, true) + + DenseParamCount(kJointDirectProjDim, + use_monotone_map ? 6 : 1, + true) + : (use_jointwise_direct_head + ? DenseParamCount(hidden, 2 * kJointDirectProjDim, true) + + DenseParamCount(kJointDirectProjDim, + kJointDirectProjDim, + true) + + DenseParamCount(kJointDirectProjDim, 1, true) + : 0ull); + const size_t out_block_params = + (use_jointwise_direct_head || jointwise_command_head) + ? 0ull + : BlockParamCount(hidden, static_cast(out_dim), emb); + const size_t global_command_params = + global_command_head ? DenseParamCount(hidden + dof, hidden, true) : 0ull; + const size_t external_out_dim = use_jointwise_conditioning ? 1ull : d; + const size_t external_block_params = + predict_external_torque + ? BlockParamCount(hidden, static_cast(external_out_dim), emb) + : 0ull; + return stems + ln_stems + middle_blocks + joint_global + global_command_params + + out_block_params + + direct_head_params + + external_block_params; + } + + static bool load_norm_stats_from_buffer(const std::vector& buf, + int dof, + int history_steps, + NormStats& stats) { + const int input_dim = dof * history_steps; + const size_t per_dof = 6ull * static_cast(dof); + const size_t per_step = 6ull * static_cast(input_dim); + if (buf.size() != per_dof && buf.size() != per_step) return false; + + const float eps = 1e-6f; + V var_q_full; + V var_qd_full; + V var_tau_full; + + const auto fill_full = [&](size_t idx, V& dst) { + dst.resize(input_dim); + if (buf.size() == per_step) { + const float* p = buf.data() + static_cast(idx * input_dim); + for (int i = 0; i < input_dim; ++i) dst(i) = p[i]; + return; + } + const float* p = buf.data() + static_cast(idx * dof); + for (int h = 0; h < history_steps; ++h) { + for (int j = 0; j < dof; ++j) { + dst(h * dof + j) = p[j]; + } + } + }; + + fill_full(0, stats.mean_q); + fill_full(1, stats.mean_qd); + fill_full(2, stats.mean_tau); + fill_full(3, var_q_full); + fill_full(4, var_qd_full); + fill_full(5, var_tau_full); + + stats.inv_std_q.resize(input_dim); + stats.inv_std_qd.resize(input_dim); + stats.inv_std_tau.resize(input_dim); + for (int i = 0; i < input_dim; ++i) { + stats.inv_std_q(i) = 1.0f / std::sqrt(var_q_full(i) + eps); + stats.inv_std_qd(i) = 1.0f / std::sqrt(var_qd_full(i) + eps); + stats.inv_std_tau(i) = 1.0f / std::sqrt(var_tau_full(i) + eps); + } + stats.enabled = true; + return true; + } + + static bool load_norm_stats_stream(std::istream& in, + size_t float_count, + int dof, + int history_steps, + NormStats& stats) { + std::vector buf(float_count); + in.read(reinterpret_cast(buf.data()), + static_cast(float_count * sizeof(float))); + if (!in.good()) return false; + return load_norm_stats_from_buffer(buf, dof, history_steps, stats); + } +}; + +} // namespace adaptor