diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..d645695673349e3947e8e5ae42332d0ac3164cd7 --- /dev/null +++ b/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..05a1e504cd811155aeb2ee06d296f3b13ce289b2 --- /dev/null +++ b/README.md @@ -0,0 +1,100 @@ +--- +license: apache-2.0 +library_name: actioncodec2 +tags: +- robotics +- action-tokenization +- transformers +- extended-edition +--- + +# ActionCodec2 Extended Edition — 1st Order + +This tokenizer was trained on a curated corpus with lower-quality datasets and most stationary-arm segments removed. The data curation concerns tokenizer training, independently of a downstream policy's training data. + +This is a complete tokenizer artifact for continuous robot actions, rather than +a neural policy checkpoint. Keep all files together when copying it. + +- Encoding order: 1 +- Codec rate: 15 Hz +- Codebook size: 2048 per profile +- Fitted profiles: eef +- Gripper representation: binary + +- Trajectory tracking precision tier: 1mm + +The [ActionCodec2 collection](https://huggingface.co/collections/ZibinDong/actioncodec2-6ab6c4860ee90a881dda0a8b) +contains the paper and extended editions. Edition and encoding order are separate choices. + +## Install and enable C++ acceleration + +```bash +hf download ZibinDong/ActionCodec2-Extended-1mm-2k --local-dir ./actioncodec2-tokenizer +pip install -r ./actioncodec2-tokenizer/requirements.txt +pip install ./actioncodec2-tokenizer +``` + +The last command compiles the bundled C++17 CPU kernels. It requires a C++17 +compiler and installs `actioncodec2-kernels`, independently of the G0.5 package. +The processor automatically uses installed kernels; no GPU is needed for +tokenization. An installed copy of the +[source repository](https://github.com/ZibinDong/actioncodec2) also supplies them. +Loading without kernels uses the bundled Python implementation. Set +`ACTIONCODEC2_NO_NATIVE=1` to explicitly select that reference path. + +Acceleration covers physical quantization, Set-BPE encode/decode, trajectory +resampling and pose operations. Equal-horizon batches use C++ worker threads, +which release the Python GIL. Thread counts adapt to measured chunk cost and +the process CPU allowance. Forked DataLoader workers use one thread each by +default; `ACTIONCODEC2_NUM_THREADS` sets an explicit per-process budget. + +## Encode and decode + +```python +import numpy as np +from transformers import AutoProcessor + +codec = AutoProcessor.from_pretrained( + "ZibinDong/ActionCodec2-Extended-1mm-2k", trust_remote_code=True, +) +# Pin revision to a Hub commit for reproducible downloads. +actions = np.zeros((2, 12, 7), dtype=np.float32) +actions[..., 6] = 1.0 + +tokens = codec.encode(actions, fps=15) +reconstructed = codec.decode(tokens, fps=15) +``` + +Inputs are `(T, D)` or `(B, T, D)`, with a common horizon within a batch. +Tokens are ragged integer rows; decode returns CPU `torch.float32` tensors with +shape `(B, T_out, D)`. Pass the actual recording rate as `fps`; resampling can +change the decoded horizon. Batch chunks together for encoding throughput. + +The default layout is `single_eef_delta`: position increments in metres, +body-frame rotation-vector increments in radians, then gripper (open at >= 0.8). +Rotation follows `R_next = R_previous @ Exp(rotvec)`. Convert normalized +controller commands to these units before encoding. Inspect other layouts with +`codec.print_action_spaces()` and select with `codec.for_action_space(name)`. +Absolute layouts require the same `current_state` of shape `(B, D)` during +encoding and decoding. Second-order profiles additionally require explicit +`previous_first_order` boundaries at the codec rate. + +`codec.grammar(horizon, fps=...)` constrains generated token coverage. Remove +model prompt, EOS and padding IDs before decoding. Each profile occupies its +own namespace; `router_config.yaml` records token offsets. A policy must have +enough action vocabulary rows for its selected profile or routed namespace. +The source repository's G0.5 adapter currently accepts at most 4096 per profile. + +## Validation scope + +This configuration has not been evaluated in the selected extended-edition LIBERO comparison. + +The extended edition currently contains first-order profiles. At 1 mm, EEF +starts at 2k and joint profiles at 4k because of the primitive vocabulary size; +5 mm and 10 mm provide 1k, 2k, 4k, 8k and 16k. Only six first-order configurations +have been tested on LIBERO. The full range has not been validated downstream. +Precision tiers describe tokenizer trajectory tracking, not robot execution +accuracy. Selected benchmark results do not establish a general advantage for +the extended edition or predict performance on other tasks or real robots. + +See [configuration and measured results](https://github.com/ZibinDong/actioncodec2). diff --git a/actioncodec2_kernels/__init__.py b/actioncodec2_kernels/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0102cb43af492d4ad6d149af86f29b56fb8bc5e8 --- /dev/null +++ b/actioncodec2_kernels/__init__.py @@ -0,0 +1 @@ +"""Portable ActionCodec2 CPU kernels.""" diff --git a/actioncodec2_kernels/module.cpp b/actioncodec2_kernels/module.cpp new file mode 100644 index 0000000000000000000000000000000000000000..d00217eb968834cd7da73320ac62c62d699679bb --- /dev/null +++ b/actioncodec2_kernels/module.cpp @@ -0,0 +1,90 @@ +// pybind11 entry point for the ActionCodec2 physical-stage kernels. +// +// These are the encode/decode kernels the runtime bundled with an ActionCodec2 +// Hugging Face artifact calls when a native module is available. Each one has +// a bit-identical NumPy implementation in that runtime. + +#include +#include + +#include +#include + +#include "parallel.hpp" + +namespace py = pybind11; + +py::array_t so3_second_order_encode_batch( + const py::array_t&, + const py::array_t&, + const py::array_t&, + const py::array_t&, + int32_t, bool, double, double, bool, int); + +py::array_t additive_second_order_encode_batch( + const py::array_t&, + const py::array_t&, + const py::array_t&, + const py::array_t&, + const py::array_t&, bool, int); + +py::array_t physical_decode_batch( + const py::array_t&, const py::array_t&, const py::array_t&, + const py::array_t&, const py::array_t&, + const py::array_t&, + const py::array_t&, + const py::array_t&, int32_t, int); + +py::array_t physical_encode_batch( + const py::array_t&, const py::array_t&, const py::array_t&, + const py::array_t&, const py::array_t&, + const py::array_t&, + const py::array_t&, + const py::array_t&, int32_t, int, + const py::array_t&, bool); + +py::array_t so3_encode_batch( + const py::array_t&, + const py::array_t&, + const py::array_t&, int32_t, bool, bool, + double, bool, int); + +void register_setbpe(py::module_& module); +void register_resample(py::module_& module); + +PYBIND11_MODULE(_actioncodec2_native, module) { + module.doc() = "ActionCodec2 encode/decode kernels: physical stage, resampling, Set-BPE"; + register_setbpe(module); + register_resample(module); + module.def( + "set_thread_budget", + [](int threads) { + if (threads < 0) throw std::invalid_argument("thread budget must be non-negative"); + ac2::thread_budget().store(threads); + }, + py::arg("threads"), + "Cap threads per call for this process; 0 uses every CPU it may run on."); + module.def("thread_budget", [] { return ac2::thread_budget().load(); }); + module.def("plan_threads", &ac2::plan_threads, py::arg("count"), py::arg("requested"), + py::arg("seconds_each"), "Threads a call would use for ``count`` items."); + module.def("so3_second_order_encode_batch", &so3_second_order_encode_batch, + py::arg("table"), py::arg("bins"), py::arg("quaternions"), + py::arg("previous_first_order"), py::arg("input_order"), py::arg("left"), + py::arg("limit"), py::arg("radius"), py::arg("smooth"), py::arg("threads") = 0); + module.def("additive_second_order_encode_batch", &additive_second_order_encode_batch, + py::arg("table"), py::arg("bins"), py::arg("positions"), + py::arg("previous_first_order"), py::arg("epsilon"), py::arg("smooth"), + py::arg("threads") = 0); + module.def("physical_decode_batch", &physical_decode_batch, py::arg("table"), + py::arg("bins"), py::arg("binary"), py::arg("lower"), py::arg("upper"), + py::arg("cells"), py::arg("initial"), py::arg("initial_velocity"), + py::arg("order"), py::arg("threads") = 0); + module.def("physical_encode_batch", &physical_encode_batch, py::arg("table"), py::arg("bins"), + py::arg("binary"), py::arg("lower"), py::arg("upper"), py::arg("positions"), + py::arg("initial"), py::arg("initial_velocity"), py::arg("order"), + py::arg("threads") = 0, py::arg("epsilon") = py::array_t(), + py::arg("smooth") = false); + module.def("so3_encode_batch", &so3_encode_batch, py::arg("table"), py::arg("bins"), + py::arg("actions"), py::arg("input_order"), py::arg("rotation_vector"), + py::arg("left"), py::arg("limit"), py::arg("smooth"), py::arg("threads") = 0); +} diff --git a/actioncodec2_kernels/parallel.hpp b/actioncodec2_kernels/parallel.hpp new file mode 100644 index 0000000000000000000000000000000000000000..f848b7bebc04646ee80e7a91034bad2812593fab --- /dev/null +++ b/actioncodec2_kernels/parallel.hpp @@ -0,0 +1,116 @@ +// Fork-safe data parallelism for the runtime encode/decode kernels. +// +// These kernels run inside forked processes (for example PyTorch DataLoader +// workers). GNU OpenMP keeps a process-wide thread pool that does not survive +// fork(): once the parent has entered a parallel region -- ours, or any other +// library's sharing libgomp -- the child hangs or crashes when it enters one. +// This helper instead starts its threads per call and joins them before +// returning, so no thread state outlives a call and fork() is always safe. +// +// Iterations are split into contiguous static blocks, like OpenMP's +// ``schedule(static)``, and every iteration writes only its own outputs, so +// results do not depend on the thread count. +// +// Starting a thread costs tens of microseconds, and one iteration costs from +// well under a microsecond to about a hundred depending on the kernel and the +// profile's primitive counts. So each call times its first iteration on the +// calling thread and sizes the split of the remaining work ``W`` from it: the +// calling thread starts threads one after another at ``s`` seconds each, so a +// split over ``w`` threads takes about ``W / w + s * w``, least at +// ``w = sqrt(W / s)``. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef __linux__ +#include +#endif + +namespace ac2 { + +// Process-wide cap on threads per call; 0 means every CPU the process may use. +inline std::atomic& thread_budget() { + static std::atomic value{0}; + return value; +} + +inline int available_cpus() { +#ifdef __linux__ + cpu_set_t allowed; + CPU_ZERO(&allowed); + if (sched_getaffinity(0, sizeof(allowed), &allowed) == 0) { + return std::max(1, CPU_COUNT(&allowed)); + } +#endif + return std::max(1, static_cast(std::thread::hardware_concurrency())); +} + +// Measured cost of starting and joining one thread (seconds). +constexpr double kThreadStartSeconds = 30e-6; + +// Threads for ``count`` iterations of ``seconds_each``: at most ``requested`` +// when positive, at most the process budget, and at most ``sqrt(W / s)``. +inline int64_t plan_threads(int64_t count, int requested, double seconds_each) { + const int budget = thread_budget().load(); + int64_t threads = budget > 0 ? budget : available_cpus(); + if (requested > 0) threads = std::min(threads, requested); + const double by_work = std::sqrt(count * seconds_each / kThreadStartSeconds); + if (by_work < static_cast(threads)) threads = static_cast(by_work); + return std::max(1, std::min(threads, count)); +} + +// Call ``body(index)`` for every index in [0, count). Iteration 0 runs first +// on the calling thread and sizes the split of the rest; the calling thread +// also runs the first block of the rest. The first exception thrown by any +// block is rethrown after all threads have joined. +template +void parallel_for(int64_t count, int threads, Body&& body) { + if (count <= 0) return; + const auto start = std::chrono::steady_clock::now(); + body(0); + const double seconds_each = + std::chrono::duration(std::chrono::steady_clock::now() - start).count(); + const int64_t rest = count - 1; + if (rest == 0) return; + const int64_t workers = plan_threads(rest, threads, seconds_each); + auto run = [&](int64_t worker) { + const int64_t begin = 1 + rest * worker / workers; + const int64_t end = 1 + rest * (worker + 1) / workers; + for (int64_t index = begin; index < end; ++index) body(index); + }; + if (workers == 1) { + run(0); + return; + } + std::exception_ptr failure; + std::mutex failure_lock; + auto guarded = [&](int64_t worker) { + try { + run(worker); + } catch (...) { + std::lock_guard hold(failure_lock); + if (!failure) failure = std::current_exception(); + } + }; + std::vector pool; + pool.reserve(static_cast(workers - 1)); + try { + for (int64_t worker = 1; worker < workers; ++worker) pool.emplace_back(guarded, worker); + } catch (...) { + for (auto& thread : pool) thread.join(); + throw; + } + guarded(0); + for (auto& thread : pool) thread.join(); + if (failure) std::rethrow_exception(failure); +} + +} // namespace ac2 diff --git a/actioncodec2_kernels/physical.cpp b/actioncodec2_kernels/physical.cpp new file mode 100644 index 0000000000000000000000000000000000000000..8fe5735380aa25a1bbb064566730a6ec951fe87d --- /dev/null +++ b/actioncodec2_kernels/physical.cpp @@ -0,0 +1,407 @@ +// Native physical-stage recursions. +// +// The encoders are strictly sequential in time -- each step's request depends +// on the reconstruction error the previous step left behind -- so the win here +// is not vectorization but removing per-step interpreter and NumPy dispatch +// overhead. That overhead dominates exactly where it hurts most: caching a +// dataset as fixed-horizon chunks calls the encoder millions of times with +// T ~ 20, where a NumPy implementation spends nearly all of its time in call +// setup. The batch entry points additionally spread chunks over threads. + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "parallel.hpp" + +namespace py = pybind11; + +namespace { + +struct SecondOrderAxis { + const float* values; + int32_t count; + bool ordered; + int32_t nonnegative; + int32_t nonpositive; +}; + +template +inline int32_t pick_second_order_axis(const SecondOrderAxis& axis, double request, + double direction, double epsilon) { + const float* values = axis.values; + if (axis.ordered) { + const int32_t insertion = std::lower_bound(values, values + axis.count, request) - values; + const int32_t left = std::max(0, insertion - 1); + const int32_t right = std::min(insertion, axis.count - 1); + int32_t nearest = std::fabs(request - values[left]) <= std::fabs(request - values[right]) + ? left : right; + // Rounded distance plateaus must keep the original first-cell tie-break. + while (nearest > 0 && std::fabs(request - values[nearest - 1]) == std::fabs(request - values[nearest])) + --nearest; + if constexpr (!Smooth) return nearest; + const int32_t preferred = direction > 0 ? std::max(nearest, axis.nonnegative) + : direction < 0 ? std::min(nearest, axis.nonpositive) : nearest; + const double primitive = values[preferred]; + const double sign = primitive > 0 ? 1 : primitive < 0 ? -1 : 0; + return sign * direction >= 0 && std::fabs(request - primitive) <= epsilon ? preferred : nearest; + } + // Small tables are faster to scan; arbitrary stored order also uses this path. + int32_t nearest = 0, feasible = -1, preferred = -1; + double nearest_distance = INFINITY, feasible_distance = INFINITY, preferred_distance = INFINITY; + for (int32_t k = 0; k < axis.count; ++k) { + const double primitive = values[k]; + const double distance = std::fabs(request - primitive); + if (distance < nearest_distance) { nearest = k; nearest_distance = distance; } + if constexpr (!Smooth) continue; + if (distance > epsilon) continue; + if (distance < feasible_distance) { feasible = k; feasible_distance = distance; } + const double sign = primitive > 0 ? 1 : primitive < 0 ? -1 : 0; + if (sign * direction >= 0 && distance < preferred_distance) { + preferred = k; + preferred_distance = distance; + } + } + return preferred >= 0 ? preferred : feasible >= 0 ? feasible : nearest; +} + +} // namespace + +py::array_t additive_second_order_encode_batch( + const py::array_t& table, + const py::array_t& bins, + const py::array_t& positions, + const py::array_t& previous_first_order, + const py::array_t& epsilon, + bool smooth, int threads) { + if (positions.ndim() != 3 || table.ndim() != 2 || previous_first_order.ndim() != 2) + throw std::invalid_argument("invalid second-order additive batch shapes"); + const int64_t batch = positions.shape(0), steps = positions.shape(1), dim = positions.shape(2); + if (table.shape(0) != dim || bins.size() != dim || epsilon.size() != dim || + previous_first_order.shape(0) != batch || previous_first_order.shape(1) != dim) + throw std::invalid_argument("inconsistent second-order additive batch dimensions"); + for (int64_t j = 0; j < dim; ++j) { + if (bins.data()[j] < 1 || bins.data()[j] > table.shape(1)) + throw std::invalid_argument("primitive count exceeds table width"); + } + py::array_t output({batch, steps, dim}); + const auto* targets = positions.data(); + const auto* initial = previous_first_order.data(); + const auto* axes = table.data(); + const auto* counts = bins.data(); + const auto* limits = epsilon.data(); + auto* cells = output.mutable_data(); + const int64_t width = table.shape(1); + std::vector search; + for (int64_t j = 0; j < dim; ++j) { + const float* begin = axes + j * width; + const float* end = begin + counts[j]; + const bool ordered = counts[j] > 16 && + std::adjacent_find(begin, end, [](float a, float b) { return a >= b; }) == end; + const int32_t nonnegative = ordered ? std::lower_bound(begin, end, 0.0f) - begin : 0; + const int32_t nonpositive = ordered ? std::upper_bound(begin, end, 0.0f) - begin - 1 : 0; + search.push_back({begin, counts[j], ordered, std::min(nonnegative, counts[j] - 1), + std::max(nonpositive, 0)}); + } + py::gil_scoped_release release; + ac2::parallel_for(batch, threads, [&](int64_t i) { + std::vector position(dim, 0), velocity(initial+i*dim, initial+(i+1)*dim), direction(dim, 0); + for (int64_t t = 0; t < steps; ++t) { + for (int64_t j = 0; j < dim; ++j) { + const double request = targets[(i*steps+t)*dim+j] - (position[j]+velocity[j]); + const int32_t choice = smooth + ? pick_second_order_axis(search[j], request, direction[j], limits[j]) + : pick_second_order_axis(search[j], request, direction[j], limits[j]); + const double primitive = axes[j*width+choice]; + velocity[j] += primitive; + position[j] += velocity[j]; + if (primitive != 0) direction[j] = primitive > 0 ? 1 : -1; + cells[(i*steps+t)*dim+j] = choice; + } + } + }); + return output; +} + +namespace { + +struct Layout { + int32_t dimension; + int32_t width; // padded primitive-table stride + const float* table; // [dimension * width], +inf padding + const int32_t* bins; // [dimension] + const uint8_t* binary; // [dimension], 1 for absolute two-valued axes + const float* lower; + const float* upper; + const float* epsilon; // [dimension]; null unless `smooth` is set + bool smooth; // pick by FirstOrderCodec._select_smooth +}; + +// Matches PhysicalCodecBase.SOURCE_TOLERANCE. +constexpr float kSourceTolerance = 8.0F * FLT_EPSILON; + +inline int32_t nearest(const Layout& layout, int32_t axis, float request) { + const float* values = layout.table + static_cast(axis) * layout.width; + const int32_t count = layout.bins[axis]; + int32_t best = 0; + float best_distance = std::fabs(values[0] - request); + for (int32_t index = 1; index < count; ++index) { + const float distance = std::fabs(values[index] - request); + if (distance < best_distance) { + best_distance = distance; + best = index; + } + } + return best; +} + +// FirstOrderCodec._select_smooth, primitive for primitive. Lexicographic in +// (does not reverse `direction`, distance to the request), restricted to the +// primitives already within epsilon; nothing feasible falls back to `nearest`. +inline int32_t select_smooth(const Layout& layout, int32_t axis, float request, float direction) { + const float* values = layout.table + static_cast(axis) * layout.width; + const int32_t count = layout.bins[axis]; + const float radius = layout.epsilon[axis] + kSourceTolerance; + int32_t nearest_index = 0; + float nearest_distance = std::fabs(values[0] - request); + bool any_feasible = false; + int32_t feasible_index = 0; + float feasible_distance = 0.0F; + bool any_preferred = false; + int32_t preferred_index = 0; + float preferred_distance = 0.0F; + for (int32_t index = 0; index < count; ++index) { + const float distance = std::fabs(values[index] - request); + if (index > 0 && distance < nearest_distance) { + nearest_distance = distance; + nearest_index = index; + } + if (!(distance <= radius)) continue; + if (!any_feasible || distance < feasible_distance) { + any_feasible = true; + feasible_distance = distance; + feasible_index = index; + } + // np.sign(value) * direction >= 0: a zero primitive never reverses. + const float sign = values[index] > 0.0F ? 1.0F : (values[index] < 0.0F ? -1.0F : 0.0F); + if (!(sign * direction >= 0.0F)) continue; + if (!any_preferred || distance < preferred_distance) { + any_preferred = true; + preferred_distance = distance; + preferred_index = index; + } + } + if (any_preferred) return preferred_index; + if (any_feasible) return feasible_index; + return nearest_index; +} + +void encode_first_order(const Layout& layout, const float* positions, int64_t steps, + const float* initial, int32_t* cells) { + const int32_t dimension = layout.dimension; + std::vector state(initial, initial + dimension); + // Sign of the last non-zero primitive each coordinate emitted. + std::vector direction(dimension, 0.0F); + for (int64_t step = 0; step < steps; ++step) { + const float* target = positions + step * dimension; + int32_t* row = cells + step * dimension; + for (int32_t axis = 0; axis < dimension; ++axis) { + const float* values = layout.table + static_cast(axis) * layout.width; + if (layout.binary[axis]) { + const int32_t cell = nearest(layout, axis, target[axis]); + state[axis] = values[cell]; + row[axis] = cell; + continue; + } + const float request = target[axis] - state[axis]; + const int32_t cell = layout.smooth ? select_smooth(layout, axis, request, direction[axis]) + : nearest(layout, axis, request); + const float picked = values[cell]; + state[axis] = std::min(std::max(state[axis] + picked, layout.lower[axis]), + layout.upper[axis]); + if (picked != 0.0F) direction[axis] = picked > 0.0F ? 1.0F : -1.0F; + row[axis] = cell; + } + } +} + +void decode_first_order(const Layout& layout, const int32_t* cells, int64_t steps, + const float* initial, float* positions) { + const int32_t dimension = layout.dimension; + std::vector state(initial, initial + dimension); + for (int64_t step = 0; step < steps; ++step) { + const int32_t* row = cells + step * dimension; + float* out = positions + step * dimension; + for (int32_t axis = 0; axis < dimension; ++axis) { + const float* values = layout.table + static_cast(axis) * layout.width; + const float primitive = values[row[axis]]; + if (layout.binary[axis]) { + state[axis] = primitive; + } else { + state[axis] = std::min(std::max(state[axis] + primitive, layout.lower[axis]), + layout.upper[axis]); + } + out[axis] = state[axis]; + } + } +} + +void encode_second_order(const Layout& layout, const float* positions, int64_t steps, + const float* initial, const float* initial_velocity, int32_t* cells) { + const int32_t dimension = layout.dimension; + std::vector source(initial, initial + dimension); + std::vector source_velocity(initial_velocity, initial_velocity + dimension); + std::vector recon(initial, initial + dimension); + std::vector recon_velocity(initial_velocity, initial_velocity + dimension); + std::vector error(dimension, 0.0F); + std::vector previous(dimension, 0.0F); + for (int64_t step = 0; step < steps; ++step) { + const float* target = positions + step * dimension; + int32_t* row = cells + step * dimension; + for (int32_t axis = 0; axis < dimension; ++axis) { + const float* values = layout.table + static_cast(axis) * layout.width; + if (layout.binary[axis]) { + row[axis] = nearest(layout, axis, target[axis]); + continue; + } + const float velocity = target[axis] - source[axis]; + const float request = + 2.0F * error[axis] - previous[axis] + (velocity - source_velocity[axis]); + const int32_t cell = nearest(layout, axis, request); + recon_velocity[axis] += values[cell]; + recon[axis] += recon_velocity[axis]; + source[axis] = target[axis]; + source_velocity[axis] = velocity; + previous[axis] = error[axis]; + error[axis] = source[axis] - recon[axis]; + row[axis] = cell; + } + } +} + +void decode_second_order(const Layout& layout, const int32_t* cells, int64_t steps, + const float* initial, const float* initial_velocity, float* positions) { + const int32_t dimension = layout.dimension; + std::vector state(initial, initial + dimension); + std::vector velocity(initial_velocity, initial_velocity + dimension); + for (int64_t step = 0; step < steps; ++step) { + const int32_t* row = cells + step * dimension; + float* out = positions + step * dimension; + for (int32_t axis = 0; axis < dimension; ++axis) { + const float primitive = layout.table[static_cast(axis) * layout.width + row[axis]]; + if (layout.binary[axis]) { + out[axis] = primitive; + continue; + } + velocity[axis] += primitive; + state[axis] += velocity[axis]; + out[axis] = state[axis]; + } + } +} + +Layout make_layout(const py::array_t& table, const py::array_t& bins, + const py::array_t& binary, const py::array_t& lower, + const py::array_t& upper, const py::array_t& epsilon, + bool smooth) { + Layout layout; + layout.dimension = static_cast(table.shape(0)); + layout.width = static_cast(table.shape(1)); + layout.table = table.data(); + layout.bins = bins.data(); + layout.binary = binary.data(); + layout.lower = lower.data(); + layout.upper = upper.data(); + layout.epsilon = epsilon.size() == 0 ? nullptr : epsilon.data(); + layout.smooth = smooth && layout.epsilon != nullptr; + if (smooth && layout.epsilon == nullptr) { + throw std::invalid_argument("smooth selection needs a per-coordinate epsilon"); + } + if (layout.epsilon != nullptr && epsilon.size() != layout.dimension) { + throw std::invalid_argument("epsilon must have one entry per coordinate"); + } + return layout; +} + +using Array2D = py::array_t; +using Array3D = py::array_t; +using Cells3D = py::array_t; + +} // namespace + +py::array_t physical_decode_batch(const py::array_t& table, + const py::array_t& bins, + const py::array_t& binary, + const py::array_t& lower, + const py::array_t& upper, const Cells3D& cells, + const Array2D& initial, + const Array2D& initial_velocity, int32_t order, + int threads) { + if (cells.ndim() != 3 || cells.shape(2) != table.shape(0) || initial.ndim() != 2 || + initial_velocity.ndim() != 2 || initial.shape(0) != cells.shape(0) || + initial_velocity.shape(0) != cells.shape(0) || initial.shape(1) != cells.shape(2) || + initial_velocity.shape(1) != cells.shape(2) || (order != 1 && order != 2)) { + throw std::invalid_argument("invalid physical decode batch shapes or order"); + } + const Layout layout = + make_layout(table, bins, binary, lower, upper, py::array_t(), false); + const int64_t batch = cells.shape(0); + const int64_t steps = cells.shape(1); + const int64_t dimension = layout.dimension; + py::array_t positions({batch, steps, dimension}); + { + py::gil_scoped_release release; + ac2::parallel_for(batch, threads, [&](int64_t index) { + const int32_t* source = cells.data() + index * steps * dimension; + float* target = positions.mutable_data() + index * steps * dimension; + if (order == 1) { + decode_first_order(layout, source, steps, initial.data() + index * dimension, target); + } else { + decode_second_order(layout, source, steps, initial.data() + index * dimension, + initial_velocity.data() + index * dimension, target); + } + }); + } + return positions; +} + +// One thread per chunk. ``positions`` is [B, T, D] and ``initial`` is [B, D]. +py::array_t physical_encode_batch(const py::array_t& table, + const py::array_t& bins, + const py::array_t& binary, + const py::array_t& lower, + const py::array_t& upper, + const Array3D& positions, const Array2D& initial, + const Array2D& initial_velocity, int32_t order, + int threads, const py::array_t& epsilon, + bool smooth) { + const Layout layout = make_layout(table, bins, binary, lower, upper, epsilon, smooth); + const int64_t batch = positions.shape(0); + const int64_t steps = positions.shape(1); + const int64_t dimension = layout.dimension; + py::array_t cells({batch, steps, dimension}); + const float* source = positions.data(); + const float* starts = initial.data(); + const float* velocities = initial_velocity.data(); + int32_t* output = cells.mutable_data(); + { + py::gil_scoped_release release; + ac2::parallel_for(batch, threads, [&](int64_t index) { + const float* chunk = source + index * steps * dimension; + int32_t* target = output + index * steps * dimension; + if (order == 1) { + encode_first_order(layout, chunk, steps, starts + index * dimension, target); + } else { + encode_second_order(layout, chunk, steps, starts + index * dimension, + velocities + index * dimension, target); + } + }); + } + return cells; +} diff --git a/actioncodec2_kernels/resample.cpp b/actioncodec2_kernels/resample.cpp new file mode 100644 index 0000000000000000000000000000000000000000..750388b58ebf477f2462ea4ad99e06b2ea5005d5 --- /dev/null +++ b/actioncodec2_kernels/resample.cpp @@ -0,0 +1,833 @@ +// Rate conversion and pose algebra for canonical ActionCodec2 trajectories. +// +// A port of the runtime's ``physical/semantic.py`` resampling and pose helpers, +// including the scipy pieces they call (natural CubicSpline, PchipInterpolator, +// PPoly evaluation, RotationSpline). Operations follow the NumPy/scipy order +// of evaluation, and the LAPACK solvers are the reference algorithms, so +// results agree with the Python runtime to the last bits of float64; the +// float32 outputs therefore match except where a value sits on a rounding +// boundary. Batch entry points run rows in parallel without the GIL. + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "parallel.hpp" + +namespace py = pybind11; + +namespace { + +using Vec = std::vector; + +struct Quat { + double x, y, z, w; +}; + +// ---- runtime ``_so3`` quaternion math (xyzw), NumPy evaluation order ---- + +Quat so3_normalize(const Quat& q) { + const double norm = std::sqrt(q.x * q.x + q.y * q.y + q.z * q.z + q.w * q.w); + if (norm < 1e-12) throw std::invalid_argument("a pose contains a zero quaternion"); + return {q.x / norm, q.y / norm, q.z / norm, q.w / norm}; +} + +Quat so3_multiply(const Quat& a, const Quat& b) { + const double cx = a.y * b.z - a.z * b.y; + const double cy = a.z * b.x - a.x * b.z; + const double cz = a.x * b.y - a.y * b.x; + return so3_normalize({a.w * b.x + b.w * a.x + cx, a.w * b.y + b.w * a.y + cy, + a.w * b.z + b.w * a.z + cz, + a.w * b.w - (a.x * b.x + a.y * b.y + a.z * b.z)}); +} + +Quat so3_from_rotvec(double x, double y, double z) { + const double angle = std::sqrt(x * x + y * y + z * z); + const double scale = angle >= 1e-12 ? std::sin(0.5 * angle) / angle : 0.5; + return so3_normalize({x * scale, y * scale, z * scale, std::cos(0.5 * angle)}); +} + +void so3_to_rotvec(const Quat& raw, double* out) { + Quat q = so3_normalize(raw); + if (q.w < 0.0) q = {-q.x, -q.y, -q.z, -q.w}; + const double vector_norm = std::sqrt(q.x * q.x + q.y * q.y + q.z * q.z); + const double angle = 2.0 * std::atan2(vector_norm, std::min(std::max(q.w, 0.0), 1.0)); + const double scale = vector_norm >= 1e-12 ? angle / vector_norm : 2.0; + out[0] = q.x * scale; + out[1] = q.y * scale; + out[2] = q.z * scale; +} + +// ---- scipy ``Rotation`` operations used by RotationSpline ---- + +Quat sp_from_quat(const Quat& q) { + const double norm = std::sqrt(q.x * q.x + q.y * q.y + q.z * q.z + q.w * q.w); + if (!(norm > 0.0)) throw std::invalid_argument("found zero norm quaternions in `quat`"); + return {q.x / norm, q.y / norm, q.z / norm, q.w / norm}; +} + +Quat sp_inv(const Quat& q) { return {-q.x, -q.y, -q.z, q.w}; } + +Quat sp_compose(const Quat& p, const Quat& q) { + const double cx = p.y * q.z - p.z * q.y; + const double cy = p.z * q.x - p.x * q.z; + const double cz = p.x * q.y - p.y * q.x; + const Quat r{p.w * q.x + q.w * p.x + cx, p.w * q.y + q.w * p.y + cy, p.w * q.z + q.w * p.z + cz, + p.w * q.w - (p.x * q.x + p.y * q.y + p.z * q.z)}; + return sp_from_quat(r); +} + +void sp_as_rotvec(Quat q, double* out) { + if (q.w < 0.0) q = {-q.x, -q.y, -q.z, -q.w}; + const double angle = 2.0 * std::atan2(std::sqrt(q.x * q.x + q.y * q.y + q.z * q.z), q.w); + double scale; + if (angle <= 1e-3) { + const double angle2 = angle * angle; + scale = 2.0 + angle2 / 12.0 + 7.0 * angle2 * angle2 / 2880.0; + } else { + scale = angle / std::sin(angle / 2.0); + } + out[0] = scale * q.x; + out[1] = scale * q.y; + out[2] = scale * q.z; +} + +Quat sp_from_rotvec(const double* v) { + const double angle = std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]); + double scale; + if (angle <= 1e-3) { + const double angle2 = angle * angle; + scale = 0.5 - angle2 / 48.0 + angle2 * angle2 / 3840.0; + } else { + scale = std::sin(angle / 2.0) / angle; + } + return {scale * v[0], scale * v[1], scale * v[2], std::cos(angle / 2.0)}; +} + +// ---- piecewise cubics (scipy PPoly with CubicHermiteSpline coefficients) ---- + +// PPoly interval: x[i] <= t < x[i+1], extrapolating with the end intervals. +int64_t ppoly_interval(const Vec& x, double t) { + const int64_t n = static_cast(x.size()); + if (t >= x[n - 1]) return n - 2; + if (t < x[0]) return 0; + return static_cast(std::upper_bound(x.begin(), x.end(), t) - x.begin()) - 1; +} + +// scipy ``evaluate_poly1`` for a cubic: sum of c[k-1-p] * s**p, powers built by +// repeated multiplication. +double ppoly_value(const double* c, double s) { + double result = 0.0, z = 1.0; + result = result + c[3] * z; + z *= s; + result = result + c[2] * z; + z *= s; + result = result + c[1] * z; + z *= s; + result = result + c[0] * z; + return result; +} + +// Hermite coefficients per interval and column: c[(k*4 + p)*cols + a] with +// p = 0..3 the descending powers, as CubicHermiteSpline stacks them. +Vec hermite_coefficients(const Vec& x, const Vec& y, const Vec& slopes, int64_t cols) { + const int64_t n = static_cast(x.size()); + Vec c(static_cast((n - 1) * 4 * cols)); + for (int64_t k = 0; k < n - 1; ++k) { + const double dx = x[k + 1] - x[k]; + for (int64_t a = 0; a < cols; ++a) { + const double y0 = y[k * cols + a], y1 = y[(k + 1) * cols + a]; + const double d0 = slopes[k * cols + a], d1 = slopes[(k + 1) * cols + a]; + const double slope = (y1 - y0) / dx; + const double t = (d0 + d1 - 2.0 * slope) / dx; + c[(k * 4 + 0) * cols + a] = t / dx; + c[(k * 4 + 1) * cols + a] = (slope - d0) / dx - t; + c[(k * 4 + 2) * cols + a] = d0; + c[(k * 4 + 3) * cols + a] = y0; + } + } + return c; +} + +void ppoly_evaluate(const Vec& x, const Vec& c, int64_t cols, const Vec& times, int64_t out_cols, + int64_t col_offset, double* out) { + for (size_t i = 0; i < times.size(); ++i) { + const int64_t k = ppoly_interval(x, times[i]); + const double s = times[i] - x[k]; + for (int64_t a = 0; a < cols; ++a) { + const double coeff[4] = {c[(k * 4 + 0) * cols + a], c[(k * 4 + 1) * cols + a], + c[(k * 4 + 2) * cols + a], c[(k * 4 + 3) * cols + a]}; + out[i * out_cols + col_offset + a] = ppoly_value(coeff, s); + } + } +} + +// LAPACK dgtsv (reference algorithm) for several right-hand sides b[n][cols]. +void gtsv(Vec dl, Vec d, Vec du, Vec& b, int64_t cols) { + const int64_t n = static_cast(d.size()); + auto pivot = [&](int64_t i, bool last) { + if (std::fabs(d[i]) >= std::fabs(dl[i])) { + if (d[i] == 0.0) throw std::invalid_argument("singular matrix"); + const double fact = dl[i] / d[i]; + d[i + 1] = d[i + 1] - fact * du[i]; + for (int64_t j = 0; j < cols; ++j) b[(i + 1) * cols + j] = b[(i + 1) * cols + j] - fact * b[i * cols + j]; + if (!last) dl[i] = 0.0; + } else { + const double fact = d[i] / dl[i]; + d[i] = dl[i]; + const double temp = d[i + 1]; + d[i + 1] = du[i] - fact * temp; + if (!last) { + dl[i] = du[i + 1]; + du[i + 1] = -fact * dl[i]; + } + du[i] = temp; + for (int64_t j = 0; j < cols; ++j) { + const double swap = b[i * cols + j]; + b[i * cols + j] = b[(i + 1) * cols + j]; + b[(i + 1) * cols + j] = swap - fact * b[(i + 1) * cols + j]; + } + } + }; + for (int64_t i = 0; i + 2 < n; ++i) pivot(i, false); + if (n > 1) pivot(n - 2, true); + if (d[n - 1] == 0.0) throw std::invalid_argument("singular matrix"); + for (int64_t j = 0; j < cols; ++j) { + b[(n - 1) * cols + j] = b[(n - 1) * cols + j] / d[n - 1]; + if (n > 1) + b[(n - 2) * cols + j] = (b[(n - 2) * cols + j] - du[n - 2] * b[(n - 1) * cols + j]) / d[n - 2]; + for (int64_t i = n - 3; i >= 0; --i) { + b[i * cols + j] = + (b[i * cols + j] - du[i] * b[(i + 1) * cols + j] - dl[i] * b[(i + 2) * cols + j]) / d[i]; + } + } +} + +// scipy CubicSpline(bc_type="natural") slopes at the knots. +Vec natural_cubic_slopes(const Vec& x, const Vec& y, int64_t cols) { + const int64_t n = static_cast(x.size()); + Vec dx(n - 1), slope((n - 1) * cols); + for (int64_t k = 0; k < n - 1; ++k) { + dx[k] = x[k + 1] - x[k]; + for (int64_t a = 0; a < cols; ++a) slope[k * cols + a] = (y[(k + 1) * cols + a] - y[k * cols + a]) / dx[k]; + } + // Banded A (3 x n): A[0] superdiagonal, A[1] diagonal, A[2] subdiagonal. + Vec upper(n, 0.0), diagonal(n, 0.0), lower(n, 0.0), b(n * cols); + for (int64_t k = 1; k < n - 1; ++k) diagonal[k] = 2.0 * (dx[k - 1] + dx[k]); + for (int64_t k = 2; k < n; ++k) upper[k] = dx[k - 2]; + for (int64_t k = 0; k + 2 < n; ++k) lower[k] = dx[k + 1]; + for (int64_t k = 1; k < n - 1; ++k) { + for (int64_t a = 0; a < cols; ++a) + b[k * cols + a] = 3.0 * (dx[k] * slope[(k - 1) * cols + a] + dx[k - 1] * slope[k * cols + a]); + } + // Natural ends: second derivative zero (bc value 0.0). + diagonal[0] = 2.0 * dx[0]; + upper[1] = dx[0]; + diagonal[n - 1] = 2.0 * dx[n - 2]; + lower[n - 2] = dx[n - 2]; + for (int64_t a = 0; a < cols; ++a) { + b[a] = -0.5 * 0.0 * (dx[0] * dx[0]) + 3.0 * (y[cols + a] - y[a]); + b[(n - 1) * cols + a] = + 0.5 * 0.0 * (dx[n - 2] * dx[n - 2]) + 3.0 * (y[(n - 1) * cols + a] - y[(n - 2) * cols + a]); + } + Vec dl(lower.begin(), lower.end() - 1), du(upper.begin() + 1, upper.end()); + gtsv(dl, diagonal, du, b, cols); + return b; +} + +double sign(double v) { return v > 0.0 ? 1.0 : (v < 0.0 ? -1.0 : (v == 0.0 ? 0.0 : v)); } + +double pchip_edge(double h0, double h1, double m0, double m1) { + double d = ((2.0 * h0 + h1) * m0 - h0 * m1) / (h0 + h1); + const bool mask = sign(d) != sign(m0); + const bool mask2 = (sign(m0) != sign(m1)) && (std::fabs(d) > 3.0 * std::fabs(m0)); + if (mask) d = 0.0; + else if (mask2) d = 3.0 * m0; + return d; +} + +// scipy PchipInterpolator derivatives at the knots. +Vec pchip_slopes(const Vec& x, const Vec& y, int64_t cols) { + const int64_t n = static_cast(x.size()); + Vec h(n - 1), m((n - 1) * cols), dk(n * cols, 0.0); + for (int64_t k = 0; k < n - 1; ++k) { + h[k] = x[k + 1] - x[k]; + for (int64_t a = 0; a < cols; ++a) m[k * cols + a] = (y[(k + 1) * cols + a] - y[k * cols + a]) / h[k]; + } + if (n == 2) { + for (int64_t a = 0; a < cols; ++a) dk[a] = dk[cols + a] = m[a]; + return dk; + } + for (int64_t k = 1; k < n - 1; ++k) { + const double w1 = 2.0 * h[k] + h[k - 1]; + const double w2 = h[k] + 2.0 * h[k - 1]; + for (int64_t a = 0; a < cols; ++a) { + const double left = m[(k - 1) * cols + a], right = m[k * cols + a]; + const bool condition = sign(right) != sign(left) || right == 0.0 || left == 0.0; + if (condition) { + dk[k * cols + a] = 0.0; + } else { + const double whmean = (w1 / left + w2 / right) / (w1 + w2); + dk[k * cols + a] = 1.0 / whmean; + } + } + } + for (int64_t a = 0; a < cols; ++a) { + dk[a] = pchip_edge(h[0], h[1], m[a], m[cols + a]); + dk[(n - 1) * cols + a] = pchip_edge(h[n - 2], h[n - 3], m[(n - 2) * cols + a], m[(n - 3) * cols + a]); + } + return dk; +} + +// ---- scipy RotationSpline ---- + +using Mat3 = std::array; + +Mat3 skew(const double* v) { return {0.0, -v[2], v[1], v[2], 0.0, -v[0], -v[1], v[0], 0.0}; } + +Mat3 matmul(const Mat3& a, const Mat3& b) { + Mat3 r{}; + for (int i = 0; i < 3; ++i) + for (int j = 0; j < 3; ++j) + r[i * 3 + j] = a[i * 3 + 0] * b[0 * 3 + j] + a[i * 3 + 1] * b[1 * 3 + j] + a[i * 3 + 2] * b[2 * 3 + j]; + return r; +} + +void matvec(const Mat3& a, const double* v, double* out) { + for (int i = 0; i < 3; ++i) out[i] = a[i * 3 + 0] * v[0] + a[i * 3 + 1] * v[1] + a[i * 3 + 2] * v[2]; +} + +double norm3(const double* v) { return std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]); } + +void cross(const double* a, const double* b, double* out) { + out[0] = a[1] * b[2] - a[2] * b[1]; + out[1] = a[2] * b[0] - a[0] * b[2]; + out[2] = a[0] * b[1] - a[1] * b[0]; +} + +Mat3 rate_to_rotvec_dot(const double* rotvec) { + const double nm = norm3(rotvec); + const double k = nm > 1e-4 ? (1.0 - 0.5 * nm / std::tan(0.5 * nm)) / (nm * nm) + : 1.0 / 12.0 + 1.0 / 720.0 * (nm * nm); + const Mat3 s = skew(rotvec), ss = matmul(s, s); + Mat3 r{}; + for (int i = 0; i < 9; ++i) r[i] = ((i % 4 == 0 ? 1.0 : 0.0) + 0.5 * s[i]) + k * ss[i]; + return r; +} + +Mat3 rotvec_dot_to_rate(const double* rotvec) { + const double nm = norm3(rotvec); + double k1, k2; + if (nm > 1e-4) { + k1 = (1.0 - std::cos(nm)) / (nm * nm); + k2 = (nm - std::sin(nm)) / std::pow(nm, 3.0); + } else { + k1 = 0.5 - (nm * nm) / 24.0; + k2 = 1.0 / 6.0 - (nm * nm) / 120.0; + } + const Mat3 s = skew(rotvec), ss = matmul(s, s); + Mat3 r{}; + for (int i = 0; i < 9; ++i) r[i] = ((i % 4 == 0 ? 1.0 : 0.0) - k1 * s[i]) + k2 * ss[i]; + return r; +} + +void nonlinear_term(const double* rv, const double* rvd, double* out) { + const double nm = norm3(rv); + const double dp = rv[0] * rvd[0] + rv[1] * rvd[1] + rv[2] * rvd[2]; + double cp[3], ccp[3], dccp[3]; + cross(rv, rvd, cp); + cross(rv, cp, ccp); + cross(rvd, cp, dccp); + double k1, k2, k3; + if (nm > 1e-4) { + const double sn = std::sin(nm), cs = std::cos(nm); + k1 = (-nm * sn - 2.0 * (cs - 1.0)) / std::pow(nm, 4.0); + k2 = (-2.0 * nm + 3.0 * sn - nm * cs) / std::pow(nm, 5.0); + k3 = (nm - sn) / std::pow(nm, 3.0); + } else { + k1 = 1.0 / 12.0 - (nm * nm) / 180.0; + k2 = -1.0 / 60.0 + (nm * nm) / 12604.0; + k3 = 1.0 / 6.0 - (nm * nm) / 120.0; + } + for (int i = 0; i < 3; ++i) out[i] = dp * (k1 * cp[i] + k2 * ccp[i]) + k3 * dccp[i]; +} + +// LAPACK dgbsv (reference dgbtf2 + dgbtrs) for one right-hand side. ``ab`` +// holds A in band storage with ``kl`` extra rows for fill-in: ldab = +// 2*kl+ku+1 and ab[(kl+ku+i-j) + ldab*j] = A(i, j), all zero-based. +void gbsv(int64_t n, int64_t kl, int64_t ku, Vec& ab, Vec& b) { + const int64_t ldab = 2 * kl + ku + 1, kv = ku + kl; + auto at = [&](int64_t row, int64_t col) -> double& { return ab[row + ldab * col]; }; + std::vector pivots(n); + int64_t ju = 0; + for (int64_t j = 0; j < n; ++j) { + const int64_t km = std::min(kl, n - 1 - j); + int64_t jp = 0; + double best = std::fabs(at(kv, j)); + for (int64_t i = 1; i <= km; ++i) { + if (std::fabs(at(kv + i, j)) > best) { best = std::fabs(at(kv + i, j)); jp = i; } + } + pivots[j] = jp + j; + if (at(kv + jp, j) == 0.0) throw std::invalid_argument("singular matrix"); + ju = std::max(ju, std::min(j + ku + jp, n - 1)); + if (jp != 0) { + for (int64_t c = 0; c <= ju - j; ++c) std::swap(at(kv + jp - c, j + c), at(kv - c, j + c)); + } + if (km > 0) { + const double reciprocal = 1.0 / at(kv, j); + for (int64_t i = 1; i <= km; ++i) at(kv + i, j) = reciprocal * at(kv + i, j); + for (int64_t c = 1; c <= ju - j; ++c) { + const double y = at(kv - c, j + c); + if (y != 0.0) { + const double temp = -1.0 * y; + for (int64_t i = 1; i <= km; ++i) at(kv - c + i, j + c) = at(kv - c + i, j + c) + at(kv + i, j) * temp; + } + } + } + } + // Forward: L, with the row interchanges. + for (int64_t j = 0; j + 1 < n; ++j) { + const int64_t lm = std::min(kl, n - 1 - j); + if (pivots[j] != j) std::swap(b[pivots[j]], b[j]); + if (b[j] != 0.0) { + const double temp = -1.0 * b[j]; + for (int64_t i = 1; i <= lm; ++i) b[j + i] = b[j + i] + at(kv + i, j) * temp; + } + } + // Backward: U, a band of width kl+ku above the diagonal (dtbsv). + const int64_t k = kl + ku; + for (int64_t j = n - 1; j >= 0; --j) { + if (b[j] != 0.0) { + b[j] = b[j] / at(k, j); + const double temp = b[j]; + for (int64_t i = j - 1; i >= std::max(0, j - k); --i) b[i] = b[i] - temp * at(k + i - j, j); + } + } +} + +struct RotationSpline { + Vec times; + std::vector rotations; + Vec coefficients; // ppoly layout, 3 columns + + RotationSpline(const Vec& knots, const std::vector& raw) : times(knots) { + const int64_t n = static_cast(knots.size()); + rotations.reserve(n); + for (const Quat& q : raw) rotations.push_back(sp_from_quat(q)); + const int64_t m = n - 1; + Vec dt(m), rotvecs(m * 3), rates(m * 3), rotvecs_dot; + for (int64_t k = 0; k < m; ++k) { + dt[k] = times[k + 1] - times[k]; + sp_as_rotvec(sp_compose(sp_inv(rotations[k]), rotations[k + 1]), &rotvecs[k * 3]); + for (int a = 0; a < 3; ++a) rates[k * 3 + a] = rotvecs[k * 3 + a] / dt[k]; + } + if (n == 2) { + rotvecs_dot = rates; + } else { + solve_rates(dt, rates, rotvecs, rotvecs_dot); + } + coefficients.assign(m * 4 * 3, 0.0); + for (int64_t k = 0; k < m; ++k) { + for (int a = 0; a < 3; ++a) { + const double rv = rotvecs[k * 3 + a], ar = rates[k * 3 + a], rd = rotvecs_dot[k * 3 + a]; + coefficients[(k * 4 + 0) * 3 + a] = (-2.0 * rv + dt[k] * ar + dt[k] * rd) / std::pow(dt[k], 3.0); + coefficients[(k * 4 + 1) * 3 + a] = (3.0 * rv - 2.0 * dt[k] * ar - dt[k] * rd) / (dt[k] * dt[k]); + coefficients[(k * 4 + 2) * 3 + a] = ar; + coefficients[(k * 4 + 3) * 3 + a] = 0.0; + } + } + } + + // scipy RotationSpline._solve_for_angular_rates; ``rates`` comes back as the + // knot angular rates, ``rotvecs_dot`` as the segment rotation-vector rates. + static void solve_rates(const Vec& dt, Vec& rates, const Vec& rotvecs, Vec& rotvecs_dot) { + const int64_t m = static_cast(dt.size()); + const double first[3] = {rates[0], rates[1], rates[2]}; + std::vector a(m), a_inv(m); + for (int64_t k = 0; k < m; ++k) { + a[k] = rate_to_rotvec_dot(&rotvecs[k * 3]); + a_inv[k] = rotvec_dot_to_rate(&rotvecs[k * 3]); + } + // Block-tridiagonal system of m-1 unknown 3-vectors in (5, 5) band form. + const int64_t size = 3 * (m - 1), kl = 5, ku = 5, ldab = 2 * kl + ku + 1; + Vec band(ldab * size, 0.0); + auto set = [&](int64_t i, int64_t j, double value) { band[(kl + ku + i - j) + ldab * j] = value; }; + for (int64_t k = 0; k < m - 2; ++k) { + for (int r = 0; r < 3; ++r) { + for (int c = 0; c < 3; ++c) { + set(3 * (k + 1) + r, 3 * k + c, 2.0 * a_inv[k + 1][r * 3 + c] / dt[k + 1]); + set(3 * k + r, 3 * (k + 1) + c, 2.0 * a[k + 1][r * 3 + c] / dt[k + 1]); + } + } + } + for (int64_t k = 0; k < m - 1; ++k) { + const double diagonal = 4.0 * (1.0 / dt[k] + 1.0 / dt[k + 1]); + for (int r = 0; r < 3; ++r) set(3 * k + r, 3 * k + r, diagonal); + } + Vec b0(size); + for (int64_t k = 0; k < m - 1; ++k) { + for (int r = 0; r < 3; ++r) + b0[3 * k + r] = 6.0 * (rotvecs[k * 3 + r] * std::pow(dt[k], -2.0) + + rotvecs[(k + 1) * 3 + r] * std::pow(dt[k + 1], -2.0)); + } + double product[3]; + matvec(a_inv[0], first, product); + for (int r = 0; r < 3; ++r) b0[r] -= 2.0 / dt[0] * product[r]; + matvec(a[m - 1], &rates[(m - 1) * 3], product); + for (int r = 0; r < 3; ++r) b0[size - 3 + r] -= 2.0 / dt[m - 1] * product[r]; + + rotvecs_dot.assign(m * 3, 0.0); + for (int iteration = 0; iteration < 10; ++iteration) { + for (int64_t k = 0; k < m; ++k) matvec(a[k], &rates[k * 3], &rotvecs_dot[k * 3]); + Vec b(size); + for (int64_t k = 0; k < m - 1; ++k) { + double delta_beta[3]; + nonlinear_term(&rotvecs[k * 3], &rotvecs_dot[k * 3], delta_beta); + for (int r = 0; r < 3; ++r) b[3 * k + r] = b0[3 * k + r] - delta_beta[r]; + } + Vec work = band; + gbsv(size, kl, ku, work, b); + bool converged = true; + for (int64_t i = 0; i < size; ++i) { + const double delta = std::fabs(b[i] - rates[i]); + rates[i] = b[i]; + if (!(delta < 1e-9 * (1.0 + std::fabs(b[i])))) converged = false; + } + if (converged) break; + } + for (int64_t k = 0; k < m; ++k) matvec(a[k], &rates[k * 3], &rotvecs_dot[k * 3]); + // Knot rates: the fixed first rate, then the solved interior ones. + Vec shifted(m * 3); + for (int r = 0; r < 3; ++r) shifted[r] = first[r]; + for (int64_t i = 0; i < 3 * (m - 1); ++i) shifted[3 + i] = rates[i]; + rates = shifted; + } + + Quat operator()(double t) const { + const int64_t m = static_cast(times.size()) - 1; + const int64_t k = ppoly_interval(times, t); + const double s = t - times[k]; + double rotvec[3]; + for (int a = 0; a < 3; ++a) { + const double coeff[4] = {coefficients[(k * 4 + 0) * 3 + a], coefficients[(k * 4 + 1) * 3 + a], + coefficients[(k * 4 + 2) * 3 + a], coefficients[(k * 4 + 3) * 3 + a]}; + rotvec[a] = ppoly_value(coeff, s); + } + int64_t index = static_cast(std::upper_bound(times.begin(), times.end(), t) - times.begin()) - 1; + index = std::min(std::max(index, 0), m - 1); + return sp_compose(rotations[index], sp_from_rotvec(rotvec)); + } +}; + +Quat slerp(const Quat& left, const Quat& right, double weight) { + const Quat first = so3_normalize(left); + Quat second = so3_normalize(right); + double dot = first.x * second.x + first.y * second.y + first.z * second.z + first.w * second.w; + if (dot < 0.0) second = {-second.x, -second.y, -second.z, -second.w}; + dot = std::min(std::max(std::fabs(dot), 0.0), 1.0); + Quat out; + if (dot > 0.9995) { + out = {first.x + weight * (second.x - first.x), first.y + weight * (second.y - first.y), + first.z + weight * (second.z - first.z), first.w + weight * (second.w - first.w)}; + } else { + const double angle = std::acos(dot), denominator = std::sin(angle); + const double a = std::sin((1.0 - weight) * angle) / denominator, b = std::sin(weight * angle) / denominator; + out = {a * first.x + b * second.x, a * first.y + b * second.y, a * first.z + b * second.z, + a * first.w + b * second.w}; + } + return so3_normalize(out); +} + +// ---- runtime resample_uniform ---- + +enum Kind { kLinear = 0, kHold = 1, kPose = 2 }; +enum Method { kLinearMethod = 0, kPchip = 1, kCubic = 2 }; + +int64_t resampled_horizon(int64_t length, double source_hz, double target_hz) { + return std::max(1, static_cast(std::floor(length * target_hz / source_hz + 0.5))); +} + +int64_t output_width(int kind, int64_t width) { return kind == kPose ? 7 : width; } + +// One row: float32 [T, D] in, float32 [count, D'] out. +void resample_row(const float* input, int64_t steps, int64_t width, double source_hz, double target_hz, + int kind, int64_t count, bool identity, int method, float* output) { + Vec values(steps * width); + for (int64_t i = 0; i < steps * width; ++i) values[i] = input[i]; + const int64_t out_width = output_width(kind, width); + if (identity) { + for (int64_t i = 0; i < steps * width; ++i) output[i] = static_cast(values[i]); + return; + } + Vec target_time(count), source_time(steps), interpolation_time(count); + for (int64_t i = 0; i < count; ++i) target_time[i] = static_cast(i) / target_hz; + for (int64_t i = 0; i < steps; ++i) source_time[i] = static_cast(i) / source_hz; + std::vector left(count), right(count); + Vec weight(count); + for (int64_t i = 0; i < count; ++i) { + interpolation_time[i] = std::min(target_time[i], source_time[steps - 1]); + const double position = target_time[i] * source_hz; + left[i] = std::min(static_cast(std::floor(position)), steps - 1); + right[i] = std::min(left[i] + 1, steps - 1); + weight[i] = position - static_cast(left[i]); + } + Vec result(count * out_width); + if (kind == kHold) { + for (int64_t i = 0; i < count; ++i) + for (int64_t a = 0; a < width; ++a) result[i * out_width + a] = values[left[i] * width + a]; + } else { + const bool pose = kind == kPose; + const int64_t cols = pose ? 3 : width; + Vec continuous(steps * cols); + for (int64_t t = 0; t < steps; ++t) + for (int64_t a = 0; a < cols; ++a) continuous[t * cols + a] = values[t * width + a]; + if (steps == 1) { + for (int64_t i = 0; i < count; ++i) + for (int64_t a = 0; a < cols; ++a) result[i * out_width + a] = continuous[a]; + } else if (method == kPchip) { + const Vec c = hermite_coefficients(source_time, continuous, pchip_slopes(source_time, continuous, cols), cols); + ppoly_evaluate(source_time, c, cols, interpolation_time, out_width, 0, result.data()); + } else if (method == kCubic) { + const Vec c = hermite_coefficients(source_time, continuous, natural_cubic_slopes(source_time, continuous, cols), cols); + ppoly_evaluate(source_time, c, cols, interpolation_time, out_width, 0, result.data()); + // Guard: fall back to PCHIP on axes that materially leave their envelope. + std::vector fallback; + for (int64_t a = 0; a < cols; ++a) { + double lower = continuous[a], upper = continuous[a]; + for (int64_t t = 1; t < steps; ++t) { + lower = std::min(lower, continuous[t * cols + a]); + upper = std::max(upper, continuous[t * cols + a]); + } + const double tolerance = 1e-9 + 0.05 * (upper - lower); + for (int64_t i = 0; i < count; ++i) { + const double v = result[i * out_width + a]; + if (v < lower - tolerance || v > upper + tolerance) { fallback.push_back(a); break; } + } + } + if (!fallback.empty()) { + const int64_t f = static_cast(fallback.size()); + Vec subset(steps * f); + for (int64_t t = 0; t < steps; ++t) + for (int64_t j = 0; j < f; ++j) subset[t * f + j] = continuous[t * cols + fallback[j]]; + const Vec c = hermite_coefficients(source_time, subset, pchip_slopes(source_time, subset, f), f); + Vec sampled(count * f); + ppoly_evaluate(source_time, c, f, interpolation_time, f, 0, sampled.data()); + for (int64_t i = 0; i < count; ++i) + for (int64_t j = 0; j < f; ++j) result[i * out_width + fallback[j]] = sampled[i * f + j]; + } + } else { + for (int64_t i = 0; i < count; ++i) + for (int64_t a = 0; a < cols; ++a) { + const double lo = values[left[i] * width + a], hi = values[right[i] * width + a]; + result[i * out_width + a] = lo + weight[i] * (hi - lo); + } + } + if (pose) { + std::vector quaternions(steps); + for (int64_t t = 0; t < steps; ++t) + quaternions[t] = {values[t * width + 3], values[t * width + 4], values[t * width + 5], values[t * width + 6]}; + auto store = [&](int64_t i, const Quat& q) { + result[i * out_width + 3] = q.x; + result[i * out_width + 4] = q.y; + result[i * out_width + 5] = q.z; + result[i * out_width + 6] = q.w; + }; + if (steps == 1) { + for (int64_t i = 0; i < count; ++i) store(i, quaternions[0]); + } else if (method == kLinearMethod) { + for (int64_t i = 0; i < count; ++i) store(i, slerp(quaternions[left[i]], quaternions[right[i]], weight[i])); + } else { + const RotationSpline spline(source_time, quaternions); + for (int64_t i = 0; i < count; ++i) store(i, spline(interpolation_time[i])); + } + } + } + for (int64_t i = 0; i < count * out_width; ++i) output[i] = static_cast(result[i]); +} + +int parse_kind(const std::string& kind) { + if (kind == "linear") return kLinear; + if (kind == "hold") return kHold; + if (kind == "pose_xyzw") return kPose; + throw std::invalid_argument("unsupported resampling kind"); +} + +int parse_method(const std::string& method) { + if (method == "linear") return kLinearMethod; + if (method == "pchip") return kPchip; + if (method == "cubic") return kCubic; + throw std::invalid_argument("method must be one of ['cubic', 'linear', 'pchip']"); +} + +using F32 = py::array_t; + +void check_finite(const F32& values) { + const float* data = values.data(); + for (py::ssize_t i = 0; i < values.size(); ++i) + if (!std::isfinite(data[i])) throw std::invalid_argument("values must be a non-empty finite [T, D] array"); +} + +F32 resample_uniform_batch(const F32& values, double source_hz, double target_hz, const std::string& kind_name, + int64_t target_length, const std::string& method_name, int threads) { + if (values.ndim() != 3 || values.shape(0) < 1 || values.shape(1) < 1) + throw std::invalid_argument("values must have non-empty shape [B, T, D]"); + if (!(std::isfinite(source_hz) && source_hz > 0.0)) throw std::invalid_argument("source_hz must be positive and finite"); + if (!(std::isfinite(target_hz) && target_hz > 0.0)) throw std::invalid_argument("target_hz must be positive and finite"); + check_finite(values); + const int kind = parse_kind(kind_name), method = parse_method(method_name); + const int64_t batch = values.shape(0), steps = values.shape(1), width = values.shape(2); + if (kind == kPose && width != 7) throw std::invalid_argument("pose_xyzw values must have shape [T, 7]"); + const int64_t standard = resampled_horizon(steps, source_hz, target_hz); + const bool identity = source_hz == target_hz && (target_length < 0 || target_length == standard); + const int64_t count = identity ? steps : (target_length < 0 ? standard : target_length); + if (count < 1) throw std::invalid_argument("target_length must be a positive integer"); + const int64_t out_width = output_width(kind, width); + F32 output({batch, count, out_width}); + const float* in = values.data(); + float* out = output.mutable_data(); + { + py::gil_scoped_release release; + ac2::parallel_for(batch, threads, [&](int64_t row) { + resample_row(in + row * steps * width, steps, width, source_hz, target_hz, kind, count, identity, method, + out + row * count * out_width); + }); + } + return output; +} + +// Elementwise pose maps over [B, T, D] float32 arrays. +template +F32 map_rows(const F32& values, int64_t in_width, int64_t out_width, int threads, Row&& row) { + if (values.ndim() != 3 || values.shape(2) != in_width || values.shape(1) < 1) + throw std::invalid_argument("expected shape [B, T, " + std::to_string(in_width) + "]"); + check_finite(values); + const int64_t batch = values.shape(0), steps = values.shape(1); + F32 output({batch, steps, out_width}); + const float* in = values.data(); + float* out = output.mutable_data(); + { + py::gil_scoped_release release; + ac2::parallel_for(batch, threads, [&](int64_t b) { + row(in + b * steps * in_width, steps, out + b * steps * out_width); + }); + } + return output; +} + +// runtime integrate_eef_deltas: [T, 6] body increments -> [T, 7] absolute poses. +F32 integrate_eef_deltas_batch(const F32& deltas, int threads) { + return map_rows(deltas, 6, 7, threads, [](const float* in, int64_t steps, float* out) { + double position[3] = {0.0, 0.0, 0.0}; + Quat current{0.0, 0.0, 0.0, 1.0}; + for (int64_t t = 0; t < steps; ++t) { + for (int a = 0; a < 3; ++a) { + position[a] = position[a] + static_cast(in[t * 6 + a]); + out[t * 7 + a] = static_cast(position[a]); + } + const Quat increment = so3_from_rotvec(in[t * 6 + 3], in[t * 6 + 4], in[t * 6 + 5]); + current = so3_multiply(current, increment); + out[t * 7 + 3] = static_cast(current.x); + out[t * 7 + 4] = static_cast(current.y); + out[t * 7 + 5] = static_cast(current.z); + out[t * 7 + 6] = static_cast(current.w); + } + }); +} + +// runtime eef_pose_from_components(..., "rotation_vector"). +F32 eef_pose_from_rotvec_batch(const F32& poses, int threads) { + return map_rows(poses, 6, 7, threads, [](const float* in, int64_t steps, float* out) { + for (int64_t t = 0; t < steps; ++t) { + for (int a = 0; a < 3; ++a) out[t * 7 + a] = in[t * 6 + a]; + const Quat q = so3_from_rotvec(in[t * 6 + 3], in[t * 6 + 4], in[t * 6 + 5]); + out[t * 7 + 3] = static_cast(q.x); + out[t * 7 + 4] = static_cast(q.y); + out[t * 7 + 5] = static_cast(q.z); + out[t * 7 + 6] = static_cast(q.w); + } + }); +} + +// runtime eef_pose_absolute: [T, 7] xyz + xyzw -> [T, 6] xyz + rotation vector. +F32 eef_pose_absolute_batch(const F32& poses, int threads) { + return map_rows(poses, 7, 6, threads, [](const float* in, int64_t steps, float* out) { + for (int64_t t = 0; t < steps; ++t) { + for (int a = 0; a < 3; ++a) out[t * 6 + a] = in[t * 7 + a]; + double rotvec[3]; + so3_to_rotvec({in[t * 7 + 3], in[t * 7 + 4], in[t * 7 + 5], in[t * 7 + 6]}, rotvec); + for (int a = 0; a < 3; ++a) out[t * 6 + 3 + a] = static_cast(rotvec[a]); + } + }); +} + +// runtime eef_pose_deltas: local [T, 7] path -> [T, 6] body increments. +F32 eef_pose_deltas_batch(const F32& poses, int threads) { + return map_rows(poses, 7, 6, threads, [](const float* in, int64_t steps, float* out) { + Quat previous{0.0, 0.0, 0.0, 1.0}; + for (int64_t t = 0; t < steps; ++t) { + for (int a = 0; a < 3; ++a) { + const double before = t == 0 ? 0.0 : static_cast(in[(t - 1) * 7 + a]); + out[t * 6 + a] = static_cast(static_cast(in[t * 7 + a]) - before); + } + const Quat q = so3_normalize({in[t * 7 + 3], in[t * 7 + 4], in[t * 7 + 5], in[t * 7 + 6]}); + double rotvec[3]; + so3_to_rotvec(so3_multiply({-previous.x, -previous.y, -previous.z, previous.w}, q), rotvec); + for (int a = 0; a < 3; ++a) out[t * 6 + 3 + a] = static_cast(rotvec[a]); + previous = q; + } + }); +} + +// Re-express per-step rotation increments between world-frame (left) and +// body-frame (right) multiplication, as scipy ``Rotation`` composes them. +F32 rotvec_left_to_right_batch(const F32& increments, int threads) { + return map_rows(increments, 3, 3, threads, [](const float* in, int64_t steps, float* out) { + Quat current{0.0, 0.0, 0.0, 1.0}; + for (int64_t t = 0; t < steps; ++t) { + const double v[3] = {in[t * 3], in[t * 3 + 1], in[t * 3 + 2]}; + const Quat following = sp_compose(sp_from_rotvec(v), current); + double rotvec[3]; + sp_as_rotvec(sp_compose(sp_inv(current), following), rotvec); + for (int a = 0; a < 3; ++a) out[t * 3 + a] = static_cast(rotvec[a]); + current = following; + } + }); +} + +F32 rotvec_right_to_left_batch(const F32& increments, int threads) { + return map_rows(increments, 3, 3, threads, [](const float* in, int64_t steps, float* out) { + Quat current{0.0, 0.0, 0.0, 1.0}; + for (int64_t t = 0; t < steps; ++t) { + const double v[3] = {in[t * 3], in[t * 3 + 1], in[t * 3 + 2]}; + const Quat following = sp_compose(current, sp_from_rotvec(v)); + double rotvec[3]; + sp_as_rotvec(sp_compose(following, sp_inv(current)), rotvec); + for (int a = 0; a < 3; ++a) out[t * 3 + a] = static_cast(rotvec[a]); + current = following; + } + }); +} + +} // namespace + +void register_resample(py::module_& module) { + module.def("resample_uniform_batch", &resample_uniform_batch, py::arg("values"), py::arg("source_hz"), + py::arg("target_hz"), py::arg("kind"), py::arg("target_length") = -1, + py::arg("method") = "linear", py::arg("threads") = 0, + "runtime resample_uniform over each [T, D] row of a [B, T, D] batch."); + module.def("integrate_eef_deltas_batch", &integrate_eef_deltas_batch, py::arg("deltas"), py::arg("threads") = 0); + module.def("eef_pose_from_rotvec_batch", &eef_pose_from_rotvec_batch, py::arg("poses"), py::arg("threads") = 0); + module.def("eef_pose_absolute_batch", &eef_pose_absolute_batch, py::arg("poses"), py::arg("threads") = 0); + module.def("eef_pose_deltas_batch", &eef_pose_deltas_batch, py::arg("poses"), py::arg("threads") = 0); + module.def("rotvec_left_to_right_batch", &rotvec_left_to_right_batch, py::arg("increments"), + py::arg("threads") = 0); + module.def("rotvec_right_to_left_batch", &rotvec_right_to_left_batch, py::arg("increments"), + py::arg("threads") = 0); +} diff --git a/actioncodec2_kernels/setbpe.cpp b/actioncodec2_kernels/setbpe.cpp new file mode 100644 index 0000000000000000000000000000000000000000..9a9731dd4ba1a1d6cfd4a6f2027afe49b3e2936d --- /dev/null +++ b/actioncodec2_kernels/setbpe.cpp @@ -0,0 +1,394 @@ +// Set-BPE serialization: quantized cell trajectories <-> token streams. +// +// A port of the runtime's SpatialVocabulary / TemporalVocabulary (Python, +// ``tokenization/setbpe_{spatial,temporal}.py``). Every merge is applied from +// a priority queue ordered exactly as the Python heaps order their tuples, so +// the token streams are identical, not merely equivalent. + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "parallel.hpp" + +namespace py = pybind11; + +namespace { + +int lowest_bit(uint64_t mask) { return mask ? __builtin_ctzll(mask) : -1; } + +uint64_t pair_key(int64_t left, int64_t right) { + return (static_cast(left) << 32) | static_cast(right); +} + +struct Rule { + int64_t child; + int64_t rank; +}; + +class SetBPE { + public: + SetBPE(std::vector bins, int64_t lookahead, + const std::vector>& spatial_merges, + const std::vector>& temporal_merges) + : bins_(std::move(bins)), lookahead_(lookahead) { + const int64_t dimension = static_cast(bins_.size()); + if (dimension < 1 || dimension > 63) throw std::invalid_argument("bins must have 1..63 slots"); + if (lookahead_ < 1) throw std::invalid_argument("lookahead must be positive"); + offsets_.assign(dimension + 1, 0); + for (int64_t d = 0; d < dimension; ++d) { + if (bins_[d] < 1) throw std::invalid_argument("every dimension needs at least one bin"); + offsets_[d + 1] = offsets_[d] + bins_[d]; + } + atom_count_ = offsets_[dimension]; + spatial_size_ = atom_count_ + static_cast(spatial_merges.size()); + support_.assign(spatial_size_, 0); + expansion_.assign(spatial_size_ * dimension, -1); + for (int64_t d = 0; d < dimension; ++d) { + for (int64_t v = 0; v < bins_[d]; ++v) { + const int64_t token = offsets_[d] + v; + support_[token] = uint64_t{1} << d; + expansion_[token * dimension + d] = v; + } + } + for (size_t rank = 0; rank < spatial_merges.size(); ++rank) { + const auto [left, right] = spatial_merges[rank]; + const int64_t token = atom_count_ + static_cast(rank); + if (left < 0 || right < 0 || left >= token || right >= token) + throw std::invalid_argument("a merge rule must reference already-created tokens"); + if (support_[left] & support_[right]) + throw std::invalid_argument("a merge rule joins two overlapping supports"); + support_[token] = support_[left] | support_[right]; + for (int64_t d = 0; d < dimension; ++d) { + expansion_[token * dimension + d] = + std::max(expansion_[left * dimension + d], expansion_[right * dimension + d]); + } + spatial_rules_[pair_key(std::min(left, right), std::max(left, right))] = + Rule{token, static_cast(rank)}; + } + spatial_anchor_.resize(spatial_size_); + for (int64_t token = 0; token < spatial_size_; ++token) + spatial_anchor_[token] = lowest_bit(support_[token]); + + // Rooted temporal tokens: a pattern of (time offset, spatial token) pairs. + vocab_size_ = spatial_size_ + static_cast(temporal_merges.size()); + patterns_.resize(vocab_size_); + root_.resize(vocab_size_); + for (int64_t token = 0; token < spatial_size_; ++token) { + patterns_[token] = {{0, token}}; + root_[token] = token; + } + for (size_t rank = 0; rank < temporal_merges.size(); ++rank) { + const auto [left, right, offset] = temporal_merges[rank]; + const int64_t child = spatial_size_ + static_cast(rank); + if (left < 0 || right < 0 || left >= child || right >= child) + throw std::invalid_argument("a temporal merge must reference already-created tokens"); + if (offset < 1 || offset >= lookahead_) + throw std::invalid_argument("a temporal merge offset must be in [1, lookahead)"); + std::vector> pattern = patterns_[left]; + for (const auto& [time, token] : patterns_[right]) pattern.emplace_back(time + offset, token); + std::sort(pattern.begin(), pattern.end()); + const int64_t span = 1 + pattern.back().first; + if (span > lookahead_) throw std::invalid_argument("temporal merge exceeds the lookahead horizon"); + std::vector masks(span, 0); + int roots = 0; + for (const auto& [time, token] : pattern) { + if (masks[time] & support_[token]) + throw std::invalid_argument("temporal merge joins overlapping spatial supports"); + masks[time] |= support_[token]; + roots += time == 0; + } + if (roots != 1) + throw std::invalid_argument("a rooted temporal token must have exactly one time-zero spatial token"); + patterns_[child] = std::move(pattern); + root_[child] = root_[left]; + const auto key = std::make_tuple(left, right, offset); + if (temporal_rules_.count(key)) throw std::invalid_argument("duplicate temporal merge rule"); + temporal_rules_[key] = Rule{child, static_cast(rank)}; + } + } + + int64_t vocab_size() const { return vocab_size_; } + + // One dense [T, D] trajectory with -1 for semantically absent slots. + std::vector encode(const int64_t* cells, int64_t steps) const { + const int64_t dimension = static_cast(bins_.size()); + if (steps < 1) throw std::invalid_argument("every episode must have at least one step"); + uint64_t first_support = 0; + for (int64_t t = 0; t < steps; ++t) { + uint64_t support = 0; + for (int64_t d = 0; d < dimension; ++d) { + const int64_t value = cells[t * dimension + d]; + if (value < -1 || value >= bins_[d]) + throw std::invalid_argument("episodes must contain -1 or an in-range cell index"); + if (value >= 0) support |= uint64_t{1} << d; + } + if (!support) throw std::invalid_argument("every timestep needs at least one active coordinate"); + if (t == 0) first_support = support; + else if (support != first_support) + throw std::invalid_argument("active dimensions must remain constant within an episode"); + } + + // Instances, as in the Python encoder: token, anchor step, liveness. + std::vector tokens, anchors; + std::vector active; + std::vector>> at(steps); // (token, instance) + for (int64_t t = 0; t < steps; ++t) { + for (int64_t token : spatial_encode(cells + t * dimension)) { + at[t].emplace_back(token, static_cast(tokens.size())); + tokens.push_back(token); + anchors.push_back(t); + active.push_back(1); + } + } + + // (rank, anchor, root anchor dimension, left instance, right instance, child) + using Offer = std::tuple; + std::priority_queue, std::greater> frontier; + auto offer = [&](int64_t left_instance, int64_t right_instance) { + const auto found = temporal_rules_.find(std::make_tuple( + tokens[left_instance], tokens[right_instance], + anchors[right_instance] - anchors[left_instance])); + if (found == temporal_rules_.end()) return; + const Rule& rule = found->second; + frontier.emplace(rule.rank, anchors[left_instance], + spatial_anchor_[root_[rule.child]], left_instance, right_instance, + rule.child); + }; + for (int64_t anchor = 0; anchor < steps; ++anchor) { + for (int64_t offset = 1; offset < std::min(lookahead_, steps - anchor); ++offset) { + for (const auto& [left_token, left_instance] : at[anchor]) { + for (const auto& [right_token, right_instance] : at[anchor + offset]) { + offer(left_instance, right_instance); + } + } + } + } + auto erase = [&](int64_t step, int64_t token) { + auto& items = at[step]; + for (auto it = items.begin(); it != items.end(); ++it) { + if (it->first == token) { items.erase(it); return; } + } + }; + while (!frontier.empty()) { + const auto [rank, anchor, root_anchor, left_instance, right_instance, child] = frontier.top(); + frontier.pop(); + (void)rank; (void)root_anchor; + if (!active[left_instance] || !active[right_instance]) continue; + if (anchors[left_instance] != anchor) continue; + const auto resolved = temporal_rules_.find(std::make_tuple( + tokens[left_instance], tokens[right_instance], anchors[right_instance] - anchor)); + if (resolved == temporal_rules_.end() || resolved->second.child != child) continue; + active[left_instance] = 0; + active[right_instance] = 0; + erase(anchor, tokens[left_instance]); + erase(anchors[right_instance], tokens[right_instance]); + const int64_t new_instance = static_cast(tokens.size()); + tokens.push_back(child); + anchors.push_back(anchor); + active.push_back(1); + at[anchor].emplace_back(child, new_instance); + for (int64_t offset = 1; offset < lookahead_; ++offset) { + const int64_t future = anchor + offset; + if (future < steps) { + for (const auto& [token, other] : at[future]) offer(new_instance, other); + } + const int64_t past = anchor - offset; + if (past >= 0) { + for (const auto& [token, other] : at[past]) offer(other, new_instance); + } + } + } + + std::vector> ordered; + for (int64_t anchor = 0; anchor < steps; ++anchor) { + for (const auto& [token, instance] : at[anchor]) { + if (active[instance]) ordered.emplace_back(anchor, spatial_anchor_[root_[token]], token); + } + } + std::sort(ordered.begin(), ordered.end()); + std::vector stream; + stream.reserve(ordered.size()); + for (const auto& item : ordered) stream.push_back(std::get<2>(item)); + return stream; + } + + // Tokens -> dense [T, D] cells (-1 outside ``full``), validating like the runtime. + std::vector decode(const int64_t* stream, int64_t length, uint64_t full) const { + const int64_t dimension = static_cast(bins_.size()); + std::vector coverage{0}; + std::vector> spatial_tokens(1); + size_t front = 0; + for (int64_t index = 0; index < length; ++index) { + const int64_t token = stream[index]; + if (token < 0 || token >= vocab_size_) + throw std::invalid_argument("temporal token " + std::to_string(token) + " is outside the vocabulary"); + while (front < coverage.size() && coverage[front] == full) ++front; + if (front == coverage.size()) { + coverage.push_back(0); + spatial_tokens.emplace_back(); + } + const uint64_t uncovered = full & ~coverage[front]; + if (spatial_anchor_[root_[token]] != lowest_bit(uncovered)) + throw std::invalid_argument("temporal stream violates root order at the earliest incomplete step"); + for (const auto& [offset, spatial_token] : patterns_[token]) { + const size_t step = front + static_cast(offset); + while (step >= coverage.size()) { + coverage.push_back(0); + spatial_tokens.emplace_back(); + } + if (coverage[step] & support_[spatial_token]) + throw std::invalid_argument("temporal stream contains overlapping spatial support"); + coverage[step] |= support_[spatial_token]; + spatial_tokens[step].push_back(spatial_token); + } + } + for (uint64_t mask : coverage) { + if (mask != full) throw std::invalid_argument("temporal stream ended with an incomplete control step"); + } + std::vector cells(coverage.size() * dimension, -1); + for (size_t step = 0; step < coverage.size(); ++step) { + for (int64_t token : spatial_tokens[step]) { + for (int64_t d = 0; d < dimension; ++d) { + const int64_t value = expansion_[token * dimension + d]; + if (value >= 0) cells[step * dimension + d] = value; + } + } + } + return cells; + } + + private: + // Spatial Set-BPE of one dense row: merges in rank order, then anchor order. + std::vector spatial_encode(const int64_t* row) const { + const int64_t dimension = static_cast(bins_.size()); + std::vector present; + for (int64_t d = 0; d < dimension; ++d) { + if (row[d] >= 0) present.push_back(offsets_[d] + row[d]); + } + std::sort(present.begin(), present.end()); + using Candidate = std::tuple; // rank, child, left, right + std::priority_queue, std::greater> heap; + for (size_t i = 0; i < present.size(); ++i) { + for (size_t j = i + 1; j < present.size(); ++j) { + const auto found = spatial_rules_.find(pair_key(present[i], present[j])); + if (found != spatial_rules_.end()) + heap.emplace(found->second.rank, found->second.child, present[i], present[j]); + } + } + std::vector tokens = present; // a small set; linear lookups are fastest + auto contains = [&](int64_t token) { + return std::find(tokens.begin(), tokens.end(), token) != tokens.end(); + }; + while (!heap.empty()) { + const auto [rank, child, left, right] = heap.top(); + heap.pop(); + (void)rank; + if (!contains(left) || !contains(right)) continue; + tokens.erase(std::find(tokens.begin(), tokens.end(), left)); + tokens.erase(std::find(tokens.begin(), tokens.end(), right)); + for (int64_t other : tokens) { + const int64_t low = std::min(child, other), high = std::max(child, other); + const auto found = spatial_rules_.find(pair_key(low, high)); + if (found != spatial_rules_.end()) + heap.emplace(found->second.rank, found->second.child, low, high); + } + tokens.push_back(child); + } + std::sort(tokens.begin(), tokens.end(), [&](int64_t a, int64_t b) { + return std::make_pair(spatial_anchor_[a], a) < std::make_pair(spatial_anchor_[b], b); + }); + return tokens; + } + + struct TupleHash { + size_t operator()(const std::tuple& key) const { + const auto [a, b, c] = key; + return std::hash()((static_cast(a) * 1000003u) ^ + (static_cast(b) << 20) ^ static_cast(c)); + } + }; + + std::vector bins_; + int64_t lookahead_; + std::vector offsets_; + int64_t atom_count_ = 0, spatial_size_ = 0, vocab_size_ = 0; + std::vector support_; + std::vector expansion_; + std::vector spatial_anchor_; + std::unordered_map spatial_rules_; + std::vector>> patterns_; + std::vector root_; + std::unordered_map, Rule, TupleHash> temporal_rules_; +}; + +// Encode B dense [T, D] trajectories in parallel; returns one stream per row. +py::list encode_batch(const SetBPE& codec, + const py::array_t& cells, + int threads) { + if (cells.ndim() != 3) throw std::invalid_argument("cells must have shape [B, T, D]"); + const int64_t batch = cells.shape(0), steps = cells.shape(1), width = cells.shape(2); + const int64_t* data = cells.data(); + std::vector> streams(batch); + { + py::gil_scoped_release release; + ac2::parallel_for(batch, threads, [&](int64_t index) { + streams[index] = codec.encode(data + index * steps * width, steps); + }); + } + py::list output(batch); + for (int64_t index = 0; index < batch; ++index) output[index] = py::cast(streams[index]); + return output; +} + +// Decode B equal-horizon streams to dense [B, T, D] cells (-1 outside ``full``). +py::array_t decode_batch(const SetBPE& codec, const std::vector>& streams, + int64_t width, uint64_t full, int threads) { + const int64_t batch = static_cast(streams.size()); + if (batch < 1) throw std::invalid_argument("streams must not be empty"); + std::vector> cells(batch); + { + py::gil_scoped_release release; + ac2::parallel_for(batch, threads, [&](int64_t index) { + if (streams[index].empty()) throw std::invalid_argument("token stream must not be empty"); + cells[index] = codec.decode(streams[index].data(), static_cast(streams[index].size()), full); + }); + } + const int64_t steps = static_cast(cells[0].size()) / width; + for (const auto& row : cells) { + if (static_cast(row.size()) != steps * width) + throw std::invalid_argument("all token streams in a batch must decode to the same length"); + } + py::array_t output({batch, steps, width}); + int64_t* out = output.mutable_data(); + for (int64_t index = 0; index < batch; ++index) + std::copy(cells[index].begin(), cells[index].end(), out + index * steps * width); + return output; +} + +} // namespace + +void register_setbpe(py::module_& module) { + py::class_(module, "SetBPE") + .def(py::init, int64_t, const std::vector>&, + const std::vector>&>(), + py::arg("bins"), py::arg("lookahead"), py::arg("spatial_merges"), + py::arg("temporal_merges")) + .def_property_readonly("vocab_size", &SetBPE::vocab_size) + .def("encode_batch", &encode_batch, py::arg("cells"), py::arg("threads") = 0, + "Dense [B, T, D] cells (-1 = absent) -> one token stream per row.") + .def("decode_batch", &decode_batch, py::arg("streams"), py::arg("width"), + py::arg("full_support"), py::arg("threads") = 0, + "Equal-horizon streams -> dense [B, T, D] cells, -1 outside full_support."); +} diff --git a/actioncodec2_kernels/so3.cpp b/actioncodec2_kernels/so3.cpp new file mode 100644 index 0000000000000000000000000000000000000000..4fde99c6b5be37137d5ca0a508ef0ae5fb5dfd13 --- /dev/null +++ b/actioncodec2_kernels/so3.cpp @@ -0,0 +1,313 @@ +// Intrinsic first- and stateful second-order SO(3) error-feedback encoders. + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "parallel.hpp" + +namespace py = pybind11; + +namespace { + +struct Quaternion { + double x; + double y; + double z; + double w; +}; + +Quaternion normalize(Quaternion q) { + const double norm = std::sqrt(q.x * q.x + q.y * q.y + q.z * q.z + q.w * q.w); + q.x /= norm; + q.y /= norm; + q.z /= norm; + q.w /= norm; + return q; +} + +Quaternion multiply(const Quaternion& a, const Quaternion& b) { + return normalize({a.w * b.x + a.x * b.w + a.y * b.z - a.z * b.y, + a.w * b.y - a.x * b.z + a.y * b.w + a.z * b.x, + a.w * b.z + a.x * b.y - a.y * b.x + a.z * b.w, + a.w * b.w - a.x * b.x - a.y * b.y - a.z * b.z}); +} + +Quaternion conjugate(const Quaternion& q) { return {-q.x, -q.y, -q.z, q.w}; } + +Quaternion from_euler(const float* euler) { + const double roll = 0.5 * static_cast(euler[0]); + const double pitch = 0.5 * static_cast(euler[1]); + const double yaw = 0.5 * static_cast(euler[2]); + const double sr = std::sin(roll), cr = std::cos(roll); + const double sp = std::sin(pitch), cp = std::cos(pitch); + const double sy = std::sin(yaw), cy = std::cos(yaw); + return normalize({sr * cp * cy - cr * sp * sy, cr * sp * cy + sr * cp * sy, + cr * cp * sy - sr * sp * cy, cr * cp * cy + sr * sp * sy}); +} + +void to_rotvec(Quaternion q, double* output) { + if (q.w < 0.0) { + q.x = -q.x; + q.y = -q.y; + q.z = -q.z; + q.w = -q.w; + } + const double norm = std::sqrt(q.x * q.x + q.y * q.y + q.z * q.z); + const double scale = norm < 1e-12 ? 2.0 : 2.0 * std::atan2(norm, std::max(0.0, q.w)) / norm; + output[0] = q.x * scale; + output[1] = q.y * scale; + output[2] = q.z * scale; +} + +Quaternion from_rotvec(const double* value) { + const double angle = std::sqrt(value[0] * value[0] + value[1] * value[1] + + value[2] * value[2]); + const double scale = angle < 1e-12 ? 0.5 : std::sin(0.5 * angle) / angle; + return normalize({value[0] * scale, value[1] * scale, value[2] * scale, + std::cos(0.5 * angle)}); +} + +Quaternion from_input(const float* value, bool rotation_vector) { + if (!rotation_vector) return from_euler(value); + const double converted[3] = {static_cast(value[0]), static_cast(value[1]), + static_cast(value[2])}; + return from_rotvec(converted); +} + +Quaternion residual(const Quaternion& target, const Quaternion& reconstruction, bool left) { + return left ? multiply(target, conjugate(reconstruction)) + : multiply(conjugate(reconstruction), target); +} + +Quaternion apply(const Quaternion& state, const Quaternion& increment, bool left) { + return left ? multiply(increment, state) : multiply(state, increment); +} + +std::array as_rotvec(const Quaternion& value) { + std::array output{}; + to_rotvec(value, output.data()); + return output; +} + +struct SO3Layout { + const double* table; + const int32_t* bins; + int64_t width; + double limit; + bool smooth; + bool rotation_vector; + bool left; + int32_t input_order; + double projection_radius = std::numeric_limits::infinity(); +}; + +void project_primitive(double* value, double radius) { + if (!std::isfinite(radius)) return; + const double norm = std::sqrt(value[0]*value[0] + value[1]*value[1] + value[2]*value[2]); + if (norm > radius) { + const double scale = radius / norm; + for (int j = 0; j < 3; ++j) value[j] *= scale; + } +} + +inline double primitive(const SO3Layout& layout, int axis, int32_t index) { + return layout.table[static_cast(axis) * layout.width + index]; +} + +std::array nearest_choice(const SO3Layout& layout, + const std::array& request) { + std::array choice{}; + for (int axis = 0; axis < 3; ++axis) { + double best = std::numeric_limits::infinity(); + for (int32_t index = 0; index < layout.bins[axis]; ++index) { + const double distance = std::fabs(request[axis] - primitive(layout, axis, index)); + if (distance < best) { + best = distance; + choice[axis] = index; + } + } + } + return choice; +} + +std::vector smooth_axis(const SO3Layout& layout, int axis, double request, + int32_t nearest) { + std::vector output; + bool has_nearest = false; + const double radius = 2.0 * layout.limit + 1e-12; + for (int32_t index = 0; index < layout.bins[axis]; ++index) { + if (std::fabs(primitive(layout, axis, index) - request) <= radius) { + output.push_back(index); + has_nearest = has_nearest || index == nearest; + } + } + if (!has_nearest) output.push_back(nearest); + return output; +} + +bool earlier(double error, const std::array& choice, double best_error, + const std::array& best_choice, bool found) { + return !found || error < best_error || (error == best_error && choice < best_choice); +} + +std::array smooth_choice(const SO3Layout& layout, + const std::array& request, + const Quaternion& reconstruction, + const Quaternion& target, + const std::array& direction, + const std::array& nearest) { + const auto x = smooth_axis(layout, 0, request[0], nearest[0]); + const auto y = smooth_axis(layout, 1, request[1], nearest[1]); + const auto z = smooth_axis(layout, 2, request[2], nearest[2]); + std::array best_any{}, best_preferred{}; + double any_error = 0.0, preferred_error = 0.0; + bool found_any = false, found_preferred = false; + for (const int32_t ix : x) { + for (const int32_t iy : y) { + for (const int32_t iz : z) { + const std::array choice{ix, iy, iz}; + double vector[3] = {primitive(layout, 0, ix), primitive(layout, 1, iy), + primitive(layout, 2, iz)}; + project_primitive(vector, layout.projection_radius); + const Quaternion rebuilt = apply(reconstruction, from_rotvec(vector), layout.left); + const auto remaining = as_rotvec(residual(target, rebuilt, layout.left)); + const double error = std::sqrt(remaining[0] * remaining[0] + remaining[1] * remaining[1] + + remaining[2] * remaining[2]); + if (error > layout.limit + 1e-12) continue; + if (earlier(error, choice, any_error, best_any, found_any)) { + best_any = choice; + any_error = error; + found_any = true; + } + const double dot = vector[0] * direction[0] + vector[1] * direction[1] + + vector[2] * direction[2]; + if (dot >= -1e-15 && earlier(error, choice, preferred_error, best_preferred, + found_preferred)) { + best_preferred = choice; + preferred_error = error; + found_preferred = true; + } + } + } + } + return found_preferred ? best_preferred : (found_any ? best_any : nearest); +} + +void encode_trajectory(const SO3Layout& layout, const float* actions, int64_t steps, + int32_t* cells) { + Quaternion truth{0.0, 0.0, 0.0, 1.0}; + Quaternion reconstruction = truth; + std::array direction{}; + for (int64_t step = 0; step < steps; ++step) { + const Quaternion value = from_input(actions + 3 * step, layout.rotation_vector); + truth = layout.input_order == 0 ? value : apply(truth, value, layout.left); + const auto request = as_rotvec(residual(truth, reconstruction, layout.left)); + auto choice = nearest_choice(layout, request); + if (layout.smooth) { + choice = smooth_choice(layout, request, reconstruction, truth, direction, choice); + } + const double vector[3] = {primitive(layout, 0, choice[0]), + primitive(layout, 1, choice[1]), + primitive(layout, 2, choice[2])}; + reconstruction = apply(reconstruction, from_rotvec(vector), layout.left); + if (vector[0] != 0.0 || vector[1] != 0.0 || vector[2] != 0.0) { + direction = {vector[0], vector[1], vector[2]}; + } + for (int axis = 0; axis < 3; ++axis) cells[3 * step + axis] = choice[axis]; + } +} + +} // namespace + +py::array_t so3_second_order_encode_batch( + const py::array_t& table, + const py::array_t& bins, + const py::array_t& quaternions, + const py::array_t& previous_first_order, + int32_t input_order, bool left, double limit, double radius, bool smooth, int threads) { + if (table.ndim() != 2 || table.shape(0) != 3 || bins.size() != 3 || + quaternions.ndim() != 3 || quaternions.shape(2) != 4 || + previous_first_order.ndim() != 2 || previous_first_order.shape(1) != 4 || + previous_first_order.shape(0) != quaternions.shape(0) || + (input_order != 0 && input_order != 1) || !(limit > 0) || !(radius > 0)) { + throw std::invalid_argument("invalid second-order SO(3) batch arguments"); + } + for (int j = 0; j < 3; ++j) { + if (bins.data()[j] < 1 || bins.data()[j] > table.shape(1)) + throw std::invalid_argument("SO(3) primitive count exceeds table width"); + } + const SO3Layout layout{table.data(), bins.data(), table.shape(1), limit, smooth, + true, left, input_order, radius}; + const int64_t batch = quaternions.shape(0), steps = quaternions.shape(1); + py::array_t output({batch, steps, int64_t{3}}); + const double* values = quaternions.data(); + const double* previous = previous_first_order.data(); + int32_t* result = output.mutable_data(); + const Quaternion identity{0, 0, 0, 1}; + py::gil_scoped_release release; + ac2::parallel_for(batch, threads, [&](int64_t i) { + Quaternion truth = identity, reconstruction = identity; + Quaternion increment{previous[4*i], previous[4*i+1], previous[4*i+2], previous[4*i+3]}; + std::array direction{}; + for (int64_t t = 0; t < steps; ++t) { + const double* raw = values + (i*steps+t)*4; + const Quaternion value{raw[0], raw[1], raw[2], raw[3]}; + truth = input_order == 0 ? value : apply(truth, value, left); + const auto desired = residual(truth, reconstruction, left); + const auto correction = residual(desired, increment, left); + const auto request = as_rotvec(correction); + auto choice = nearest_choice(layout, request); + if (smooth) { + // Python scores U^-1 W for both conventions; use exactly that order. + SO3Layout selection = layout; + selection.left = false; + choice = smooth_choice(selection, request, identity, correction, direction, choice); + } + double vector[3] = {primitive(layout, 0, choice[0]), primitive(layout, 1, choice[1]), primitive(layout, 2, choice[2])}; + project_primitive(vector, radius); + increment = apply(increment, from_rotvec(vector), left); + reconstruction = apply(reconstruction, increment, left); + if (vector[0] != 0 || vector[1] != 0 || vector[2] != 0) + direction = {vector[0], vector[1], vector[2]}; + for (int j = 0; j < 3; ++j) result[(i*steps+t)*3+j] = choice[j]; + } + }); + return output; +} + +py::array_t so3_encode_batch( + const py::array_t& table, + const py::array_t& bins, + const py::array_t& actions, + int32_t input_order, bool rotation_vector, bool left, double limit, bool smooth, + int threads) { + if (table.ndim() != 2 || table.shape(0) != 3 || bins.size() != 3 || actions.ndim() != 3 || + actions.shape(2) != 3 || (input_order != 0 && input_order != 1) || limit <= 0.0) { + throw std::invalid_argument("invalid SO(3) batch encoder arguments"); + } + for (int axis = 0; axis < 3; ++axis) { + if (bins.data()[axis] < 1 || bins.data()[axis] > table.shape(1)) { + throw std::invalid_argument("SO(3) primitive count exceeds table width"); + } + } + const SO3Layout layout{table.data(), bins.data(), table.shape(1), limit, smooth, + rotation_vector, left, input_order}; + const int64_t batch = actions.shape(0); + const int64_t steps = actions.shape(1); + py::array_t output({batch, steps, int64_t{3}}); + { + py::gil_scoped_release release; + ac2::parallel_for(batch, threads, [&](int64_t index) { + encode_trajectory(layout, actions.data() + index * steps * 3, steps, + output.mutable_data() + index * steps * 3); + }); + } + return output; +} diff --git a/config.json b/config.json new file mode 100644 index 0000000000000000000000000000000000000000..fe98b52e0a07faed76f56cff0aee215091a4a2d3 --- /dev/null +++ b/config.json @@ -0,0 +1,4 @@ +{ + "model_type": "actioncodec2", + "processor_class": "ActionCodec2" +} diff --git a/processing_actioncodec2.py b/processing_actioncodec2.py new file mode 100644 index 0000000000000000000000000000000000000000..ca47be24297dc12733d9cdc358f02bb663d03546 --- /dev/null +++ b/processing_actioncodec2.py @@ -0,0 +1,94 @@ +"""Small Hugging Face entry point for a grouped ActionCodec2 artifact. + +Transformers copies only top-level Python files from a Hub repository into its +dynamic-module cache. This entry point locates the artifact, then loads the +versioned runtime stored in ``runtime/`` without installing another package. +""" + +from __future__ import annotations + +import hashlib +import importlib +import importlib.util +import sys +from pathlib import Path + + +class ActionCodec2: + """Load the processor implementation bundled with a pretrained artifact.""" + + @classmethod + def register_for_auto_class(cls, auto_class="AutoProcessor"): + """Satisfy Transformers' dynamic class hook for this thin loader.""" + return cls + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): + source = Path(pretrained_model_name_or_path) + subfolder = str(kwargs.pop("subfolder", "")) + action_space = kwargs.pop("action_space", None) + kwargs.pop("_from_auto", None) + kwargs.pop("trust_remote_code", None) + allowed = { + "cache_dir", + "force_download", + "local_files_only", + "token", + "revision", + "repo_type", + } + unknown = sorted(set(kwargs) - allowed) + if unknown: + raise TypeError( + f"unsupported from_pretrained keyword(s): {', '.join(unknown)}" + ) + if source.is_dir(): + root = source / subfolder + else: + from huggingface_hub import snapshot_download + + prefix = f"{subfolder.rstrip('/')}/" if subfolder else "" + patterns = [ + "config.json", + "processor_config.json", + "router_config.yaml", + "profiles/**", + "runtime/**", + "actioncodec2_kernels/**", + "setup.py", "pyproject.toml", "LICENSE", + "README.md", + "requirements.txt", + "fit_report.json", + ] + snapshot = snapshot_download( + repo_id=str(pretrained_model_name_or_path), + allow_patterns=[prefix + pattern for pattern in patterns], + **kwargs, + ) + root = Path(snapshot) / subfolder + + runtime = root / "runtime" + package_file = runtime / "__init__.py" + if not package_file.is_file(): + raise FileNotFoundError( + f"ActionCodec2 runtime is missing from {root}; copy the whole artifact" + ) + identity = hashlib.sha256(str(runtime.resolve()).encode()).hexdigest()[:16] + package_name = f"_actioncodec2_artifact_{identity}" + if package_name not in sys.modules: + spec = importlib.util.spec_from_file_location( + package_name, package_file, submodule_search_locations=[str(runtime)] + ) + if spec is None or spec.loader is None: + raise ImportError(f"cannot load ActionCodec2 runtime from {runtime}") + package = importlib.util.module_from_spec(spec) + sys.modules[package_name] = package + try: + spec.loader.exec_module(package) + except BaseException: + del sys.modules[package_name] + raise + processor = importlib.import_module(f"{package_name}.processing_actioncodec2") + acceleration = importlib.import_module(f"{package_name}.hub_acceleration") + acceleration.install(package_name) + return processor.ActionCodec2.from_pretrained(root, action_space=action_space) diff --git a/processor_config.json b/processor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..35af7cf9cc192da8937d3b1ce54023bb6ab2ec8a --- /dev/null +++ b/processor_config.json @@ -0,0 +1,10 @@ +{ + "processor_class": "ActionCodec2", + "action_space": "single_eef_delta", + "fitted_profiles": [ + "eef" + ], + "auto_map": { + "AutoProcessor": "processing_actioncodec2.ActionCodec2" + } +} diff --git a/profiles/eef/bpe.json b/profiles/eef/bpe.json new file mode 100644 index 0000000000000000000000000000000000000000..e3152c2be10de521c1dbbd6bd872eabb0b54c605 --- /dev/null +++ b/profiles/eef/bpe.json @@ -0,0 +1,5950 @@ +{ + "spatial_merges": [ + [ + 40, + 121 + ], + [ + 202, + 720 + ], + [ + 338, + 816 + ], + [ + 529, + 818 + ], + [ + 819, + 821 + ], + [ + 338, + 817 + ], + [ + 820, + 822 + ], + [ + 337, + 816 + ], + [ + 339, + 816 + ], + [ + 201, + 720 + ], + [ + 822, + 823 + ], + [ + 41, + 121 + ], + [ + 529, + 819 + ], + [ + 39, + 121 + ], + [ + 40, + 120 + ], + [ + 202, + 721 + ], + [ + 337, + 817 + ], + [ + 40, + 122 + ], + [ + 202, + 719 + ], + [ + 339, + 817 + ], + [ + 203, + 720 + ], + [ + 528, + 818 + ], + [ + 336, + 816 + ], + [ + 340, + 816 + ], + [ + 530, + 818 + ], + [ + 201, + 721 + ], + [ + 820, + 821 + ], + [ + 201, + 719 + ], + [ + 41, + 120 + ], + [ + 41, + 122 + ], + [ + 335, + 816 + ], + [ + 341, + 816 + ], + [ + 819, + 820 + ], + [ + 203, + 719 + ], + [ + 39, + 120 + ], + [ + 39, + 122 + ], + [ + 529, + 820 + ], + [ + 203, + 721 + ], + [ + 336, + 817 + ], + [ + 340, + 817 + ], + [ + 200, + 720 + ], + [ + 821, + 823 + ], + [ + 820, + 830 + ], + [ + 42, + 121 + ], + [ + 334, + 816 + ], + [ + 342, + 816 + ], + [ + 528, + 827 + ], + [ + 819, + 823 + ], + [ + 823, + 830 + ], + [ + 204, + 720 + ], + [ + 529, + 823 + ], + [ + 530, + 820 + ], + [ + 38, + 121 + ], + [ + 202, + 722 + ], + [ + 202, + 718 + ], + [ + 200, + 721 + ], + [ + 200, + 719 + ], + [ + 333, + 816 + ], + [ + 40, + 119 + ], + [ + 343, + 816 + ], + [ + 40, + 123 + ], + [ + 528, + 820 + ], + [ + 335, + 817 + ], + [ + 341, + 817 + ], + [ + 529, + 825 + ], + [ + 529, + 826 + ], + [ + 820, + 839 + ], + [ + 42, + 120 + ], + [ + 531, + 820 + ], + [ + 42, + 122 + ], + [ + 201, + 722 + ], + [ + 527, + 820 + ], + [ + 204, + 719 + ], + [ + 201, + 718 + ], + [ + 530, + 825 + ], + [ + 530, + 826 + ], + [ + 41, + 119 + ], + [ + 204, + 721 + ], + [ + 528, + 819 + ], + [ + 41, + 123 + ], + [ + 38, + 122 + ], + [ + 38, + 120 + ], + [ + 332, + 816 + ], + [ + 820, + 842 + ], + [ + 203, + 718 + ], + [ + 344, + 816 + ], + [ + 39, + 123 + ], + [ + 39, + 119 + ], + [ + 203, + 722 + ], + [ + 832, + 834 + ], + [ + 835, + 837 + ], + [ + 839, + 850 + ], + [ + 43, + 121 + ], + [ + 205, + 720 + ], + [ + 530, + 823 + ], + [ + 199, + 720 + ], + [ + 528, + 829 + ], + [ + 827, + 844 + ], + [ + 334, + 817 + ], + [ + 823, + 839 + ], + [ + 342, + 817 + ], + [ + 37, + 121 + ], + [ + 528, + 825 + ], + [ + 531, + 825 + ], + [ + 531, + 826 + ], + [ + 529, + 833 + ], + [ + 842, + 850 + ], + [ + 200, + 722 + ], + [ + 527, + 818 + ], + [ + 529, + 836 + ], + [ + 202, + 717 + ], + [ + 528, + 826 + ], + [ + 345, + 816 + ], + [ + 202, + 723 + ], + [ + 331, + 816 + ], + [ + 530, + 819 + ], + [ + 200, + 718 + ], + [ + 822, + 825 + ], + [ + 528, + 823 + ], + [ + 822, + 826 + ], + [ + 833, + 844 + ], + [ + 532, + 820 + ], + [ + 836, + 844 + ], + [ + 199, + 721 + ], + [ + 42, + 119 + ], + [ + 40, + 118 + ], + [ + 529, + 827 + ], + [ + 40, + 124 + ], + [ + 205, + 719 + ], + [ + 822, + 834 + ], + [ + 527, + 825 + ], + [ + 42, + 123 + ], + [ + 839, + 865 + ], + [ + 527, + 826 + ], + [ + 204, + 718 + ], + [ + 199, + 719 + ], + [ + 822, + 837 + ], + [ + 530, + 838 + ], + [ + 526, + 820 + ], + [ + 205, + 721 + ], + [ + 38, + 123 + ], + [ + 201, + 717 + ], + [ + 201, + 723 + ], + [ + 825, + 830 + ], + [ + 821, + 825 + ], + [ + 821, + 826 + ], + [ + 826, + 830 + ], + [ + 43, + 120 + ], + [ + 204, + 722 + ], + [ + 38, + 119 + ], + [ + 838, + 844 + ], + [ + 43, + 122 + ], + [ + 827, + 829 + ], + [ + 346, + 816 + ], + [ + 41, + 118 + ], + [ + 330, + 816 + ], + [ + 41, + 124 + ], + [ + 823, + 842 + ], + [ + 831, + 838 + ], + [ + 819, + 839 + ], + [ + 37, + 122 + ], + [ + 37, + 120 + ], + [ + 333, + 817 + ], + [ + 527, + 823 + ], + [ + 343, + 817 + ], + [ + 39, + 118 + ], + [ + 827, + 859 + ], + [ + 531, + 823 + ], + [ + 831, + 860 + ], + [ + 39, + 124 + ], + [ + 530, + 840 + ], + [ + 203, + 717 + ], + [ + 203, + 723 + ], + [ + 829, + 834 + ], + [ + 529, + 840 + ], + [ + 530, + 841 + ], + [ + 833, + 859 + ], + [ + 842, + 865 + ], + [ + 827, + 884 + ], + [ + 836, + 859 + ], + [ + 528, + 840 + ], + [ + 44, + 121 + ], + [ + 206, + 720 + ], + [ + 529, + 841 + ], + [ + 528, + 841 + ], + [ + 819, + 842 + ], + [ + 821, + 834 + ], + [ + 532, + 826 + ], + [ + 532, + 825 + ], + [ + 531, + 840 + ], + [ + 821, + 837 + ], + [ + 531, + 841 + ], + [ + 529, + 837 + ], + [ + 820, + 864 + ], + [ + 829, + 860 + ], + [ + 36, + 121 + ], + [ + 199, + 722 + ], + [ + 832, + 860 + ], + [ + 835, + 860 + ], + [ + 834, + 846 + ], + [ + 202, + 716 + ], + [ + 829, + 866 + ], + [ + 832, + 866 + ], + [ + 347, + 816 + ], + [ + 835, + 866 + ], + [ + 530, + 837 + ] + ], + "temporal_merges": [ + [ + 824, + 824, + 1 + ], + [ + 1024, + 1024, + 2 + ], + [ + 828, + 828, + 1 + ], + [ + 1025, + 1025, + 4 + ], + [ + 816, + 816, + 1 + ], + [ + 817, + 817, + 1 + ], + [ + 831, + 831, + 2 + ], + [ + 527, + 527, + 1 + ], + [ + 829, + 829, + 2 + ], + [ + 1026, + 1026, + 2 + ], + [ + 526, + 526, + 1 + ], + [ + 832, + 832, + 2 + ], + [ + 528, + 528, + 1 + ], + [ + 835, + 835, + 2 + ], + [ + 530, + 530, + 1 + ], + [ + 531, + 531, + 1 + ], + [ + 848, + 848, + 1 + ], + [ + 849, + 849, + 1 + ], + [ + 529, + 529, + 1 + ], + [ + 834, + 834, + 1 + ], + [ + 532, + 532, + 1 + ], + [ + 847, + 847, + 2 + ], + [ + 533, + 533, + 1 + ], + [ + 837, + 837, + 1 + ], + [ + 856, + 856, + 1 + ], + [ + 838, + 838, + 2 + ], + [ + 525, + 525, + 1 + ], + [ + 198, + 198, + 1 + ], + [ + 857, + 857, + 1 + ], + [ + 840, + 840, + 1 + ], + [ + 833, + 833, + 2 + ], + [ + 841, + 841, + 1 + ], + [ + 843, + 843, + 2 + ], + [ + 117, + 117, + 1 + ], + [ + 125, + 125, + 1 + ], + [ + 836, + 836, + 2 + ], + [ + 846, + 846, + 2 + ], + [ + 852, + 852, + 2 + ], + [ + 1024, + 824, + 2 + ], + [ + 853, + 853, + 2 + ], + [ + 825, + 825, + 1 + ], + [ + 845, + 845, + 2 + ], + [ + 826, + 826, + 1 + ], + [ + 820, + 820, + 1 + ], + [ + 818, + 818, + 2 + ], + [ + 200, + 200, + 1 + ], + [ + 206, + 206, + 1 + ], + [ + 202, + 202, + 1 + ], + [ + 37, + 37, + 1 + ], + [ + 861, + 861, + 2 + ], + [ + 724, + 724, + 1 + ], + [ + 205, + 205, + 1 + ], + [ + 43, + 43, + 1 + ], + [ + 862, + 862, + 1 + ], + [ + 827, + 827, + 2 + ], + [ + 42, + 42, + 1 + ], + [ + 118, + 118, + 1 + ], + [ + 863, + 863, + 1 + ], + [ + 38, + 38, + 1 + ], + [ + 1028, + 1028, + 2 + ], + [ + 36, + 36, + 1 + ], + [ + 40, + 40, + 2 + ], + [ + 1029, + 1029, + 2 + ], + [ + 528, + 527, + 1 + ], + [ + 124, + 124, + 1 + ], + [ + 201, + 201, + 1 + ], + [ + 121, + 121, + 1 + ], + [ + 119, + 119, + 1 + ], + [ + 526, + 527, + 1 + ], + [ + 120, + 120, + 1 + ], + [ + 204, + 204, + 1 + ], + [ + 858, + 858, + 2 + ], + [ + 534, + 534, + 1 + ], + [ + 123, + 123, + 1 + ], + [ + 530, + 529, + 1 + ], + [ + 122, + 122, + 1 + ], + [ + 44, + 44, + 1 + ], + [ + 824, + 824, + 2 + ], + [ + 851, + 851, + 2 + ], + [ + 199, + 199, + 1 + ], + [ + 718, + 718, + 1 + ], + [ + 717, + 717, + 1 + ], + [ + 203, + 203, + 1 + ], + [ + 723, + 723, + 1 + ], + [ + 828, + 828, + 2 + ], + [ + 207, + 207, + 1 + ], + [ + 530, + 531, + 1 + ], + [ + 528, + 529, + 1 + ], + [ + 524, + 524, + 1 + ], + [ + 39, + 39, + 1 + ], + [ + 526, + 525, + 1 + ], + [ + 41, + 41, + 1 + ], + [ + 720, + 720, + 1 + ], + [ + 855, + 855, + 2 + ], + [ + 870, + 870, + 2 + ], + [ + 716, + 716, + 1 + ], + [ + 532, + 531, + 1 + ], + [ + 35, + 35, + 1 + ], + [ + 880, + 880, + 1 + ], + [ + 197, + 197, + 1 + ], + [ + 831, + 831, + 3 + ], + [ + 45, + 45, + 1 + ], + [ + 116, + 116, + 1 + ], + [ + 881, + 881, + 1 + ], + [ + 126, + 126, + 1 + ], + [ + 823, + 823, + 1 + ], + [ + 719, + 719, + 1 + ], + [ + 867, + 867, + 2 + ], + [ + 532, + 533, + 1 + ], + [ + 721, + 721, + 1 + ], + [ + 722, + 722, + 1 + ], + [ + 875, + 875, + 1 + ], + [ + 856, + 834, + 1 + ], + [ + 877, + 877, + 1 + ], + [ + 848, + 862, + 1 + ], + [ + 876, + 876, + 2 + ], + [ + 848, + 840, + 1 + ], + [ + 849, + 863, + 1 + ], + [ + 878, + 878, + 2 + ], + [ + 857, + 837, + 1 + ], + [ + 829, + 829, + 3 + ], + [ + 849, + 841, + 1 + ], + [ + 724, + 725, + 1 + ], + [ + 534, + 533, + 1 + ], + [ + 715, + 715, + 1 + ], + [ + 833, + 833, + 1 + ], + [ + 836, + 836, + 1 + ], + [ + 1026, + 828, + 2 + ], + [ + 832, + 832, + 3 + ], + [ + 1033, + 1033, + 4 + ], + [ + 854, + 854, + 2 + ], + [ + 830, + 830, + 2 + ], + [ + 835, + 835, + 3 + ], + [ + 535, + 535, + 1 + ], + [ + 524, + 525, + 1 + ], + [ + 885, + 885, + 2 + ], + [ + 887, + 887, + 2 + ], + [ + 825, + 840, + 1 + ], + [ + 826, + 841, + 1 + ], + [ + 523, + 523, + 1 + ], + [ + 838, + 838, + 3 + ], + [ + 843, + 827, + 2 + ], + [ + 208, + 208, + 1 + ], + [ + 845, + 827, + 2 + ], + [ + 873, + 873, + 2 + ], + [ + 894, + 894, + 2 + ], + [ + 871, + 871, + 1 + ], + [ + 725, + 725, + 1 + ], + [ + 897, + 897, + 2 + ], + [ + 872, + 872, + 1 + ], + [ + 856, + 880, + 1 + ], + [ + 34, + 34, + 1 + ], + [ + 874, + 874, + 1 + ], + [ + 857, + 881, + 1 + ], + [ + 900, + 900, + 1 + ], + [ + 903, + 903, + 1 + ], + [ + 818, + 818, + 1 + ], + [ + 898, + 898, + 2 + ], + [ + 899, + 899, + 2 + ], + [ + 206, + 205, + 1 + ], + [ + 1031, + 527, + 2 + ], + [ + 831, + 831, + 1 + ], + [ + 829, + 829, + 1 + ], + [ + 825, + 820, + 1 + ], + [ + 826, + 820, + 1 + ], + [ + 837, + 837, + 2 + ], + [ + 869, + 869, + 2 + ], + [ + 875, + 862, + 1 + ], + [ + 202, + 201, + 1 + ], + [ + 534, + 535, + 1 + ], + [ + 206, + 207, + 1 + ], + [ + 838, + 838, + 1 + ], + [ + 127, + 127, + 1 + ], + [ + 877, + 863, + 1 + ], + [ + 905, + 905, + 2 + ], + [ + 904, + 904, + 2 + ], + [ + 46, + 46, + 1 + ], + [ + 115, + 115, + 1 + ], + [ + 717, + 716, + 1 + ], + [ + 910, + 910, + 2 + ], + [ + 833, + 833, + 3 + ], + [ + 916, + 916, + 1 + ], + [ + 843, + 843, + 1 + ], + [ + 724, + 723, + 1 + ], + [ + 834, + 834, + 2 + ], + [ + 717, + 718, + 1 + ], + [ + 200, + 199, + 1 + ], + [ + 836, + 836, + 3 + ], + [ + 886, + 886, + 1 + ], + [ + 868, + 868, + 2 + ], + [ + 918, + 918, + 1 + ], + [ + 197, + 198, + 1 + ], + [ + 196, + 196, + 1 + ], + [ + 845, + 845, + 1 + ], + [ + 117, + 116, + 1 + ], + [ + 851, + 851, + 1 + ], + [ + 715, + 714, + 1 + ], + [ + 890, + 890, + 1 + ], + [ + 889, + 889, + 1 + ], + [ + 715, + 716, + 1 + ], + [ + 37, + 36, + 1 + ], + [ + 125, + 126, + 1 + ], + [ + 524, + 523, + 1 + ], + [ + 847, + 835, + 2 + ], + [ + 37, + 38, + 1 + ], + [ + 43, + 42, + 1 + ], + [ + 200, + 201, + 1 + ], + [ + 536, + 536, + 1 + ], + [ + 726, + 726, + 1 + ], + [ + 204, + 203, + 1 + ], + [ + 858, + 858, + 1 + ], + [ + 855, + 855, + 1 + ], + [ + 531, + 530, + 1 + ], + [ + 529, + 530, + 1 + ], + [ + 529, + 528, + 1 + ], + [ + 117, + 118, + 1 + ], + [ + 202, + 203, + 1 + ], + [ + 204, + 205, + 1 + ], + [ + 35, + 36, + 1 + ], + [ + 125, + 124, + 1 + ], + [ + 824, + 1024, + 2 + ], + [ + 895, + 895, + 1 + ], + [ + 832, + 832, + 1 + ], + [ + 121, + 120, + 1 + ], + [ + 43, + 44, + 1 + ], + [ + 119, + 118, + 1 + ], + [ + 888, + 888, + 1 + ], + [ + 527, + 528, + 1 + ], + [ + 522, + 522, + 1 + ], + [ + 121, + 122, + 1 + ], + [ + 853, + 835, + 2 + ], + [ + 831, + 831, + 4 + ], + [ + 847, + 829, + 2 + ], + [ + 891, + 891, + 1 + ], + [ + 526, + 1050, + 1 + ], + [ + 123, + 124, + 1 + ], + [ + 45, + 44, + 1 + ], + [ + 719, + 718, + 1 + ], + [ + 911, + 911, + 1 + ], + [ + 722, + 723, + 1 + ], + [ + 1051, + 198, + 2 + ], + [ + 919, + 919, + 2 + ], + [ + 39, + 38, + 1 + ], + [ + 824, + 824, + 3 + ], + [ + 846, + 846, + 3 + ], + [ + 119, + 120, + 1 + ], + [ + 41, + 42, + 1 + ], + [ + 930, + 930, + 1 + ], + [ + 932, + 932, + 1 + ], + [ + 864, + 864, + 2 + ], + [ + 1039, + 531, + 2 + ], + [ + 720, + 721, + 1 + ], + [ + 913, + 913, + 1 + ], + [ + 875, + 900, + 1 + ], + [ + 123, + 122, + 1 + ], + [ + 877, + 903, + 1 + ], + [ + 1028, + 816, + 2 + ], + [ + 852, + 852, + 3 + ], + [ + 714, + 714, + 1 + ], + [ + 867, + 867, + 1 + ], + [ + 1037, + 835, + 4 + ], + [ + 1046, + 533, + 2 + ], + [ + 208, + 207, + 1 + ], + [ + 209, + 209, + 1 + ], + [ + 907, + 907, + 2 + ], + [ + 879, + 879, + 1 + ], + [ + 722, + 721, + 1 + ], + [ + 33, + 33, + 1 + ], + [ + 854, + 854, + 1 + ], + [ + 199, + 198, + 1 + ], + [ + 526, + 1031, + 1 + ], + [ + 908, + 908, + 2 + ], + [ + 828, + 828, + 3 + ], + [ + 1086, + 1086, + 4 + ], + [ + 1030, + 831, + 4 + ], + [ + 827, + 843, + 2 + ], + [ + 823, + 837, + 1 + ], + [ + 726, + 725, + 1 + ], + [ + 892, + 892, + 1 + ], + [ + 833, + 833, + 4 + ], + [ + 861, + 861, + 1 + ], + [ + 1043, + 834, + 2 + ], + [ + 943, + 943, + 2 + ], + [ + 882, + 882, + 2 + ], + [ + 827, + 845, + 2 + ], + [ + 893, + 893, + 1 + ], + [ + 902, + 902, + 1 + ], + [ + 836, + 836, + 4 + ], + [ + 832, + 832, + 4 + ], + [ + 848, + 1077, + 1 + ], + [ + 838, + 838, + 4 + ], + [ + 1025, + 1024, + 4 + ], + [ + 35, + 34, + 1 + ], + [ + 945, + 945, + 2 + ], + [ + 536, + 535, + 1 + ], + [ + 916, + 880, + 1 + ], + [ + 883, + 883, + 2 + ], + [ + 720, + 719, + 1 + ], + [ + 531, + 532, + 1 + ], + [ + 849, + 1081, + 1 + ], + [ + 128, + 128, + 1 + ], + [ + 906, + 906, + 1 + ], + [ + 47, + 47, + 1 + ], + [ + 918, + 881, + 1 + ], + [ + 40, + 40, + 1 + ], + [ + 823, + 834, + 1 + ], + [ + 1025, + 824, + 4 + ], + [ + 942, + 942, + 2 + ], + [ + 847, + 847, + 3 + ], + [ + 536, + 537, + 1 + ], + [ + 873, + 873, + 1 + ], + [ + 528, + 528, + 2 + ], + [ + 846, + 846, + 1 + ], + [ + 527, + 526, + 1 + ], + [ + 1053, + 840, + 2 + ], + [ + 522, + 523, + 1 + ], + [ + 828, + 1026, + 2 + ], + [ + 525, + 526, + 1 + ], + [ + 1055, + 841, + 2 + ], + [ + 949, + 949, + 2 + ], + [ + 853, + 853, + 3 + ], + [ + 114, + 114, + 1 + ], + [ + 726, + 727, + 1 + ], + [ + 1067, + 820, + 2 + ], + [ + 1047, + 837, + 2 + ], + [ + 980, + 980, + 1 + ], + [ + 852, + 852, + 1 + ], + [ + 825, + 825, + 2 + ], + [ + 830, + 830, + 1 + ], + [ + 870, + 870, + 1 + ], + [ + 1083, + 1083, + 4 + ], + [ + 874, + 874, + 2 + ], + [ + 1032, + 829, + 4 + ], + [ + 826, + 826, + 2 + ], + [ + 982, + 982, + 1 + ], + [ + 713, + 713, + 1 + ], + [ + 869, + 869, + 1 + ], + [ + 1029, + 817, + 2 + ], + [ + 528, + 1031, + 1 + ], + [ + 1042, + 529, + 2 + ], + [ + 864, + 827, + 2 + ], + [ + 819, + 819, + 2 + ], + [ + 115, + 116, + 1 + ], + [ + 127, + 126, + 1 + ], + [ + 1079, + 42, + 2 + ], + [ + 524, + 1050, + 1 + ], + [ + 878, + 835, + 2 + ], + [ + 45, + 46, + 1 + ], + [ + 1082, + 38, + 2 + ], + [ + 1080, + 118, + 2 + ], + [ + 843, + 843, + 3 + ], + [ + 871, + 871, + 2 + ], + [ + 533, + 532, + 1 + ], + [ + 839, + 839, + 2 + ], + [ + 958, + 958, + 2 + ], + [ + 522, + 521, + 1 + ], + [ + 872, + 872, + 2 + ], + [ + 1084, + 36, + 2 + ], + [ + 838, + 851, + 2 + ], + [ + 197, + 196, + 1 + ], + [ + 1088, + 124, + 2 + ], + [ + 1075, + 205, + 2 + ], + [ + 971, + 971, + 1 + ], + [ + 1093, + 120, + 2 + ], + [ + 1035, + 832, + 4 + ], + [ + 195, + 195, + 1 + ], + [ + 973, + 973, + 1 + ], + [ + 856, + 1043, + 1 + ], + [ + 965, + 965, + 2 + ], + [ + 915, + 909, + 1 + ], + [ + 829, + 829, + 4 + ], + [ + 533, + 534, + 1 + ], + [ + 858, + 827, + 2 + ], + [ + 1030, + 1030, + 4 + ], + [ + 967, + 967, + 2 + ], + [ + 1034, + 527, + 2 + ], + [ + 930, + 903, + 1 + ], + [ + 838, + 855, + 2 + ], + [ + 854, + 854, + 3 + ], + [ + 932, + 900, + 1 + ], + [ + 1099, + 122, + 2 + ], + [ + 972, + 972, + 2 + ], + [ + 1089, + 201, + 2 + ], + [ + 1038, + 530, + 2 + ], + [ + 969, + 969, + 2 + ], + [ + 1085, + 41, + 1 + ], + [ + 845, + 845, + 3 + ], + [ + 925, + 925, + 1 + ], + [ + 974, + 974, + 2 + ], + [ + 837, + 857, + 1 + ], + [ + 912, + 912, + 2 + ], + [ + 901, + 901, + 2 + ], + [ + 853, + 853, + 1 + ], + [ + 928, + 928, + 1 + ], + [ + 841, + 826, + 1 + ], + [ + 206, + 1109, + 1 + ], + [ + 1100, + 44, + 2 + ], + [ + 117, + 1126, + 1 + ], + [ + 840, + 825, + 1 + ], + [ + 856, + 1122, + 1 + ], + [ + 840, + 848, + 1 + ], + [ + 1085, + 39, + 1 + ], + [ + 1057, + 1057, + 2 + ], + [ + 125, + 1128, + 1 + ], + [ + 537, + 537, + 1 + ], + [ + 868, + 834, + 1 + ], + [ + 1045, + 847, + 4 + ], + [ + 1058, + 1058, + 2 + ], + [ + 931, + 931, + 1 + ], + [ + 983, + 983, + 2 + ], + [ + 532, + 1039, + 1 + ], + [ + 841, + 849, + 1 + ], + [ + 1034, + 525, + 2 + ], + [ + 197, + 1051, + 1 + ], + [ + 530, + 530, + 2 + ], + [ + 833, + 833, + 5 + ], + [ + 1024, + 824, + 3 + ], + [ + 934, + 934, + 1 + ], + [ + 979, + 979, + 2 + ], + [ + 1025, + 1062, + 4 + ], + [ + 939, + 939, + 1 + ], + [ + 329, + 329, + 1 + ], + [ + 836, + 836, + 5 + ], + [ + 1106, + 203, + 2 + ], + [ + 941, + 941, + 1 + ], + [ + 987, + 987, + 2 + ], + [ + 890, + 890, + 2 + ], + [ + 889, + 879, + 1 + ], + [ + 876, + 876, + 1 + ], + [ + 1040, + 862, + 2 + ], + [ + 728, + 727, + 1 + ], + [ + 978, + 978, + 2 + ], + [ + 32, + 32, + 1 + ], + [ + 1036, + 528, + 2 + ], + [ + 818, + 818, + 3 + ], + [ + 208, + 209, + 1 + ], + [ + 1090, + 1090, + 2 + ], + [ + 922, + 922, + 1 + ], + [ + 921, + 921, + 1 + ], + [ + 850, + 850, + 2 + ], + [ + 1041, + 863, + 2 + ], + [ + 946, + 946, + 1 + ], + [ + 999, + 999, + 1 + ], + [ + 1036, + 527, + 2 + ], + [ + 917, + 917, + 2 + ], + [ + 39, + 40, + 1 + ], + [ + 41, + 40, + 1 + ], + [ + 520, + 521, + 1 + ], + [ + 851, + 851, + 3 + ], + [ + 725, + 724, + 1 + ], + [ + 1034, + 526, + 2 + ], + [ + 713, + 714, + 1 + ], + [ + 210, + 210, + 1 + ], + [ + 1104, + 718, + 2 + ], + [ + 847, + 887, + 2 + ], + [ + 1048, + 834, + 2 + ], + [ + 868, + 837, + 1 + ], + [ + 1040, + 848, + 2 + ], + [ + 857, + 1127, + 1 + ], + [ + 858, + 873, + 2 + ], + [ + 200, + 1103, + 1 + ], + [ + 895, + 895, + 2 + ], + [ + 953, + 953, + 1 + ], + [ + 875, + 1077, + 1 + ], + [ + 957, + 957, + 1 + ], + [ + 1041, + 849, + 2 + ], + [ + 855, + 838, + 2 + ], + [ + 530, + 1039, + 1 + ], + [ + 853, + 831, + 2 + ], + [ + 524, + 526, + 1 + ], + [ + 830, + 830, + 3 + ], + [ + 956, + 956, + 1 + ], + [ + 842, + 842, + 2 + ], + [ + 48, + 48, + 1 + ], + [ + 877, + 1081, + 1 + ], + [ + 888, + 888, + 2 + ], + [ + 952, + 952, + 1 + ], + [ + 827, + 827, + 3 + ], + [ + 533, + 535, + 1 + ], + [ + 1038, + 529, + 2 + ], + [ + 846, + 829, + 2 + ], + [ + 885, + 885, + 1 + ], + [ + 1073, + 861, + 4 + ], + [ + 886, + 886, + 2 + ], + [ + 847, + 897, + 2 + ], + [ + 1071, + 1071, + 2 + ], + [ + 716, + 717, + 1 + ], + [ + 1107, + 723, + 2 + ], + [ + 896, + 896, + 2 + ], + [ + 882, + 882, + 1 + ], + [ + 538, + 538, + 1 + ], + [ + 852, + 831, + 2 + ], + [ + 532, + 1046, + 1 + ], + [ + 891, + 891, + 2 + ], + [ + 1026, + 828, + 3 + ], + [ + 916, + 980, + 1 + ], + [ + 857, + 1047, + 1 + ], + [ + 33, + 34, + 1 + ], + [ + 909, + 915, + 1 + ], + [ + 862, + 848, + 1 + ], + [ + 1072, + 1072, + 2 + ], + [ + 883, + 883, + 1 + ], + [ + 918, + 982, + 1 + ], + [ + 1044, + 531, + 2 + ], + [ + 1052, + 1052, + 2 + ], + [ + 908, + 835, + 2 + ], + [ + 913, + 198, + 1 + ], + [ + 713, + 712, + 1 + ], + [ + 1000, + 1000, + 1 + ], + [ + 863, + 849, + 1 + ], + [ + 846, + 832, + 2 + ], + [ + 930, + 971, + 1 + ], + [ + 1085, + 40, + 1 + ], + [ + 836, + 838, + 1 + ], + [ + 37, + 1084, + 1 + ], + [ + 1048, + 880, + 2 + ], + [ + 868, + 868, + 1 + ], + [ + 932, + 973, + 1 + ], + [ + 530, + 1042, + 1 + ], + [ + 534, + 1046, + 1 + ], + [ + 1036, + 529, + 2 + ], + [ + 892, + 892, + 2 + ], + [ + 878, + 897, + 2 + ], + [ + 858, + 874, + 2 + ], + [ + 818, + 835, + 2 + ], + [ + 129, + 129, + 1 + ], + [ + 894, + 894, + 1 + ], + [ + 1076, + 1076, + 2 + ], + [ + 893, + 893, + 2 + ], + [ + 852, + 832, + 2 + ], + [ + 113, + 113, + 1 + ], + [ + 727, + 727, + 1 + ], + [ + 966, + 966, + 1 + ], + [ + 1013, + 1013, + 1 + ], + [ + 838, + 836, + 1 + ], + [ + 843, + 833, + 2 + ], + [ + 833, + 843, + 1 + ], + [ + 825, + 825, + 3 + ], + [ + 1021, + 1021, + 1 + ], + [ + 1038, + 531, + 2 + ], + [ + 332, + 332, + 1 + ], + [ + 854, + 854, + 4 + ], + [ + 47, + 46, + 1 + ], + [ + 527, + 525, + 1 + ], + [ + 202, + 1089, + 1 + ], + [ + 1035, + 1035, + 4 + ], + [ + 855, + 855, + 3 + ], + [ + 344, + 344, + 1 + ], + [ + 826, + 826, + 3 + ], + [ + 43, + 1079, + 1 + ], + [ + 898, + 898, + 1 + ], + [ + 907, + 834, + 1 + ], + [ + 914, + 914, + 2 + ], + [ + 521, + 521, + 1 + ], + [ + 926, + 926, + 2 + ], + [ + 528, + 1042, + 1 + ], + [ + 1032, + 1032, + 4 + ], + [ + 836, + 845, + 1 + ], + [ + 920, + 920, + 1 + ], + [ + 843, + 833, + 1 + ], + [ + 837, + 868, + 1 + ], + [ + 1040, + 840, + 2 + ], + [ + 1091, + 1091, + 2 + ], + [ + 960, + 960, + 1 + ], + [ + 208, + 1109, + 1 + ], + [ + 959, + 959, + 1 + ], + [ + 199, + 200, + 1 + ], + [ + 115, + 114, + 1 + ], + [ + 716, + 715, + 1 + ], + [ + 879, + 879, + 2 + ], + [ + 899, + 899, + 1 + ], + [ + 206, + 1075, + 1 + ], + [ + 985, + 985, + 1 + ], + [ + 861, + 829, + 2 + ], + [ + 948, + 948, + 1 + ], + [ + 718, + 717, + 1 + ], + [ + 127, + 128, + 1 + ], + [ + 875, + 848, + 1 + ], + [ + 724, + 931, + 1 + ], + [ + 1085, + 1085, + 1 + ], + [ + 845, + 836, + 1 + ], + [ + 1041, + 841, + 2 + ], + [ + 845, + 836, + 2 + ], + [ + 1037, + 1037, + 4 + ], + [ + 929, + 929, + 1 + ], + [ + 904, + 904, + 1 + ], + [ + 877, + 849, + 1 + ], + [ + 37, + 1082, + 1 + ], + [ + 119, + 1080, + 1 + ], + [ + 831, + 831, + 5 + ], + [ + 905, + 905, + 1 + ], + [ + 838, + 855, + 1 + ], + [ + 910, + 910, + 1 + ], + [ + 951, + 951, + 1 + ], + [ + 1050, + 525, + 2 + ], + [ + 827, + 864, + 2 + ], + [ + 824, + 909, + 1 + ], + [ + 837, + 823, + 1 + ], + [ + 1133, + 721, + 2 + ], + [ + 1097, + 1097, + 2 + ], + [ + 329, + 328, + 1 + ], + [ + 819, + 819, + 1 + ], + [ + 204, + 1106, + 1 + ], + [ + 723, + 724, + 1 + ], + [ + 835, + 847, + 2 + ], + [ + 861, + 887, + 2 + ], + [ + 195, + 196, + 1 + ], + [ + 838, + 851, + 1 + ], + [ + 834, + 868, + 1 + ], + [ + 348, + 348, + 1 + ], + [ + 1113, + 1113, + 2 + ], + [ + 911, + 911, + 2 + ], + [ + 902, + 902, + 2 + ], + [ + 1116, + 1116, + 2 + ], + [ + 1115, + 1115, + 2 + ], + [ + 117, + 1080, + 1 + ], + [ + 534, + 532, + 1 + ], + [ + 838, + 833, + 1 + ], + [ + 529, + 529, + 2 + ], + [ + 923, + 923, + 2 + ], + [ + 35, + 1175, + 1 + ], + [ + 1052, + 881, + 2 + ], + [ + 1069, + 1069, + 2 + ], + [ + 728, + 729, + 1 + ], + [ + 872, + 928, + 1 + ], + [ + 889, + 956, + 1 + ], + [ + 1044, + 532, + 2 + ], + [ + 818, + 829, + 2 + ], + [ + 121, + 1093, + 1 + ], + [ + 125, + 1088, + 1 + ], + [ + 123, + 1088, + 1 + ], + [ + 45, + 1100, + 1 + ], + [ + 39, + 1082, + 1 + ], + [ + 878, + 945, + 2 + ], + [ + 824, + 924, + 1 + ], + [ + 927, + 927, + 2 + ], + [ + 906, + 906, + 2 + ], + [ + 886, + 869, + 1 + ], + [ + 1070, + 205, + 2 + ], + [ + 719, + 720, + 1 + ], + [ + 871, + 931, + 1 + ], + [ + 43, + 1100, + 1 + ], + [ + 818, + 831, + 2 + ], + [ + 871, + 833, + 1 + ], + [ + 1112, + 525, + 2 + ], + [ + 1123, + 198, + 2 + ], + [ + 851, + 838, + 2 + ], + [ + 1048, + 1048, + 2 + ], + [ + 31, + 31, + 1 + ], + [ + 968, + 924, + 1 + ], + [ + 121, + 1099, + 1 + ], + [ + 843, + 864, + 2 + ], + [ + 914, + 829, + 2 + ], + [ + 1083, + 816, + 4 + ], + [ + 1051, + 1051, + 2 + ], + [ + 524, + 1163, + 1 + ], + [ + 726, + 724, + 1 + ], + [ + 1033, + 1026, + 4 + ], + [ + 119, + 1093, + 1 + ], + [ + 206, + 1000, + 1 + ], + [ + 890, + 867, + 1 + ], + [ + 711, + 712, + 1 + ], + [ + 1063, + 853, + 4 + ], + [ + 1064, + 825, + 2 + ], + [ + 538, + 537, + 1 + ], + [ + 827, + 833, + 1 + ], + [ + 533, + 1096, + 1 + ], + [ + 520, + 519, + 1 + ], + [ + 941, + 198, + 1 + ], + [ + 127, + 1128, + 1 + ], + [ + 936, + 936, + 2 + ], + [ + 820, + 825, + 1 + ], + [ + 820, + 826, + 1 + ], + [ + 823, + 981, + 1 + ], + [ + 1034, + 1034, + 2 + ], + [ + 1066, + 826, + 2 + ], + [ + 835, + 853, + 2 + ], + [ + 115, + 1126, + 1 + ], + [ + 856, + 1017, + 1 + ], + [ + 123, + 1099, + 1 + ], + [ + 845, + 864, + 2 + ], + [ + 194, + 194, + 1 + ], + [ + 833, + 827, + 1 + ], + [ + 919, + 919, + 1 + ], + [ + 856, + 907, + 1 + ], + [ + 869, + 886, + 1 + ], + [ + 872, + 836, + 1 + ], + [ + 870, + 870, + 3 + ], + [ + 858, + 873, + 1 + ], + [ + 831, + 832, + 1 + ], + [ + 1061, + 852, + 4 + ], + [ + 832, + 831, + 1 + ], + [ + 1033, + 828, + 4 + ], + [ + 534, + 1157, + 1 + ], + [ + 989, + 989, + 1 + ], + [ + 35, + 1084, + 1 + ], + [ + 41, + 1079, + 1 + ], + [ + 523, + 525, + 1 + ], + [ + 990, + 990, + 1 + ], + [ + 1045, + 1045, + 4 + ], + [ + 895, + 867, + 1 + ], + [ + 884, + 884, + 2 + ], + [ + 1036, + 1036, + 2 + ], + [ + 204, + 1075, + 1 + ], + [ + 858, + 874, + 1 + ], + [ + 716, + 714, + 1 + ], + [ + 916, + 1122, + 1 + ], + [ + 871, + 888, + 1 + ], + [ + 824, + 854, + 1 + ], + [ + 1060, + 846, + 4 + ], + [ + 835, + 835, + 4 + ], + [ + 210, + 209, + 1 + ], + [ + 717, + 1119, + 1 + ], + [ + 872, + 891, + 1 + ], + [ + 912, + 912, + 1 + ], + [ + 1094, + 1094, + 2 + ], + [ + 851, + 838, + 1 + ], + [ + 824, + 938, + 1 + ], + [ + 33, + 32, + 1 + ], + [ + 843, + 827, + 1 + ], + [ + 538, + 539, + 1 + ], + [ + 1067, + 1067, + 2 + ], + [ + 824, + 935, + 1 + ], + [ + 913, + 953, + 1 + ], + [ + 925, + 873, + 1 + ], + [ + 839, + 839, + 1 + ], + [ + 49, + 49, + 1 + ], + [ + 824, + 937, + 1 + ], + [ + 879, + 889, + 1 + ], + [ + 829, + 847, + 2 + ], + [ + 827, + 827, + 1 + ], + [ + 533, + 531, + 1 + ], + [ + 832, + 846, + 2 + ], + [ + 824, + 940, + 1 + ], + [ + 918, + 1127, + 1 + ], + [ + 867, + 890, + 1 + ], + [ + 969, + 887, + 2 + ], + [ + 840, + 862, + 1 + ], + [ + 841, + 863, + 1 + ], + [ + 949, + 887, + 2 + ], + [ + 1040, + 1040, + 2 + ], + [ + 824, + 915, + 1 + ], + [ + 831, + 853, + 2 + ], + [ + 1070, + 207, + 2 + ], + [ + 210, + 211, + 1 + ], + [ + 934, + 874, + 1 + ], + [ + 851, + 836, + 1 + ], + [ + 1057, + 116, + 2 + ], + [ + 875, + 1178, + 1 + ], + [ + 206, + 946, + 1 + ], + [ + 1013, + 36, + 1 + ], + [ + 117, + 943, + 1 + ], + [ + 1069, + 200, + 2 + ], + [ + 999, + 44, + 1 + ], + [ + 827, + 843, + 1 + ], + [ + 833, + 838, + 1 + ], + [ + 955, + 955, + 2 + ], + [ + 1058, + 126, + 2 + ], + [ + 1121, + 1121, + 2 + ], + [ + 718, + 719, + 1 + ], + [ + 984, + 950, + 1 + ], + [ + 877, + 1179, + 1 + ], + [ + 939, + 533, + 1 + ], + [ + 327, + 328, + 1 + ], + [ + 1021, + 971, + 1 + ], + [ + 908, + 837, + 1 + ], + [ + 1052, + 837, + 2 + ], + [ + 201, + 200, + 1 + ], + [ + 867, + 895, + 1 + ], + [ + 855, + 838, + 1 + ], + [ + 830, + 830, + 4 + ], + [ + 974, + 897, + 2 + ], + [ + 1041, + 1041, + 2 + ], + [ + 857, + 1010, + 1 + ], + [ + 1103, + 199, + 2 + ], + [ + 197, + 1216, + 1 + ], + [ + 1014, + 1014, + 1 + ], + [ + 861, + 885, + 2 + ], + [ + 526, + 956, + 1 + ], + [ + 1026, + 1108, + 2 + ], + [ + 1074, + 724, + 2 + ], + [ + 878, + 904, + 2 + ], + [ + 985, + 531, + 1 + ], + [ + 125, + 945, + 1 + ], + [ + 855, + 833, + 1 + ], + [ + 45, + 1200, + 1 + ], + [ + 981, + 981, + 1 + ], + [ + 828, + 828, + 4 + ], + [ + 329, + 973, + 1 + ], + [ + 876, + 876, + 3 + ], + [ + 836, + 851, + 1 + ], + [ + 961, + 961, + 2 + ], + [ + 848, + 1053, + 1 + ], + [ + 913, + 941, + 1 + ], + [ + 1024, + 1024, + 3 + ], + [ + 1038, + 1038, + 2 + ], + [ + 849, + 1055, + 1 + ], + [ + 901, + 901, + 1 + ], + [ + 202, + 1106, + 1 + ], + [ + 1160, + 887, + 4 + ], + [ + 1044, + 1044, + 2 + ], + [ + 964, + 964, + 2 + ], + [ + 528, + 528, + 3 + ], + [ + 527, + 527, + 2 + ], + [ + 933, + 933, + 2 + ], + [ + 827, + 836, + 1 + ], + [ + 834, + 823, + 1 + ], + [ + 848, + 988, + 1 + ], + [ + 832, + 852, + 2 + ], + [ + 825, + 854, + 1 + ], + [ + 833, + 855, + 1 + ], + [ + 827, + 845, + 1 + ], + [ + 998, + 840, + 1 + ], + [ + 1109, + 207, + 2 + ], + [ + 332, + 980, + 1 + ], + [ + 1125, + 44, + 2 + ], + [ + 848, + 992, + 1 + ], + [ + 849, + 993, + 1 + ], + [ + 872, + 902, + 1 + ], + [ + 826, + 854, + 1 + ], + [ + 344, + 982, + 1 + ], + [ + 1043, + 1043, + 2 + ], + [ + 532, + 530, + 1 + ], + [ + 831, + 835, + 2 + ], + [ + 712, + 712, + 1 + ], + [ + 1009, + 1009, + 1 + ], + [ + 1071, + 201, + 2 + ], + [ + 1007, + 1007, + 1 + ], + [ + 1072, + 36, + 2 + ], + [ + 206, + 957, + 1 + ], + [ + 845, + 827, + 1 + ], + [ + 722, + 1107, + 1 + ], + [ + 952, + 205, + 1 + ], + [ + 976, + 976, + 2 + ], + [ + 1002, + 841, + 1 + ], + [ + 949, + 897, + 2 + ], + [ + 1031, + 1031, + 2 + ], + [ + 857, + 1023, + 1 + ], + [ + 1005, + 1005, + 1 + ], + [ + 1006, + 1006, + 1 + ], + [ + 871, + 906, + 1 + ], + [ + 889, + 948, + 1 + ], + [ + 831, + 852, + 2 + ], + [ + 130, + 130, + 1 + ], + [ + 117, + 972, + 1 + ], + [ + 1018, + 1018, + 1 + ], + [ + 829, + 846, + 2 + ], + [ + 953, + 198, + 1 + ], + [ + 529, + 531, + 1 + ], + [ + 892, + 882, + 1 + ], + [ + 836, + 827, + 1 + ], + [ + 889, + 951, + 1 + ], + [ + 1068, + 818, + 1 + ], + [ + 925, + 723, + 1 + ], + [ + 348, + 349, + 1 + ], + [ + 849, + 1001, + 1 + ], + [ + 1076, + 42, + 2 + ], + [ + 960, + 888, + 1 + ], + [ + 834, + 856, + 1 + ], + [ + 869, + 854, + 1 + ], + [ + 345, + 345, + 1 + ], + [ + 891, + 872, + 1 + ], + [ + 842, + 842, + 1 + ], + [ + 893, + 883, + 1 + ], + [ + 1074, + 723, + 2 + ], + [ + 1086, + 817, + 4 + ], + [ + 839, + 839, + 3 + ], + [ + 1083, + 1028, + 4 + ], + [ + 47, + 48, + 1 + ], + [ + 1070, + 1070, + 2 + ], + [ + 886, + 922, + 1 + ], + [ + 527, + 1034, + 1 + ], + [ + 837, + 908, + 1 + ], + [ + 331, + 331, + 1 + ], + [ + 528, + 864, + 1 + ], + [ + 1090, + 120, + 2 + ], + [ + 715, + 1119, + 1 + ], + [ + 913, + 858, + 1 + ], + [ + 942, + 942, + 1 + ], + [ + 833, + 843, + 2 + ], + [ + 834, + 907, + 1 + ], + [ + 823, + 823, + 2 + ], + [ + 725, + 723, + 1 + ], + [ + 886, + 921, + 1 + ], + [ + 925, + 888, + 1 + ], + [ + 1135, + 862, + 2 + ], + [ + 844, + 844, + 2 + ], + [ + 892, + 921, + 1 + ], + [ + 728, + 728, + 1 + ], + [ + 835, + 829, + 1 + ], + [ + 835, + 831, + 1 + ], + [ + 1061, + 1061, + 4 + ], + [ + 1072, + 38, + 2 + ], + [ + 725, + 727, + 1 + ], + [ + 1130, + 719, + 2 + ], + [ + 886, + 939, + 1 + ], + [ + 1077, + 862, + 2 + ], + [ + 861, + 910, + 2 + ], + [ + 715, + 717, + 1 + ], + [ + 874, + 858, + 1 + ], + [ + 889, + 826, + 1 + ], + [ + 524, + 522, + 1 + ], + [ + 893, + 922, + 1 + ], + [ + 889, + 825, + 1 + ], + [ + 1105, + 717, + 2 + ], + [ + 112, + 112, + 1 + ], + [ + 1064, + 840, + 2 + ], + [ + 1137, + 863, + 2 + ], + [ + 1131, + 867, + 1 + ], + [ + 540, + 539, + 1 + ], + [ + 966, + 205, + 1 + ], + [ + 1039, + 1039, + 2 + ], + [ + 520, + 520, + 1 + ], + [ + 129, + 128, + 1 + ], + [ + 970, + 970, + 2 + ], + [ + 722, + 1133, + 1 + ], + [ + 829, + 835, + 2 + ], + [ + 117, + 983, + 1 + ], + [ + 857, + 908, + 1 + ], + [ + 1057, + 118, + 2 + ], + [ + 1071, + 202, + 2 + ], + [ + 1063, + 1063, + 4 + ], + [ + 891, + 845, + 1 + ], + [ + 991, + 991, + 2 + ], + [ + 1081, + 863, + 2 + ], + [ + 1090, + 122, + 2 + ], + [ + 206, + 911, + 1 + ], + [ + 988, + 988, + 1 + ], + [ + 530, + 529, + 2 + ], + [ + 113, + 114, + 1 + ], + [ + 526, + 528, + 1 + ], + [ + 525, + 523, + 1 + ], + [ + 1058, + 124, + 2 + ], + [ + 879, + 854, + 1 + ], + [ + 212, + 211, + 1 + ], + [ + 1166, + 207, + 2 + ], + [ + 1060, + 1060, + 4 + ], + [ + 125, + 987, + 1 + ], + [ + 878, + 1037, + 2 + ], + [ + 1091, + 118, + 2 + ], + [ + 1094, + 203, + 2 + ], + [ + 818, + 832, + 2 + ], + [ + 861, + 1032, + 2 + ], + [ + 847, + 1032, + 2 + ], + [ + 1066, + 841, + 2 + ], + [ + 1172, + 897, + 4 + ], + [ + 530, + 528, + 1 + ], + [ + 928, + 872, + 1 + ], + [ + 846, + 885, + 2 + ], + [ + 882, + 892, + 1 + ], + [ + 1073, + 1073, + 4 + ], + [ + 975, + 975, + 2 + ], + [ + 831, + 835, + 1 + ], + [ + 888, + 871, + 1 + ], + [ + 1044, + 533, + 2 + ], + [ + 1087, + 528, + 2 + ], + [ + 1020, + 947, + 1 + ], + [ + 825, + 1053, + 1 + ], + [ + 890, + 946, + 1 + ], + [ + 1030, + 1030, + 3 + ], + [ + 1118, + 870, + 4 + ], + [ + 534, + 536, + 1 + ], + [ + 882, + 840, + 1 + ], + [ + 883, + 893, + 1 + ], + [ + 716, + 1148, + 1 + ], + [ + 873, + 858, + 1 + ], + [ + 1121, + 36, + 2 + ], + [ + 900, + 862, + 1 + ], + [ + 328, + 328, + 1 + ], + [ + 868, + 828, + 1 + ], + [ + 993, + 993, + 1 + ], + [ + 831, + 832, + 2 + ], + [ + 831, + 870, + 2 + ], + [ + 867, + 911, + 1 + ], + [ + 958, + 958, + 1 + ], + [ + 721, + 720, + 1 + ], + [ + 208, + 1287, + 1 + ], + [ + 1000, + 207, + 1 + ], + [ + 903, + 863, + 1 + ], + [ + 1022, + 954, + 1 + ], + [ + 824, + 901, + 1 + ], + [ + 859, + 995, + 1 + ], + [ + 1049, + 1049, + 4 + ], + [ + 1076, + 44, + 2 + ], + [ + 529, + 527, + 1 + ], + [ + 869, + 854, + 2 + ], + [ + 836, + 845, + 2 + ], + [ + 711, + 710, + 1 + ], + [ + 1096, + 535, + 2 + ], + [ + 826, + 1055, + 1 + ], + [ + 883, + 841, + 1 + ], + [ + 726, + 1171, + 1 + ], + [ + 31, + 30, + 1 + ], + [ + 1014, + 198, + 1 + ], + [ + 847, + 1037, + 2 + ], + [ + 875, + 932, + 1 + ], + [ + 836, + 851, + 2 + ], + [ + 877, + 930, + 1 + ], + [ + 1056, + 843, + 4 + ], + [ + 848, + 998, + 1 + ], + [ + 902, + 872, + 1 + ], + [ + 1097, + 124, + 2 + ], + [ + 829, + 835, + 1 + ], + [ + 518, + 519, + 1 + ], + [ + 924, + 968, + 1 + ], + [ + 828, + 868, + 1 + ], + [ + 833, + 855, + 2 + ], + [ + 920, + 840, + 1 + ], + [ + 818, + 829, + 1 + ], + [ + 891, + 959, + 1 + ], + [ + 730, + 729, + 1 + ], + [ + 965, + 885, + 2 + ], + [ + 894, + 894, + 3 + ], + [ + 906, + 871, + 1 + ], + [ + 824, + 986, + 1 + ], + [ + 950, + 984, + 1 + ], + [ + 944, + 944, + 2 + ], + [ + 1105, + 718, + 2 + ], + [ + 1047, + 1047, + 2 + ], + [ + 865, + 865, + 2 + ], + [ + 824, + 968, + 1 + ], + [ + 537, + 535, + 1 + ], + [ + 1125, + 1125, + 2 + ], + [ + 1026, + 1026, + 3 + ] + ] +} diff --git a/profiles/eef/config.json b/profiles/eef/config.json new file mode 100644 index 0000000000000000000000000000000000000000..8ba90c25fb80027af5144546af9676d7efc7d9f8 --- /dev/null +++ b/profiles/eef/config.json @@ -0,0 +1,4 @@ +{ + "model_type": "actioncodec2_unified", + "processor_class": "ActionCodec2Processor" +} diff --git a/profiles/eef/config.yaml b/profiles/eef/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..38417ae9a237e14157f5e75a8dd639a40182c682 --- /dev/null +++ b/profiles/eef/config.yaml @@ -0,0 +1,74 @@ +model_type: actioncodec2_unified +config_version: 1 +physical: + input_dim: 7 + components: + translation_m: + dims: + start: 0 + stop: 3 + space: additive + input_order: 0 + primitive_order: 1 + input_transform: + scale: 1.0 + clip: false + constraints: + order_0: + min: -2.0 + max: 2.0 + order_1: + min: -0.08 + max: 0.08 + distortion_limit: 0.001 + primitive_selection: smooth + orientation_so3: + dims: + start: 3 + stop: 6 + space: so3 + input_order: 0 + primitive_order: 1 + input_transform: + scale: 1.0 + clip: false + constraints: + order_0: + min: -3.141592653589793 + max: 3.141592653589793 + order_1: + min: -0.25 + max: 0.25 + distortion_limit: 0.004 + primitive_selection: smooth + representation: rotation_vector + multiplication: right + gripper_open_command: + dims: + - 6 + space: additive + input_order: 0 + primitive_order: 0 + input_transform: + scale: 1.0 + clip: true + constraints: + order_0: + min: 0.0 + max: 1.0 + distortion_limit: 0.5 + primitive_selection: nearest +set_bpe: + lookahead: 8 + allow_cross_component_merge: false +train: + primitives: + zero_primitive_components: + - translation_m + - orientation_so3 + bpe: + spatial_vocab_size: 1024 + vocab_size: 2048 + quality_control: + hard_order_1_scale: 3.0 + min_segment_length: 2 diff --git a/profiles/eef/primitives.npz b/profiles/eef/primitives.npz new file mode 100644 index 0000000000000000000000000000000000000000..c820267995e274355db50a20506f9978c3974a4d --- /dev/null +++ b/profiles/eef/primitives.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b393814f311daf6b609475c9cf943bfc9dde617b6b4249bb404df50d98b2526b +size 7406 diff --git a/profiles/eef/processor_config.json b/profiles/eef/processor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..e0091cc74d1d60cc0979bf84338cd6ef8ab117b4 --- /dev/null +++ b/profiles/eef/processor_config.json @@ -0,0 +1,3 @@ +{ + "processor_class": "ActionCodec2Processor" +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..dd0e26a7fb93782a3ad58086b12a81c13270c41b --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,16 @@ +[build-system] +requires = ["setuptools>=77", "wheel", "pybind11>=2.12"] +build-backend = "setuptools.build_meta" + +[project] +name = "actioncodec2-kernels" +version = "0.1.0" +description = "Portable C++17 encode/decode kernels for ActionCodec2 tokenizer artifacts" +requires-python = ">=3.10" +license = "Apache-2.0" + +[tool.setuptools] +packages = ["actioncodec2_kernels"] + +[tool.setuptools.package-data] +actioncodec2_kernels = ["*.cpp", "*.hpp"] diff --git a/release.json b/release.json new file mode 100644 index 0000000000000000000000000000000000000000..5a178babec23f8deed73901889ce4b92004227a2 --- /dev/null +++ b/release.json @@ -0,0 +1,12 @@ +{ + "edition": "Extended", + "encoding_order": 1, + "precision": "1mm", + "codebook_size": 2048, + "profiles": [ + "eef" + ], + "gripper": "binary", + "source_revision": "589f307", + "libero_evaluated": false +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..a140fba756c72a75eaffe68088f049d7000f7f65 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +numpy +scipy +torch>=2.3 +transformers>=4.57,<6 +huggingface-hub +pyyaml diff --git a/router_config.yaml b/router_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..73d2824d8ce21f0adb5895e1a284de42d8c159ef --- /dev/null +++ b/router_config.yaml @@ -0,0 +1,120 @@ +model_type: actioncodec2_router +config_version: 3 +profiles: +- name: eef + token_offset: 0 + token_budget: 2048 + artifact_fingerprint: 9272f9beecca4cd1bd2906f2b34c27d81f6f95ba547010843c8f2edfc7018af6 +presets: +- name: single_eef_abs + input_dim: 7 + input: absolute + source_hz: 15.0 + target_hz: 15.0 + interpolation: cubic + components: + arm: + space: eef + position_dims: + - 0 + - 1 + - 2 + rotation_dims: + - 3 + - 4 + - 5 + rotation_format: rotation_vector + gripper_dim: 6 + gripper_threshold: 0.8 + gripper_open_is_high: true +- name: single_eef_delta + input_dim: 7 + input: delta + source_hz: 15.0 + target_hz: 15.0 + interpolation: cubic + components: + arm: + space: eef + position_dims: + - 0 + - 1 + - 2 + rotation_dims: + - 3 + - 4 + - 5 + rotation_format: rotation_vector + gripper_dim: 6 + gripper_threshold: 0.8 + gripper_open_is_high: true +- name: dual_eef_abs + input_dim: 14 + input: absolute + source_hz: 15.0 + target_hz: 15.0 + interpolation: cubic + components: + left_arm: + space: eef + position_dims: + - 0 + - 1 + - 2 + rotation_dims: + - 3 + - 4 + - 5 + rotation_format: rotation_vector + gripper_dim: 6 + gripper_threshold: 0.8 + gripper_open_is_high: true + right_arm: + space: eef + position_dims: + - 7 + - 8 + - 9 + rotation_dims: + - 10 + - 11 + - 12 + rotation_format: rotation_vector + gripper_dim: 13 + gripper_threshold: 0.8 + gripper_open_is_high: true +- name: dual_eef_delta + input_dim: 14 + input: delta + source_hz: 15.0 + target_hz: 15.0 + interpolation: cubic + components: + left_arm: + space: eef + position_dims: + - 0 + - 1 + - 2 + rotation_dims: + - 3 + - 4 + - 5 + rotation_format: rotation_vector + gripper_dim: 6 + gripper_threshold: 0.8 + gripper_open_is_high: true + right_arm: + space: eef + position_dims: + - 7 + - 8 + - 9 + rotation_dims: + - 10 + - 11 + - 12 + rotation_format: rotation_vector + gripper_dim: 13 + gripper_threshold: 0.8 + gripper_open_is_high: true diff --git a/runtime/__init__.py b/runtime/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..80ac1fb2f5fdce15983f50c63e788e5f870bcccd --- /dev/null +++ b/runtime/__init__.py @@ -0,0 +1 @@ +"""Private inference runtime bundled with this ActionCodec2 artifact.""" diff --git a/runtime/backend/__init__.py b/runtime/backend/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7f82d8351a7c8342f78ecb6d99091596a4144b01 --- /dev/null +++ b/runtime/backend/__init__.py @@ -0,0 +1 @@ +"""Internal backend components for ActionCodec2.""" diff --git a/runtime/backend/native_backend.py b/runtime/backend/native_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..5e4ef5b612370f8a9b5d2d2e7e30f79ef8c2e108 --- /dev/null +++ b/runtime/backend/native_backend.py @@ -0,0 +1,39 @@ +"""Optional CPU accelerator; build with python setup.py build_ext --inplace.""" + +from __future__ import annotations + +import os + +__all__ = ["native", "have_native", "require_native", "resolve_backend"] + +try: # pragma: no cover - exercised by whichever branch the build state picks + from .. import _native as native +except ImportError: # pragma: no cover + native = None + + +def have_native() -> bool: + """True when the compiled kernels are importable and not disabled.""" + return native is not None and not os.environ.get("ACTIONCODEC2_NO_NATIVE") + + +def require_native(): + """Return the native module or explain how to get one.""" + if not have_native(): + raise RuntimeError( + "the actioncodec2 native extension is not available; build it from " + "the package native/ directory (python setup.py build_ext --inplace) or " + "select backend='python'" + ) + return native + + +def resolve_backend(backend: str) -> str: + """Turn ``'auto'`` into a concrete backend name.""" + if backend not in {"auto", "native", "python"}: + raise ValueError("backend must be 'auto', 'native' or 'python'") + if backend == "auto": + return "native" if have_native() else "python" + if backend == "native": + require_native() + return backend diff --git a/runtime/common/__init__.py b/runtime/common/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2c5964e8bf7ff68dc273e84ba3ef3b4d4c454a45 --- /dev/null +++ b/runtime/common/__init__.py @@ -0,0 +1 @@ +"""Internal common components for ActionCodec2.""" diff --git a/runtime/common/_frontend_base.py b/runtime/common/_frontend_base.py new file mode 100644 index 0000000000000000000000000000000000000000..2ca96253d24755cb46cb9f377a8eb11bd3d7a5cd --- /dev/null +++ b/runtime/common/_frontend_base.py @@ -0,0 +1,108 @@ +"""ActionCodec2-only tensor and token-row validation helpers.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np +import torch + +from ._types import ActionIdxBatch, TokenIds + + +class BaseActionTokenizer: + """Shared action tensor methods, without policy-specific sentinels.""" + + @property + def new_action_tokens(self) -> list[str]: + """Token strings in id order — downstream registers these into its VLM vocab.""" + return list(self._new_action_tokens) + + def to_token_strs(self, token_ids: TokenIds) -> list[str]: + """Render one id row as its token strings, e.g. ```` / ````. + + This is exactly the text a VLM would see after registering + :attr:`new_action_tokens` — useful for eyeballing sequences and for + prompt-template debugging. + + Args: + token_ids: One row of action-token ids in ``[0, vocab_size)``. + + Returns: + list[str]: Token strings in row order. + """ + row = self._to_id_row(token_ids).tolist() + n = len(self._new_action_tokens) + bad = [i for i in row if not 0 <= i < n] + if bad: + raise ValueError(f"token ids {bad[:5]} out of range [0, {n})") + return [self._new_action_tokens[i] for i in row] + + def _check_encode_kwargs(self, encode_kwargs: dict | None) -> dict: + """Normalize ``encode_kwargs`` to a dict, rejecting keys outside :attr:`encode_kwargs`.""" + kw = dict(encode_kwargs or {}) + unknown = sorted(set(kw) - self.encode_kwargs) + if unknown: + raise ValueError( + f"{type(self).__name__} does not consume encode_kwargs {unknown}; it declares {sorted(self.encode_kwargs)}" + ) + return kw + + @staticmethod + def _to_id_row(token_ids: TokenIds) -> torch.Tensor: + """Validate and normalize one integer id row to a 1-D long tensor.""" + if isinstance(token_ids, torch.Tensor): + row = token_ids.detach().cpu() + if ( + row.dtype == torch.bool + or row.dtype.is_floating_point + or row.dtype.is_complex + ): + raise ValueError("token ids must be integers") + row = row.to(torch.long) + else: + values = np.asarray(token_ids) + if values.dtype.kind not in {"i", "u"}: + raise ValueError("token ids must be integers") + row = torch.from_numpy(np.asarray(values, dtype=np.int64)) + if row.ndim != 1: + raise ValueError(f"token id row must be 1D, got {tuple(row.shape)}") + return row + + def _checked_decode_row(self, token_ids: TokenIds) -> torch.Tensor: + """Return one validated row without repairing invalid model output.""" + row = self._to_id_row(token_ids) + invalid = (row < 0) | (row >= int(self._action_vocab_size)) + if torch.any(invalid): + bad = row[invalid][:5].tolist() + raise ValueError( + f"token ids {bad} out of range [0, {int(self._action_vocab_size)})" + ) + return row + + def _decode_token_ids_to_actions_batch( + self, + action_token_ids: Sequence[TokenIds], + *, + time_horizon: int | None = None, + action_dim: int | None = None, + decode_kwargs: dict | None = None, + ) -> torch.Tensor: + """Decode a batch of id rows to actions ``(B, T, D)``. + + The decoder applies the configured action-space layout to each row. + """ + batch_action_idx: ActionIdxBatch = [ + self._checked_decode_row(a).tolist() for a in action_token_ids + ] + (actions_np, _) = self._decode_action_indices( + batch_action_idx, + time_horizon=time_horizon, + action_dim=action_dim, + decode_kwargs=decode_kwargs, + ) + return torch.from_numpy(actions_np) + + @property + def vocab_size(self) -> int: + return int(self._action_vocab_size) diff --git a/runtime/common/_types.py b/runtime/common/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..957a2446a646af50383079cabe5fc00380770737 --- /dev/null +++ b/runtime/common/_types.py @@ -0,0 +1,12 @@ +"""Shared lightweight types for the frontend layer.""" + +from __future__ import annotations + +import numpy as np +import torch + +ActionIdxBatch = list[list[int]] +TokenIds = list[int] | np.ndarray | torch.Tensor + + +__all__ = ["ActionIdxBatch", "TokenIds"] diff --git a/runtime/common/_validation.py b/runtime/common/_validation.py new file mode 100644 index 0000000000000000000000000000000000000000..f091b7498d43a4be4ccbfb47116825e6fdf44068 --- /dev/null +++ b/runtime/common/_validation.py @@ -0,0 +1,89 @@ +"""Shared validation primitives for ActionCodec2 public boundaries.""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from typing import TypeVar + +import numpy as np + +ErrorT = TypeVar("ErrorT", bound=Exception) + + +def active_dimensions( + active_dims: Sequence[int] | None, + dimension: int, + *, + subject: str = "active_dims", +) -> tuple[int, ...]: + """Normalize and validate a non-empty increasing subset of dimensions.""" + dims = ( + tuple(range(dimension)) + if active_dims is None + else tuple(int(item) for item in active_dims) + ) + if not dims: + raise ValueError(f"{subject} must not be empty") + if any(left >= right for left, right in zip(dims, dims[1:])): + raise ValueError(f"{subject} must be strictly increasing") + if dims[0] < 0 or dims[-1] >= dimension: + raise ValueError(f"{subject} contains a dimension outside the vocabulary") + return dims + + +def optional_finite_array( + value: np.ndarray | None, + *, + shape: tuple[int, ...], + name: str = "current_state", +) -> np.ndarray | None: + """Convert an optional float32 array and enforce its exact shape/finiteness.""" + if value is None: + return None + result = np.asarray(value, dtype=np.float32) + if result.shape != shape: + rendered = ", ".join(str(item) for item in shape) + raise ValueError(f"{name} must have shape [{rendered}]") + if not np.all(np.isfinite(result)): + raise ValueError(f"{name} contains a non-finite value") + return result + + +def broadcast_vector( + value: object, + width: int, + *, + dtype: np.dtype, + name: str, +) -> np.ndarray: + """Broadcast a scalar or width-sized value to one numerical vector.""" + array = np.asarray(value, dtype=dtype) + try: + return np.broadcast_to(array, (width,)).copy() + except ValueError as error: + raise ValueError(f"{name} must be scalar or contain {width} values") from error + + +def validate_dimension_partition( + groups: Iterable[tuple[str, Sequence[int]]], + dimension: int, + *, + label: str, + error_type: type[ErrorT] = ValueError, +) -> None: + """Validate that named dimension groups form one exact non-overlapping partition.""" + owners: dict[int, str] = {} + for name, dims in groups: + for index in dims: + if index < 0 or index >= dimension: + raise error_type( + f"{label} {index} used by {name!r} is outside input_dim" + ) + if index in owners: + raise error_type( + f"{label} {index} belongs to both {owners[index]!r} and {name!r}" + ) + owners[index] = name + missing = sorted(set(range(dimension)) - set(owners)) + if missing: + raise error_type(f"{label}(s) {missing} are not assigned") diff --git a/runtime/frontend/__init__.py b/runtime/frontend/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..66788de9f88a4f53e82e6a93ac898bd709a4f573 --- /dev/null +++ b/runtime/frontend/__init__.py @@ -0,0 +1 @@ +"""Internal frontend components for ActionCodec2.""" diff --git a/runtime/frontend/tokenizer.py b/runtime/frontend/tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..9ea24207fa907810203e3cb5fa7562df0e211099 --- /dev/null +++ b/runtime/frontend/tokenizer.py @@ -0,0 +1,398 @@ +"""Action-space-bound, arbitrary-horizon ActionCodec2 frontend.""" + +from __future__ import annotations + +import sys +import warnings +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import TextIO + +import numpy as np +import torch + +from ..common._frontend_base import BaseActionTokenizer +from ..common._types import ActionIdxBatch, TokenIds +from ..routing.presets import PresetConfig +from ..routing.routing import ActionCodec2Router +from ..physical.semantic import resampled_horizon + + +def _numpy(value): + """Convert a tensor to CPU float32 without changing NumPy input semantics.""" + return ( + value.detach().cpu().float().numpy() + if isinstance(value, torch.Tensor) + else value + ) + + +class ActionCodec2Tokenizer(BaseActionTokenizer): + """ActionCodec2 frontend bound to physical tensor semantics, not a horizon. + + The selected action space fixes the meaning of every source column. Sequence + length remains unrestricted: encode and decode infer any positive ``T`` + directly from the lossless token coverage. + + Use :meth:`from_pretrained` for the normal one-call loading path. Built-in + and registered action spaces are discoverable through :meth:`print_action_spaces` + and are included in :func:`str` output. + + Args: + backend: Router containing the shared physical processors. + action_space: Registered action-space name to use for this frontend. + """ + + encode_kwargs = frozenset({"fps", "current_state", "previous_first_order"}) + + def __init__(self, backend: ActionCodec2Router, *, action_space: str) -> None: + if not isinstance(backend, ActionCodec2Router): + raise TypeError( + f"ActionCodec2Tokenizer requires ActionCodec2Router, got {type(backend).__name__}" + ) + preset = backend.preset(action_space) + self.action_tokenizer = backend + self._action_space = preset + self._action_vocab_size = int(backend.vocab_size) + self._new_action_tokens = list(backend.new_action_tokens) + self.action_dim = preset.input_dim + self._warned_default_fps = False + self.requires_previous_first_order = any( + physical.primitive_order == 2 + and set(component.active_dims).intersection(physical.dims) + for component in preset.components + for physical in backend.profiles[ + component.space + ].processor.config.physical.components.values() + ) + self._encoded_action_dim = sum( + len(component.active_dims) for component in preset.components + ) + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str | Path, + *, + action_space: str | Path | PresetConfig | Mapping[str, object], + **kwargs: object, + ) -> ActionCodec2Tokenizer: + """Load a local/Hub router and bind one explicit physical action space. + + Args: + pretrained_model_name_or_path: Local artifact directory or Hub repo ID. + action_space: Registered name, preset YAML path, ``PresetConfig``, or + strict preset mapping. Custom configs are registered before binding. + **kwargs: Standard Hub snapshot options such as ``revision``, + ``cache_dir``, ``token``, and ``local_files_only``. + """ + router = ActionCodec2Router.from_pretrained( + pretrained_model_name_or_path, **kwargs + ) + if isinstance(action_space, (PresetConfig, Mapping, Path)): + name = router.register_action_space(action_space).name + elif isinstance(action_space, str): + candidate = Path(action_space) + name = ( + router.register_action_space(candidate).name + if candidate.is_file() + else action_space + ) + else: + raise TypeError( + "action_space must be a registered name, PresetConfig, mapping, or YAML path" + ) + return cls(router, action_space=name) + + def __str__(self) -> str: + """Show supported tensor layouts, physical units, and invocation requirements.""" + return self._describe_action_spaces() + + def _describe_action_spaces(self, action_dim: int | None = None) -> str: + return self.action_tokenizer.describe_action_spaces( + selected=self.action_space.name, action_dim=action_dim + ) + + @property + def action_space(self) -> PresetConfig: + """The immutable tensor semantics selected for this tokenizer.""" + return self._action_space + + @property + def action_spaces(self) -> Mapping[str, PresetConfig]: + """Every built-in or user-registered option in the loaded artifact.""" + return self.action_tokenizer.action_spaces + + def print_action_spaces( + self, *, action_dim: int | None = None, file: TextIO | None = None + ) -> None: + """Print all layouts or candidates matching ``action_dim`` without changing the binding.""" + description = self._describe_action_spaces(action_dim) + print(description, file=sys.stdout if file is None else file) + + def register_action_space( + self, + preset: PresetConfig | Mapping[str, object] | str | Path, + ) -> PresetConfig: + """Register a custom action space without changing this instance's binding.""" + return self.action_tokenizer.register_action_space(preset) + + def for_action_space(self, name: str) -> ActionCodec2Tokenizer: + """Return a new tokenizer sharing weights but bound to another option.""" + return type(self)(self.action_tokenizer, action_space=name) + + def _source_fps(self, fps, preset=None): + preset = self.action_space if preset is None else preset + if fps is None: + if not self._warned_default_fps: + warnings.warn( + f"fps was omitted; assuming {preset.source_hz:g} Hz from action_space " + f"{preset.name!r}. Pass the actual fps explicitly to avoid changing action timing.", + UserWarning, + stacklevel=3, + ) + self._warned_default_fps = True + return preset.source_hz + return fps + + def first_order_template(self, batch_size: int) -> dict[str, np.ndarray]: + """Explicit zero boundary templates, shaped {component: (B,D_active)}. + + Zeros are valid only when the preceding canonical increment is known to + be zero. Otherwise replace them with measured/reconstructed values at the + action space's configured codec rate. + Only active second-order components are included; first order returns {}. + """ + if type(batch_size) is not int or batch_size < 1: + raise ValueError("batch_size must be a positive integer") + return { + c.name: np.zeros((batch_size, len(c.active_dims)), dtype=np.float32) + for c in self.action_space.components + if any( + p.primitive_order == 2 and set(c.active_dims).intersection(p.dims) + for p in self.action_tokenizer.profiles[ + c.space + ].processor.config.physical.components.values() + ) + } + + def max_tokens_for_horizon(self, horizon: int, *, fps: float | None = None) -> int: + """Return the token bound after converting a source chunk to the codec clock. + + This helper describes encode input geometry. It does not store + ``horizon`` in the token stream or pass it to decode. + """ + source_fps = self._source_fps(fps) + internal = resampled_horizon( + horizon, source_hz=source_fps, target_hz=self.action_space.target_hz + ) + return internal * self._encoded_action_dim + + def grammar(self, horizon: int, *, fps: float | None = None): + """Build a grammar for the codec steps implied by source chunk geometry. + + The grammar constrains token coverage only. When ``fps`` exceeds the + codec clock, several source row counts can map to the same codec + step count; decode therefore returns the unique canonical resampling + implied by the tokens and requested ``fps``, not a hidden source count. + """ + source_fps = self._source_fps(fps) + internal = resampled_horizon( + horizon, source_hz=source_fps, target_hz=self.action_space.target_hz + ) + return self.action_tokenizer.grammar( + preset=self.action_space.name, horizon=internal + ) + + def __call__( + self, + action: torch.Tensor | np.ndarray, + *, + fps: float | None = None, + current_state: torch.Tensor | np.ndarray | None = None, + previous_first_order: Mapping[str, torch.Tensor | np.ndarray] | None = None, + encode_kwargs: dict | None = None, + ) -> ActionIdxBatch: + """Encode arbitrary-length ``(B,T,D)`` or ``(T,D)`` actions. + + Second order requires ``previous_first_order`` mapping component names + to canonical ``(B,D_active)`` increments at the configured codec rate (``B=1`` for a single + trajectory). Absolute inputs also require ``current_state(B,D)``. + """ + kwargs = self._check_encode_kwargs(encode_kwargs) + direct = { + "fps": fps, + "current_state": current_state, + "previous_first_order": previous_first_order, + } + duplicates = sorted( + key for key, value in direct.items() if value is not None and key in kwargs + ) + if duplicates: + raise ValueError(f"encode arguments provided twice: {duplicates}") + kwargs.update( + {key: value for key, value in direct.items() if value is not None} + ) + reference = kwargs.get("current_state") + if self.action_space.input == "absolute" and reference is None: + raise ValueError( + f"action space {self.action_space.name!r} is absolute; encode requires " + "current_state with shape [B, D] to preserve the first transition" + ) + if self.action_space.input == "delta" and reference is not None: + raise ValueError( + f"action space {self.action_space.name!r} is delta; omit current_state" + ) + values = torch.as_tensor(action) + if values.ndim == 2: + values = values.unsqueeze(0) + return self._encode_action_indices(values, encode_kwargs=kwargs) + + encode = __call__ + + def decode( + self, + action_token_ids: Sequence[TokenIds], + *, + fps: float | None = None, + current_state: torch.Tensor | np.ndarray | None = None, + previous_first_order: Mapping[str, torch.Tensor | np.ndarray] | None = None, + ) -> torch.Tensor: + """Decode a batch, deriving sequence length from tokens and output rate. + + At the codec rate, token coverage determines ``T`` exactly. Other output rates + deterministically resample that decoded path. No source horizon is + accepted because downsampling cannot preserve an arbitrary original row + count without encoding extra metadata. + ``previous_first_order`` maps component names to canonical + ``(B,D_active)`` codec-rate increments; second order requires the same + boundary as encoding, independently of ``current_state(B,D)``. + """ + return self._decode_token_ids_to_actions_batch( + action_token_ids, + action_dim=self.action_dim, + decode_kwargs={ + "fps": fps, + "current_state": current_state, + "previous_first_order": previous_first_order, + }, + ) + + @staticmethod + def _numpy_boundary(previous_first_order): + if previous_first_order is None: + return None + if not isinstance(previous_first_order, Mapping): + raise ValueError( + "previous_first_order must map component names to [B,D_active] arrays" + ) + return {name: _numpy(value) for name, value in previous_first_order.items()} + + def final_first_order( + self, + action_token_ids: Sequence[TokenIds], + *, + previous_first_order: Mapping[str, torch.Tensor | np.ndarray], + executed_steps: int | None = None, + ) -> dict[str, np.ndarray]: + """Return reconstructed component states after a canonical-step prefix. + + Results are ``[B,D_active]`` arrays at the configured codec rate, + suitable for the next encode/decode call. ``executed_steps`` counts + codec steps. + """ + return self.action_tokenizer.final_first_order( + [self._checked_decode_row(row).tolist() for row in action_token_ids], + preset=self.action_space.name, + previous_first_order=self._numpy_boundary(previous_first_order), + executed_steps=executed_steps, + ) + + def _encode_action_indices( + self, action: torch.Tensor, encode_kwargs: dict | None = None + ) -> ActionIdxBatch: + """Encode finite ``(B,T,D)`` actions without fixing ``T`` at construction.""" + kwargs = dict(encode_kwargs or {}) + fps = float(self._source_fps(kwargs.get("fps"))) + current_state = _numpy(kwargs.get("current_state")) + values = _numpy(action) + if values.ndim != 3 or values.shape[2] != self.action_space.input_dim: + raise ValueError( + f"actions must have shape [B, T_source, {self.action_space.input_dim}]" + ) + return self.action_tokenizer.encode_batch( + values, + preset=self.action_space.name, + fps=fps, + current_state=current_state, + previous_first_order=self._numpy_boundary( + kwargs.get("previous_first_order") + ), + ) + + def _decode_action_indices( + self, + action_indices: ActionIdxBatch, + *, + time_horizon: int | None = None, + action_dim: int | None = None, + decode_kwargs: dict | None = None, + ) -> tuple[np.ndarray, list[set[str]]]: + """Decode rows to the selected preset's original ``(B,T,D)`` layout.""" + if time_horizon is not None: + raise ValueError( + "ActionCodec2 decode infers sequence length from tokens; " + "do not provide time_horizon" + ) + if action_dim is not None and int(action_dim) != self.action_dim: + raise ValueError( + f"action_dim={action_dim} disagrees with action space " + f"{self.action_space.name!r} (D={self.action_dim})" + ) + kwargs = dict(decode_kwargs or {}) + unknown = sorted(set(kwargs) - self.encode_kwargs) + if unknown: + raise ValueError( + f"ActionCodec2Tokenizer does not consume decode_kwargs {unknown}" + ) + raw_fps = kwargs.get("fps") + fps = float(self._source_fps(raw_fps)) + current_state = _numpy(kwargs.get("current_state")) + if self.action_space.input == "absolute" and current_state is None: + raise ValueError( + f"action space {self.action_space.name!r} is absolute; decode requires " + "current_state with shape [B, D] to restore its external reference" + ) + if self.action_space.input == "delta" and current_state is not None: + raise ValueError( + f"action space {self.action_space.name!r} is delta; omit current_state " + "to decode in the registered input semantics" + ) + decoded = self.action_tokenizer.decode_batch( + action_indices, + preset=self.action_space.name, + fps=fps, + current_state=current_state, + previous_first_order=self._numpy_boundary( + kwargs.get("previous_first_order") + ), + ) + actions = self._restore_source_layout(decoded) + return actions, [set() for _ in range(actions.shape[0])] + + def _restore_source_layout(self, decoded: Mapping[str, np.ndarray]) -> np.ndarray: + """Place component outputs back into their registered source columns.""" + first = next(iter(decoded.values())) + batch, steps = first.shape[:2] + if any(values.shape[:2] != (batch, steps) for values in decoded.values()): + raise RuntimeError( + "decoded components disagree on batch or sequence length" + ) + actions = np.empty((batch, steps, self.action_dim), dtype=np.float32) + for component in self.action_space.components: + values = decoded[component.name] + motion_width = len(component.source_dims) + actions[:, :, component.source_dims] = values[:, :, :motion_width] + if component.gripper_dim is not None: + actions[:, :, component.gripper_dim] = values[:, :, motion_width] + return actions diff --git a/runtime/hub_acceleration.py b/runtime/hub_acceleration.py new file mode 100644 index 0000000000000000000000000000000000000000..a429bd5e387f91d86ddf18a4b7d4de7da614d13b --- /dev/null +++ b/runtime/hub_acceleration.py @@ -0,0 +1,239 @@ +"""Install the shared C++ kernels in a portable ActionCodec2 runtime. + +This module is copied into exported Hub artifacts. It imports runtime classes +by their private package name so each artifact retains its own Python fallback. +""" + +from __future__ import annotations + +import functools +import importlib +from numbers import Integral +import os + +import numpy as np + + +def load_native(): + """Return installed portable or VLA kernels, with a fork-safe thread budget.""" + if os.environ.get("ACTIONCODEC2_NO_NATIVE"): + return None + for name in ( + "actioncodec2_kernels._actioncodec2_native", + "actioncodec2.tokenizers._actioncodec2_native", + ): + try: + native = importlib.import_module(name) + except ImportError: + continue + _configure(native) + return native + return None + + +@functools.cache +def _configure(native): + configured = os.environ.get("ACTIONCODEC2_NUM_THREADS") + if configured is not None: + if not configured.isdigit() or int(configured) < 1: + raise ValueError("ACTIONCODEC2_NUM_THREADS must be a positive integer") + native.set_thread_budget(int(configured)) + else: + native.set_thread_budget(0) + if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=lambda: native.set_thread_budget(1)) + + +def install(package: str) -> None: + """Accelerate an artifact while preserving its profiles and public contracts.""" + backend = importlib.import_module(f"{package}.backend.native_backend") + if getattr(backend, "_portable_acceleration", False): + return + backend._portable_acceleration = True + backend.native = load_native() + semantic = importlib.import_module(f"{package}.physical.semantic") + presets = importlib.import_module(f"{package}.routing.presets") + trajectory = importlib.import_module(f"{package}.physical.trajectory") + physical = importlib.import_module(f"{package}.physical.physical_processor") + validation = importlib.import_module(f"{package}.common._validation") + + def native_now(): + return backend.native if backend.have_native() else None + + def wrap_pose(name, kernel, width): + original = getattr(semantic, name) + + @functools.wraps(original) + def wrapped(values): + native = native_now() + array = np.asarray(values) + if (native is None or array.dtype != np.float32 or array.ndim != 2 + or array.shape[1] != width or len(array) < 1 + or not np.all(np.isfinite(array))): + return original(values) + return getattr(native, kernel)(array[None])[0] + + setattr(semantic, name, wrapped) + if hasattr(presets, name): + setattr(presets, name, wrapped) + + for name, kernel, width in ( + ("integrate_eef_deltas", "integrate_eef_deltas_batch", 6), + ("eef_pose_absolute", "eef_pose_absolute_batch", 7), + ("eef_pose_deltas", "eef_pose_deltas_batch", 7), + ): + wrap_pose(name, kernel, width) + + original_resample = semantic.resample_uniform + + @functools.wraps(original_resample) + def resample(values, *, source_hz, target_hz, kind="linear", target_length=None, + method="linear"): + native = native_now() + array = np.asarray(values) + if (native is None or array.dtype != np.float32 or array.ndim != 2 + or len(array) < 1 or not np.all(np.isfinite(array)) + or (target_length is not None and ( + isinstance(target_length, bool) or not isinstance(target_length, Integral) + or target_length < 1))): + return original_resample(values, source_hz=source_hz, target_hz=target_hz, + kind=kind, target_length=target_length, method=method) + return native.resample_uniform_batch( + array[None], source_hz, target_hz, kind, + -1 if target_length is None else target_length, method, + )[0] + + semantic.resample_uniform = presets.resample_uniform = resample + original_canonicalize = presets.PresetConfig.canonicalize + + @functools.wraps(original_canonicalize) + def canonicalize(self, values, *, source_hz=None, current_state=None): + native = native_now() + array = np.asarray(values, dtype=np.float32) + if (native is None or self.input != "delta" or current_state is not None + or array.ndim != 3 or min(array.shape) < 1 + or array.shape[2] != self.input_dim or not np.all(np.isfinite(array)) + or any(c.space == "eef" and c.rotation_format != "rotation_vector" + for c in self.components)): + return original_canonicalize(self, values, source_hz=source_hz, + current_state=current_state) + rate = self.source_hz if source_hz is None else source_hz + output = {} + for component in self.components: + if component.space == "eef": + motion = np.concatenate((array[..., component.position_dims], + array[..., component.rotation_dims]), axis=2) + pose = native.integrate_eef_deltas_batch(motion) + sampled = native.resample_uniform_batch( + pose, rate, self.target_hz, "pose_xyzw", -1, self.interpolation + ) + motion = native.eef_pose_absolute_batch(sampled) + else: + motion = np.cumsum(array[..., component.joint_dims], axis=1, + dtype=np.float32) + motion = native.resample_uniform_batch( + motion, rate, self.target_hz, "linear", -1, self.interpolation + ) + if component.gripper_dim is not None: + gripper = native.resample_uniform_batch( + array[..., component.gripper_dim:component.gripper_dim + 1], + rate, self.target_hz, "hold", -1, self.interpolation, + ) + gripper = semantic.binarize_gripper_actions( + gripper.reshape(-1, 1), threshold=component.gripper_threshold, + open_is_high=component.gripper_open_is_high, + ).reshape(gripper.shape) + motion = np.concatenate((motion, gripper), axis=2) + output[component.name] = [ + trajectory.ActiveTrajectory(row, component.active_dims) for row in motion + ] + return output + + presets.PresetConfig.canonicalize = canonicalize + original_encode = physical.ActionCodec2Processor.encode_batch + original_decode = physical.ActionCodec2Processor.decode_batch + + def serializer(self): + native = native_now() + if native is None or self.serializer.scheme != "set_bpe": + return None + if not hasattr(self, "_portable_setbpe"): + saved = self.serializer + self._portable_setbpe = native.SetBPE( + saved.bins, saved.lookahead, saved.spatial_merges, saved.temporal_merges + ) + return self._portable_setbpe + + @functools.wraps(original_encode) + def encode_batch(self, actions, *, active_dims=None, previous_first_order=None): + codec = serializer(self) + if codec is None: + return original_encode(self, actions, active_dims=active_dims, + previous_first_order=previous_first_order) + values = np.asarray(actions, dtype=np.float32) + if values.ndim != 3 or min(values.shape) < 1: + raise ValueError("actions must have non-empty shape [B, T, D_active]") + dims = validation.active_dimensions(active_dims, self.action_dim) + if len(dims) != values.shape[2]: + raise ValueError("active_dims must have one entry per action coordinate") + dense = np.zeros((*values.shape[:2], self.action_dim), dtype=np.float32) + dense[..., dims] = values + previous = self._dense_previous_first_order(previous_first_order, dims, len(values)) + cells = self.quantize_batch(dense, previous_first_order=previous) + absent = np.ones(self.action_dim, dtype=bool) + absent[list(dims)] = False + cells[..., absent] = -1 + return codec.encode_batch(cells) + + @functools.wraps(original_decode) + def decode_batch(self, streams, *, active_dims=None, previous_first_order=None): + codec = serializer(self) + if codec is None: + return original_decode(self, streams, active_dims=active_dims, + previous_first_order=previous_first_order) + dims = validation.active_dimensions(active_dims, self.action_dim) + checked = [] + for stream in streams: + row = [] + for token in stream: + if isinstance(token, bool) or not isinstance(token, Integral): + raise ValueError("tokens must be integers") + if not 0 <= int(token) < self.vocab_size: + raise ValueError("token is outside the vocabulary") + row.append(int(token)) + checked.append(row) + cells = codec.decode_batch(checked, self.action_dim, sum(1 << d for d in dims)) + previous = self._dense_previous_first_order(previous_first_order, dims, len(checked)) + decoded = self.physical.decode_batch( + self._fill_inactive_cells(cells), previous_first_order=previous + ) + return np.take(decoded, dims, axis=2) + + physical.ActionCodec2Processor.encode_batch = encode_batch + physical.ActionCodec2Processor.decode_batch = decode_batch + + exporting = importlib.import_module(f"{package}.integration.exporting") + for name in ("hub_acceleration.py", "hub_entry.py"): + if name not in exporting.MODULE_FILES: + exporting.MODULE_FILES.append(name) + original_export = exporting.ArtifactExporter.export + + @functools.wraps(original_export) + def export(self): + import shutil + + files = original_export(self) + source = self.package_root.parent + for name in ("setup.py", "pyproject.toml", "LICENSE"): + path = source / name + target = self.output / name + if path.is_file() and path.resolve() != target.resolve(): + shutil.copyfile(path, target) + files.append(str(target)) + kernels = source / "actioncodec2_kernels" + if kernels.is_dir() and kernels.resolve() != (self.output / kernels.name).resolve(): + shutil.copytree(kernels, self.output / kernels.name, dirs_exist_ok=True, + ignore=shutil.ignore_patterns("__pycache__", "*.so", "*.pyc")) + return files + + exporting.ArtifactExporter.export = export diff --git a/runtime/hub_entry.py b/runtime/hub_entry.py new file mode 100644 index 0000000000000000000000000000000000000000..ca47be24297dc12733d9cdc358f02bb663d03546 --- /dev/null +++ b/runtime/hub_entry.py @@ -0,0 +1,94 @@ +"""Small Hugging Face entry point for a grouped ActionCodec2 artifact. + +Transformers copies only top-level Python files from a Hub repository into its +dynamic-module cache. This entry point locates the artifact, then loads the +versioned runtime stored in ``runtime/`` without installing another package. +""" + +from __future__ import annotations + +import hashlib +import importlib +import importlib.util +import sys +from pathlib import Path + + +class ActionCodec2: + """Load the processor implementation bundled with a pretrained artifact.""" + + @classmethod + def register_for_auto_class(cls, auto_class="AutoProcessor"): + """Satisfy Transformers' dynamic class hook for this thin loader.""" + return cls + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): + source = Path(pretrained_model_name_or_path) + subfolder = str(kwargs.pop("subfolder", "")) + action_space = kwargs.pop("action_space", None) + kwargs.pop("_from_auto", None) + kwargs.pop("trust_remote_code", None) + allowed = { + "cache_dir", + "force_download", + "local_files_only", + "token", + "revision", + "repo_type", + } + unknown = sorted(set(kwargs) - allowed) + if unknown: + raise TypeError( + f"unsupported from_pretrained keyword(s): {', '.join(unknown)}" + ) + if source.is_dir(): + root = source / subfolder + else: + from huggingface_hub import snapshot_download + + prefix = f"{subfolder.rstrip('/')}/" if subfolder else "" + patterns = [ + "config.json", + "processor_config.json", + "router_config.yaml", + "profiles/**", + "runtime/**", + "actioncodec2_kernels/**", + "setup.py", "pyproject.toml", "LICENSE", + "README.md", + "requirements.txt", + "fit_report.json", + ] + snapshot = snapshot_download( + repo_id=str(pretrained_model_name_or_path), + allow_patterns=[prefix + pattern for pattern in patterns], + **kwargs, + ) + root = Path(snapshot) / subfolder + + runtime = root / "runtime" + package_file = runtime / "__init__.py" + if not package_file.is_file(): + raise FileNotFoundError( + f"ActionCodec2 runtime is missing from {root}; copy the whole artifact" + ) + identity = hashlib.sha256(str(runtime.resolve()).encode()).hexdigest()[:16] + package_name = f"_actioncodec2_artifact_{identity}" + if package_name not in sys.modules: + spec = importlib.util.spec_from_file_location( + package_name, package_file, submodule_search_locations=[str(runtime)] + ) + if spec is None or spec.loader is None: + raise ImportError(f"cannot load ActionCodec2 runtime from {runtime}") + package = importlib.util.module_from_spec(spec) + sys.modules[package_name] = package + try: + spec.loader.exec_module(package) + except BaseException: + del sys.modules[package_name] + raise + processor = importlib.import_module(f"{package_name}.processing_actioncodec2") + acceleration = importlib.import_module(f"{package_name}.hub_acceleration") + acceleration.install(package_name) + return processor.ActionCodec2.from_pretrained(root, action_space=action_space) diff --git a/runtime/integration/__init__.py b/runtime/integration/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3dfc34f876e95750bbb267c77ddc798e1782421a --- /dev/null +++ b/runtime/integration/__init__.py @@ -0,0 +1 @@ +"""Internal integration components for ActionCodec2.""" diff --git a/runtime/integration/exporting.py b/runtime/integration/exporting.py new file mode 100644 index 0000000000000000000000000000000000000000..26e18a12735e97225ceaf11c8d40dca51bf52585 --- /dev/null +++ b/runtime/integration/exporting.py @@ -0,0 +1,401 @@ +"""Write the portable runtime and user guide for a saved ActionCodec2 artifact.""" + +import shutil +from pathlib import Path + +# Files included in a portable runtime artifact. +MODULE_FILES: list[str] +SUPPORT_FILES: dict[str, str] +USER_GUIDE: str + + +class ArtifactExporter: + """Copy portable code alongside numerical artifacts. + + Args: + directory: Destination directory. + source_directory: Existing artifact or package directory used as a + fallback for runtime files and native resources. + """ + + def __init__( + self, directory: str | Path, *, source_directory: str | Path | None = None + ) -> None: + self.output = Path(directory) + self.package_root = Path(__file__).resolve().parent.parent + self.source = Path(source_directory) if source_directory else self.package_root + + def export(self) -> list[str]: + self.output.mkdir(parents=True, exist_ok=True) + files = [] + runtime = self.output / "runtime" + runtime.mkdir(parents=True, exist_ok=True) + for name in MODULE_FILES: + # Copy the current implementation so an artifact can be moved alone. + source = self.package_root / name + if not source.is_file(): + source = self.source / name + target = runtime / name + target.parent.mkdir(parents=True, exist_ok=True) + if source.resolve() != target.resolve(): + shutil.copyfile(source, target) + files.append(str(target)) + (runtime / "__init__.py").write_text( + '"""Private inference runtime bundled with this ActionCodec2 artifact."""\n', + encoding="utf-8", + ) + entry = self.package_root / "hub_entry.py" + if not entry.is_file(): + entry = self.source / "runtime" / "hub_entry.py" + target = self.output / "processing_actioncodec2.py" + if entry.resolve() != target.resolve(): + shutil.copyfile(entry, target) + files.append(str(target)) + # A saved artifact needs runtime dependencies, but not training CLIs or + # package-build metadata. + for name in ("requirements.txt",): + content = SUPPORT_FILES[name] + target = self.output / name + target.write_text(content, encoding="utf-8") + files.append(str(target)) + if not (self.output / "README.md").exists(): + (self.output / "README.md").write_text(USER_GUIDE, encoding="utf-8") + return files + + +USER_GUIDE = """# Loading this saved artifact + +This directory is a complete, portable ActionCodec2 artifact. Install its runtime dependencies, then point Transformers at this directory. Keep every file and subdirectory together when copying or uploading it. + +```bash +python -m pip install -r requirements.txt +``` + +```python +from transformers import AutoProcessor + +codec = AutoProcessor.from_pretrained( + "ZibinDong/ActionCodec2-1st-order", + trust_remote_code=True, +) +print(codec) +``` + +# ActionCodec2 + +ActionCodec2 converts continuous robot actions into integer token IDs and decodes +them to the same column layout. Physical quantization is lossy; token +serialization is exact for a fixed fitted codec. The Python `ActionCodec2` class +is the main interface. Its input is an action tensor, not a text prompt. + +Python 3.10+ is required. + +```bash +python -m pip install . +``` + +## Choose an action space before encoding + +Action shape alone cannot identify units or column meanings. Confirm the action +order, units, rotation representation, absolute versus delta convention, +gripper convention, and sampling rate from your data source. `D=7` can match +several different layouts. + +```python +from actioncodec2 import ActionCodec2 + +# No artifact is needed to inspect the built-in layout definitions. +print(ActionCodec2.builtin_action_spaces().keys()) + +ARTIFACT = "ZibinDong/ActionCodec2-1st-order" # or a Hub repository ID +print(ActionCodec2.describe_pretrained(ARTIFACT, action_dim=7)) +codec = ActionCodec2.from_pretrained(ARTIFACT, action_space="single_eef_delta") +print(codec) # All registered layouts, column meanings, units, and requirements +codec.print_action_spaces(action_dim=7) # Candidates only; never auto-selects +``` + +A saved artifact can include its selected `action_space`; in that case +`ActionCodec2.from_pretrained(ARTIFACT)` restores it. For an unbound router, +`action_space` is required. You can inspect `codec.action_space` for the selected +contract and `codec.action_spaces` for every registered contract. Use +`codec.for_action_space("NAME")` to create another binding over the same fitted +profiles. The original object keeps its binding. + +To define a layout with explicit semantics, register a preset mapping or YAML +file. Start from a known contract only when its column meanings really match: + +```python +preset = codec.action_space.to_dict() +preset["name"] = "my_30hz_eef_delta" +preset["source_hz"] = 30.0 +codec.register_action_space(preset) +codec = codec.for_action_space("my_30hz_eef_delta") +codec.save_pretrained("./my-bound-codec") # Persists the custom layout and binding +``` + +Changing `source_hz` does not change the meaning of action columns. For a new +robot or a different physical tensor, define the complete preset instead of +renaming an existing layout. + +## Encode and decode + +```python +import numpy as np + +actions = np.asarray(your_actions, dtype=np.float32) # (T, D) or (B, T, D) +tokens = codec.encode(actions, fps=30) # list[list[int]], one row per episode +reconstructed = codec.decode(tokens, fps=30) # CPU torch.float32, (B, T_out, D) +features = codec(actions, fps=30) # Transformers BatchFeature with input_ids +``` + +`T` can vary across calls. For unequal episode lengths, encode episodes one at a +time; a single tensor batch must have equal `T`. Pass the real sampling rate with +`fps`. If omitted, the codec warns once and assumes the selected preset's +`source_hz`. Decoding at a different rate resamples the reconstructed path, so +`T_out` may differ from the original `T`. Physical quantization and time +resampling prevent exact numeric round trips. + +For independently generated rows with different decoded horizons, use +`codec.batch_decode(rows, fps=30)`. It returns a list of tensors shaped +`(1, T_i, D)`. With `errors="return"`, an invalid row becomes a `ValueError` in +its list position without aborting other rows. Ordinary `decode` expects a +regular batch with the same decoded length. + +### Absolute actions and second-order profiles + +An absolute layout needs the physical state immediately before the first action. +Pass the same `current_state` to encode and decode; it has shape `(B, D)` even +when the input is one `(T, D)` episode: + +```python +tokens = codec.encode(actions, fps=30, current_state=current_state) +reconstructed = codec.decode(tokens, fps=30, current_state=current_state) +``` + +A second-order profile additionally needs the preceding reconstructed +first-order increment for each active component. Inspect +`codec.requires_previous_first_order`. The boundary mapping has the form +`{component_name: (B, D_active)}` at the codec's target rate. It is independent +of `current_state`. + +```python +# Only use zeros when the preceding motion is known to be zero. +boundary = codec.first_order_template(batch_size=1) +tokens = codec.encode(actions, fps=30, previous_first_order=boundary) +reconstructed = codec.decode(tokens, fps=30, previous_first_order=boundary) +next_boundary = codec.final_first_order( + tokens, previous_first_order=boundary, executed_steps=None +) +``` + +`final_first_order` returns the boundary for the next chunk. `executed_steps` +can select an executed prefix in codec-rate steps. For an absolute second-order +layout, pass both `current_state` and `previous_first_order` to encode/decode. +The library never invents a boundary from the first action. + +For diagnostics, `codec.is_fitted`, `codec.fitted_profiles`, and +`codec.requires_previous_first_order` report whether the selected artifact is +ready for inference and whether the second-order boundary is mandatory. + +## Fit and save + +Fit on actions whose columns follow the selected layout. The same `ActionCodec2` +object is the training interface: + +```python +codec = ActionCodec2(action_space="single_eef_delta") +codec.fit(episodes, fps=15, backend="auto") +print(codec.is_fitted, codec.fitted_profiles) +print(codec.training_report) +codec.save_pretrained("./my-codec") +``` + +`episodes` can be one `(T, D)` array, one regular `(B, T, D)` array, or an +iterable of variable-length episodes. For mixed layouts or sampling rates, pass +records with explicit metadata: + +```python +records = [ + {"actions": episode_a, "action_space": "single_eef_delta", "fps": 15}, + {"actions": episode_b, "action_space": "single_joint_delta", "fps": 30}, +] +codec.fit(records, backend="auto") +``` + +A record may also contain `current_state` and `previous_first_order` as required +by its layout. `fit` supports `primitive_order=2` codecs, an explicit +`bootstrap_context=True` option for two measured context frames, and +`profile_configs={"joint": ..., "eef": ...}` at construction for custom +physical profile settings. `backend` is `"auto"`, `"native"`, or `"python"`; +`threads`, `batch_size`, `progress`, `max_cell_bytes`, and `work_dir` control +resource use. Fitting replaces token assignments; re-encode datasets after +refitting. `save_pretrained` writes the complete loadable artifact, including +the layout binding and model card. `save_pretrained(path, push_to_hub=True, +repo_id="ZibinDong/my-actioncodec2")` also uploads it to a Hub repository you control. + +The optional C++ extension accelerates fitting when available. `backend="auto"` +uses it if installed and otherwise uses the Python implementation. + +### Reuse separately fitted physical profiles + +If you already have fitted joint and EEF profile directories, assemble a router +without refitting. Budgets come from the recorded training configuration unless +you supply `token_budgets`: + +```python +codec = ActionCodec2.from_profiles( + "./joint-profile", "./eef-profile", action_space="single_eef_delta" +) +codec.save_pretrained("./combined-codec") +``` + +The equivalent command is: + +```bash +actioncodec2-bundle --joint ./joint-profile --eef ./eef-profile \\ + --action-space single_eef_delta --out ./combined-codec +``` + +The command also accepts `--joint-budget` and `--eef-budget`. Omit +`--action-space` only when you intentionally want to select a layout at load +time. + +## Constrain model generation + +The codec supplies a grammar for an exact action horizon; it does not generate +actions by itself. `codec.grammar(horizon, fps=...)` has +`next_token_mask(prefix)` and `is_complete(tokens)`. A truncated or otherwise +invalid stream cannot be decoded as a complete action chunk. + +```python +horizon = 12 +grammar = codec.grammar(horizon, fps=15) +max_action_tokens = codec.max_tokens_for_horizon(horizon, fps=15) +assert grammar.is_complete(tokens) +``` + +For a Transformers model, use the adapter returned by +`codec.logits_processor(...)` in `generate`: + +```python +from transformers import LogitsProcessorList + +constraint = codec.logits_processor( + horizon, + fps=15, + prompt_length=input_ids.shape[1], + eos_token_id=model.config.eos_token_id, + token_offset=action_token_offset, +) +output = model.generate( + input_ids, + logits_processor=LogitsProcessorList([constraint]), +) +``` + +The model must already have a contiguous range of action IDs beginning at +`action_token_offset`; the codec does not modify its vocabulary. EOS and any +padding ID must lie outside that range. `codec.vocab_size` is the number of +action IDs. `codec.new_action_tokens` provides their ordered names, and +`codec.to_token_strs(one_token_row)` renders a row for inspection. Remove the +model's token offset and EOS/padding IDs before passing generated action IDs to +`codec.decode`. + +## Advanced profile commands + +The recommended training path is `codec.fit`. The following commands operate +on physical profile tensors rather than raw action-space layouts. Use them when +you intentionally need to fit or audit individual joint/EEF profiles. Their +`.npy`/`.npz` inputs must already contain canonical `(T, D_profile)` or +`(B, T, D_profile)` trajectories; `actioncodec2-fit-profile` does not perform +layout routing or source-rate conversion. + +```bash +actioncodec2-fit-profile --config ./eef-config.yaml --data ./eef.npy \\ + --output-dir ./eef-profile +actioncodec2-audit-profile stats --config ./eef-config.yaml \\ + --data ./eef.npy --output ./eef-stats.json +actioncodec2-audit-profile validate --artifact ./eef-profile \\ + --data ./eef.npy --output ./eef-validation.json +``` + +`--config` is an `ActionCodec2` physical profile YAML with a `train` section. +The fit command writes a physical profile and a JSON report. `stats` summarizes +source trajectories before fitting; `validate` checks a fitted profile against +data and writes a machine-readable report. Run each command with `--help` for +resource limits, metadata fields, and overwrite controls. + +## Public interface at a glance + +| Task | Entry point | Result | +| --- | --- | --- | +| Inspect layouts | `ActionCodec2.builtin_action_spaces()`, `describe_pretrained()`, `print(codec)`, `codec.print_action_spaces()` | Explicit physical contracts | +| Load and switch | `from_pretrained()`, `for_action_space()`, `register_action_space()` | Bound `ActionCodec2` | +| Fit and package | `ActionCodec2(...)`, `fit()`, `from_profiles()`, `save_pretrained()` | Loadable artifact | +| Encode and decode | `encode()`, `decode()`, `batch_decode()`, `codec(...)` | Token rows or action tensors | +| Stateful chunks | `first_order_template()`, `final_first_order()` | Explicit second-order boundaries | +| Generate | `grammar()`, `logits_processor()`, `max_tokens_for_horizon()` | Valid token constraints | +| Inspect tokens | `vocab_size`, `new_action_tokens`, `to_token_strs()` | ID-to-name mapping | + +`AutoProcessor.from_pretrained(path, trust_remote_code=True)` loads an artifact +without importing this package first. If the package is installed and imported, +its processor registration also supports local `AutoProcessor` resolution. + +## Tests + +```bash +python -m pytest tests -q +``` + +## Artifact contents + +A saved artifact contains the metadata, fitted profiles, private runtime, and +this guide. Keep the complete directory together when moving it or uploading it +to Hugging Face. +""" + +SUPPORT_FILES = { + "requirements.txt": "numpy\nscipy\ntorch>=2.3\ntransformers>=4.57,<6\nhuggingface-hub\npyyaml\n", +} + +MODULE_FILES = [ + "backend/__init__.py", + "backend/native_backend.py", + "common/__init__.py", + "common/_frontend_base.py", + "common/_types.py", + "common/_validation.py", + "frontend/__init__.py", + "frontend/tokenizer.py", + "integration/__init__.py", + "integration/exporting.py", + "integration/hf.py", + "physical/__init__.py", + "physical/_so3.py", + "physical/boundary.py", + "physical/configuration.py", + "physical/physical_codec.py", + "physical/physical_processor.py", + "physical/primitives.py", + "physical/semantic.py", + "physical/trajectory.py", + "routing/__init__.py", + "routing/cache.py", + "routing/defaults.py", + "routing/presets.py", + "routing/routing.py", + "tokenization/__init__.py", + "tokenization/bpe.py", + "tokenization/bpe_training.py", + "tokenization/bpe_vocabulary.py", + "tokenization/generation.py", + "tokenization/grammar.py", + "tokenization/setbpe_spatial.py", + "tokenization/setbpe_spatial_training.py", + "tokenization/setbpe_temporal.py", + "tokenization/setbpe_temporal_training.py", + "training/__init__.py", + "training/fitting.py", + "training/training.py", + "processing_actioncodec2.py", +] diff --git a/runtime/integration/hf.py b/runtime/integration/hf.py new file mode 100644 index 0000000000000000000000000000000000000000..0eaa1522374f6c196fea49f470383a3191026020 --- /dev/null +++ b/runtime/integration/hf.py @@ -0,0 +1,199 @@ +"""Hugging Face AutoProcessor integration. + +ActionCodec2 is a processor, not a model: it has no trainable ``forward`` +parameters and maps continuous trajectories to/from discrete model-facing token +IDs. The numerical artifact remains ``config.yaml``/``primitives.npz``/ +``bpe.json``; the two small JSON files here are standard Hugging Face discovery +metadata only. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from transformers import AutoConfig, AutoProcessor, PretrainedConfig + +from ..physical.configuration import MODEL_TYPE +from .exporting import USER_GUIDE + +if TYPE_CHECKING: + from ..physical.physical_processor import ActionCodec2Processor + from ..routing.routing import ActionCodec2Router + +HF_CONFIG_NAME = "config.json" +PROCESSOR_CONFIG_NAME = "processor_config.json" +PROCESSOR_CLASS_NAME = "ActionCodec2Processor" +ROUTER_MODEL_TYPE = "actioncodec2_router" +ROUTER_PROCESSOR_CLASS_NAME = "ActionCodec2Router" +ROUTER_CONFIG_NAME = "router_config.yaml" +ROUTER_ALLOW_PATTERNS = [ + ROUTER_CONFIG_NAME, + HF_CONFIG_NAME, + PROCESSOR_CONFIG_NAME, + "profiles/**", + "fit_report.json", + "README.md", + "*.py", + "requirements.txt", + "runtime/**", "actioncodec2_kernels/**", "setup.py", "pyproject.toml", "LICENSE", + # Compatibility for artifacts emitted before runtime/ was introduced. + "native/*", +] + + +class ActionCodec2FrontendHFConfig(PretrainedConfig): + """AutoConfig discovery key for the action-space-bound public processor.""" + + model_type = "actioncodec2" + + def __init__(self, **kwargs: Any) -> None: + kwargs.setdefault("processor_class", "ActionCodec2") + super().__init__(**kwargs) + + +class ActionCodec2HFConfig(PretrainedConfig): + """Minimal AutoConfig key that dispatches to :class:`ActionCodec2Processor`.""" + + model_type = MODEL_TYPE + + def __init__(self, **kwargs: Any) -> None: + kwargs.setdefault("processor_class", PROCESSOR_CLASS_NAME) + super().__init__(**kwargs) + + +class ActionCodec2RouterHFConfig(PretrainedConfig): + """Minimal AutoConfig key that dispatches to :class:`ActionCodec2Router`.""" + + model_type = ROUTER_MODEL_TYPE + + def __init__(self, **kwargs: Any) -> None: + kwargs.setdefault("processor_class", ROUTER_PROCESSOR_CLASS_NAME) + super().__init__(**kwargs) + + +def save_huggingface_metadata( + directory: str | Path, + *, + model_type: str = MODEL_TYPE, + processor_class: str = PROCESSOR_CLASS_NAME, + **processor_config: Any, +) -> tuple[Path, Path]: + """Write standard local/Hub discovery files for ``AutoProcessor``.""" + root = Path(directory) + root.mkdir(parents=True, exist_ok=True) + config_path = root / HF_CONFIG_NAME + processor_path = root / PROCESSOR_CONFIG_NAME + config_path.write_text( + json.dumps( + { + "model_type": model_type, + "processor_class": processor_class, + }, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + processor_path.write_text( + json.dumps({"processor_class": processor_class, **processor_config}, indent=2) + + "\n", + encoding="utf-8", + ) + return config_path, processor_path + + +def resolve_pretrained_directory( + pretrained_model_name_or_path: str | Path, + *, + allow_patterns: list[str] | None = None, + **kwargs: Any, +) -> Path: + """Resolve a local directory or download one immutable Hub snapshot.""" + source = Path(pretrained_model_name_or_path) + kwargs.pop("_from_auto", None) + kwargs.pop("trust_remote_code", None) + subfolder = str(kwargs.pop("subfolder", "")) + allowed = { + "cache_dir", + "force_download", + "local_files_only", + "token", + "revision", + "repo_type", + } + unknown = sorted(set(kwargs) - allowed) + if unknown: + raise TypeError(f"unsupported from_pretrained keyword(s): {', '.join(unknown)}") + if source.is_dir(): + return source / subfolder + + try: + from huggingface_hub import snapshot_download + except ImportError as error: + raise ImportError( + "loading a Hub artifact requires huggingface_hub (installed with Transformers)" + ) from error + patterns = allow_patterns or [ + "config.yaml", + "primitives.npz", + "bpe.json", + HF_CONFIG_NAME, + PROCESSOR_CONFIG_NAME, + ] + if subfolder: + prefix = subfolder.rstrip("/") + patterns = [f"{prefix}/{pattern}" for pattern in patterns] + snapshot = snapshot_download( + repo_id=str(pretrained_model_name_or_path), + allow_patterns=patterns, + **kwargs, + ) + return Path(snapshot) / subfolder + + +def upload_pretrained_directory(directory: Path, **kwargs: Any) -> None: + """Upload a complete, already-saved processor artifact to the Hub.""" + from huggingface_hub import HfApi + + repo_id = kwargs.pop("repo_id", directory.name) + commit_message = kwargs.pop("commit_message", "Upload ActionCodec2 processor") + api = HfApi(token=kwargs.pop("token", None)) + api.create_repo(repo_id=repo_id, exist_ok=True, **kwargs) + api.upload_folder( + repo_id=repo_id, + folder_path=str(directory), + commit_message=commit_message, + ) + + +def register_actioncodec2_auto_processors( + processor_class: type[ActionCodec2Processor], + router_class: type[ActionCodec2Router], +) -> bool: + """Register both installed ActionCodec2 artifact types with Transformers. + + Registration is idempotent and never changes the serialized numerical artifact. + """ + AutoConfig.register(MODEL_TYPE, ActionCodec2HFConfig, exist_ok=True) + AutoProcessor.register(ActionCodec2HFConfig, processor_class, exist_ok=True) + AutoConfig.register(ROUTER_MODEL_TYPE, ActionCodec2RouterHFConfig, exist_ok=True) + AutoProcessor.register(ActionCodec2RouterHFConfig, router_class, exist_ok=True) + return True + + +def save_model_card(directory: str | Path, description: str) -> Path: + """Create an informative Hub card on first save, preserving existing user cards.""" + path = Path(directory) / "README.md" + if not path.exists(): + path.parent.mkdir(parents=True, exist_ok=True) + guide = USER_GUIDE + path.write_text( + guide + + "\n\n## Artifact-specific report\n\n```text\n" + + description + + "\n```\n", + encoding="utf-8", + ) + return path diff --git a/runtime/physical/__init__.py b/runtime/physical/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c6eb6a1f9be6571d42b792c32c4a04a12ccdb793 --- /dev/null +++ b/runtime/physical/__init__.py @@ -0,0 +1 @@ +"""Internal physical components for ActionCodec2.""" diff --git a/runtime/physical/_so3.py b/runtime/physical/_so3.py new file mode 100644 index 0000000000000000000000000000000000000000..286467ac301d1572811b2ef127975ef157a93fc8 --- /dev/null +++ b/runtime/physical/_so3.py @@ -0,0 +1,99 @@ +"""Shared NumPy quaternion operations using the xyzw convention.""" + +from __future__ import annotations + +import numpy as np + + +def normalize_quaternion(values: np.ndarray) -> np.ndarray: + """Normalize quaternion rows and reject an undefined zero rotation.""" + quaternions = np.asarray(values, dtype=np.float64) + norms = np.linalg.norm(quaternions, axis=-1, keepdims=True) + if np.any(norms < 1e-12): + raise ValueError("a pose contains a zero quaternion") + return quaternions / norms + + +def multiply_quaternions(left: np.ndarray, right: np.ndarray) -> np.ndarray: + """Return the Hamilton product of broadcastable xyzw arrays.""" + first = np.asarray(left, dtype=np.float64) + second = np.asarray(right, dtype=np.float64) + first_vector, first_scalar = first[..., :3], first[..., 3:4] + second_vector, second_scalar = second[..., :3], second[..., 3:4] + return normalize_quaternion( + np.concatenate( + ( + first_scalar * second_vector + + second_scalar * first_vector + + np.cross(first_vector, second_vector), + first_scalar * second_scalar + - np.sum(first_vector * second_vector, axis=-1, keepdims=True), + ), + axis=-1, + ) + ) + + +def conjugate_quaternion(quaternion: np.ndarray) -> np.ndarray: + """Return the conjugate of an xyzw quaternion array.""" + result = np.asarray(quaternion, dtype=np.float64).copy() + result[..., :3] *= -1.0 + return result + + +def rotvec_to_quaternion(rotvec: np.ndarray) -> np.ndarray: + """Convert rotation vectors to xyzw quaternions.""" + values = np.asarray(rotvec, dtype=np.float64) + angle = np.linalg.norm(values, axis=-1) + scale = np.full_like(angle, 0.5) + regular = angle >= 1e-12 + scale[regular] = np.sin(0.5 * angle[regular]) / angle[regular] + return normalize_quaternion( + np.concatenate( + (values * scale[..., None], np.cos(0.5 * angle)[..., None]), axis=-1 + ) + ) + + +def quaternion_to_rotvec(quaternion: np.ndarray) -> np.ndarray: + """Convert xyzw quaternions to shortest-path rotation vectors.""" + values = normalize_quaternion(quaternion).copy() + values[values[..., 3] < 0.0] *= -1.0 + vector_norm = np.linalg.norm(values[..., :3], axis=-1) + angle = 2.0 * np.arctan2(vector_norm, np.clip(values[..., 3], 0.0, 1.0)) + scale = np.full_like(angle, 2.0) + regular = vector_norm >= 1e-12 + scale[regular] = angle[regular] / vector_norm[regular] + return values[..., :3] * scale[..., None] + + +def euler_xyz_to_quaternion(euler: np.ndarray) -> np.ndarray: + """Convert ``(N,3)`` XYZ Euler angles to xyzw quaternions.""" + values = np.asarray(euler, dtype=np.float64) + if values.ndim != 2 or values.shape[1] != 3: + raise ValueError("XYZ Euler angles must have shape [N, 3]") + half = 0.5 * values + sin_roll, sin_pitch, sin_yaw = np.sin(half).T + cos_roll, cos_pitch, cos_yaw = np.cos(half).T + result = np.empty((len(values), 4), dtype=np.float64) + result[:, 0] = sin_roll * cos_pitch * cos_yaw - cos_roll * sin_pitch * sin_yaw + result[:, 1] = cos_roll * sin_pitch * cos_yaw + sin_roll * cos_pitch * sin_yaw + result[:, 2] = cos_roll * cos_pitch * sin_yaw - sin_roll * sin_pitch * cos_yaw + result[:, 3] = cos_roll * cos_pitch * cos_yaw + sin_roll * sin_pitch * sin_yaw + return normalize_quaternion(result) + + +def quaternion_to_euler_xyz(quaternion: np.ndarray) -> np.ndarray: + """Convert ``(N,4)`` xyzw quaternions to XYZ Euler angles.""" + values = normalize_quaternion(quaternion) + if values.ndim != 2 or values.shape[1] != 4: + raise ValueError("xyzw quaternions must have shape [N, 4]") + x, y, z, w = values.T + return np.stack( + ( + np.arctan2(2.0 * (w * x + y * z), 1.0 - 2.0 * (x * x + y * y)), + np.arcsin(np.clip(2.0 * (w * y - z * x), -1.0, 1.0)), + np.arctan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z)), + ), + axis=1, + ) diff --git a/runtime/physical/boundary.py b/runtime/physical/boundary.py new file mode 100644 index 0000000000000000000000000000000000000000..20835d6008f27ba7938155c24c24ef0f35d61224 --- /dev/null +++ b/runtime/physical/boundary.py @@ -0,0 +1,153 @@ +"""Explicit measured-context preparation for stateful trajectory fitting.""" + +from __future__ import annotations + +import numpy as np + +from .configuration import ActionCodec2Config +from .physical_codec import _scale, _SO3Runtime +from .trajectory import ( + ActiveTrajectory, + TrajectoryLike, + prepare_trajectory, + trajectory_boundary, +) + + +def source_differences( + trajectory: TrajectoryLike, config: ActionCodec2Config +) -> dict[str, tuple[tuple[int, ...], np.ndarray, np.ndarray]]: + """Measure physical first/second differences with the declared group law. + + Returns component names mapped to global dims, ``[T,d]`` first-order + values, and second-order values. If no first-order boundary is supplied + for a first-order profile, the first second difference is unobserved and + omitted. Second-order profiles require their explicit boundary. + """ + values, dims = prepare_trajectory( + trajectory, dimension=config.physical.input_dim, index=0 + ) + boundary = trajectory_boundary(trajectory, config, dims) + output = {} + for name, component in config.physical.components.items(): + present = tuple(d for d in component.dims if d in dims) + if not present or component.primitive_order == 0: + continue + local = [dims.index(d) for d in present] + indices = [component.dims.index(d) for d in present] + scale = np.take(_scale(component), indices) + raw = values[:, local].astype(np.float64) * scale + previous = ( + None + if boundary is None + else boundary[list(present)].astype(np.float64) * scale + ) + if component.space == "additive": + if component.input_order == 0: + first = np.diff(raw, axis=0, prepend=np.zeros((1, len(present)))) + elif component.input_order == 1: + first = raw + else: + if previous is None: + raise ValueError( + "order-2 input statistics require previous_first_order" + ) + first = np.cumsum(raw, axis=0) + previous + second = ( + np.diff(first, axis=0) + if previous is None + else np.diff(first, axis=0, prepend=previous[None]) + ) + else: + if len(present) != 3: + raise ValueError( + "differential statistics require complete SO(3) components" + ) + runtime = _SO3Runtime(component, (np.zeros(1),) * 3) + quaternion = runtime._from_representation(raw) + if component.input_order == 0: + anchors = np.concatenate( + (np.array([[0.0, 0.0, 0.0, 1.0]]), quaternion[:-1]) + ) + increments = runtime._residual(quaternion, anchors) + else: + increments = quaternion + first = runtime._to_representation(increments) + if previous is None: + changes = runtime._residual(increments[1:], increments[:-1]) + else: + anchors = np.concatenate( + (runtime._from_representation(previous[None]), increments[:-1]) + ) + changes = runtime._residual(increments, anchors) + second = runtime._to_representation(changes) + output[name] = present, first, second + return output + + +def bootstrap_second_order_trajectory( + trajectory: ActiveTrajectory, config: ActionCodec2Config +) -> ActiveTrajectory: + """Use two observed positions as context and encode the remaining suffix. + + Args: + trajectory: Canonical ``[T,D_active]`` absolute trajectory, ``T >= 3``. + The first two frames establish an exact measured first-order + boundary; they are context, not encoded targets. + config: Physical profile with order-0 input components. + + Returns: + ActiveTrajectory: ``[T-2,D_active]`` local trajectory anchored at the + second observed frame, with explicit ``previous_first_order``. This is + an offline initialization helper, not a replacement for carrying the + reconstructed boundary between deployment chunks. + """ + values, dims = prepare_trajectory( + trajectory, dimension=config.physical.input_dim, index=0 + ) + if len(values) < 3: + raise ValueError( + "second-order bootstrap requires two context frames and at least one target" + ) + if trajectory.previous_first_order is not None: + raise ValueError( + "trajectory already has previous_first_order; do not bootstrap it again" + ) + result = values[2:].copy() + previous = np.zeros(len(dims), dtype=np.float32) + for component in config.physical.components.values(): + local = [dims.index(d) for d in component.dims if d in dims] + if not local: + continue + if component.input_order != 0: + raise ValueError("measured-context bootstrap requires input_order=0") + if component.primitive_order == 0: + continue + if component.space == "additive": + result[:, local] -= values[1, local] + previous[local] = values[1, local] - values[0, local] + else: + if len(local) != 3: + raise ValueError("bootstrap requires complete SO(3) components") + runtime = _SO3Runtime(component, (np.zeros(1),) * 3) + quaternion = runtime._from_representation( + values[:, local] * _scale(component) + ) + increment = runtime._residual(quaternion[1:2], quaternion[:1]) + # Rebase by a fixed group action that preserves the configured + # increment convention (body/right versus world/left). + from ._so3 import conjugate_quaternion, multiply_quaternions + + anchor = conjugate_quaternion(quaternion[1:2]) + local_path = ( + multiply_quaternions(quaternion[2:], anchor) + if runtime.left + else multiply_quaternions(anchor, quaternion[2:]) + ) + result[:, local] = runtime._to_representation(local_path) / _scale( + component + ) + previous[local] = runtime._to_representation(increment)[0] / _scale( + component + ) + return ActiveTrajectory(result, dims, previous_first_order=previous) diff --git a/runtime/physical/configuration.py b/runtime/physical/configuration.py new file mode 100644 index 0000000000000000000000000000000000000000..24f8e179737e91d8c35c31e664a2541b4cc01423 --- /dev/null +++ b/runtime/physical/configuration.py @@ -0,0 +1,794 @@ +"""Validated ``config.yaml`` model contract for ActionCodec2.""" + +from __future__ import annotations + +import math +import re +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType + +import yaml + +from ..common._validation import validate_dimension_partition + +CONFIG_NAME = "config.yaml" +# On-disk format identifier; retained because it participates in artifact fingerprints. +MODEL_TYPE = "actioncodec2_unified" +CONFIG_VERSION = 1 +_SPACES = ("additive", "so3") +_SO3_REPRESENTATIONS = ("rotation_vector", "euler_xyz") +_SO3_MULTIPLICATIONS = ("left", "right") +_PRIMITIVE_SELECTIONS = ("nearest", "smooth") +_ORDER_KEY = re.compile(r"order_(0|[1-9][0-9]*)\Z") + +NumberOrVector = float | tuple[float, ...] + + +class ConfigError(ValueError): + """The serialized model definition is incomplete or inconsistent.""" + + +class _UniqueKeyLoader(yaml.SafeLoader): + """Safe YAML loader that rejects duplicate mapping keys.""" + + +def _construct_unique_mapping( + loader: _UniqueKeyLoader, node: yaml.MappingNode, deep: bool = False +) -> dict[object, object]: + loader.flatten_mapping(node) + output: dict[object, object] = {} + for key_node, value_node in node.value: + key = loader.construct_object(key_node, deep=deep) + try: + duplicate = key in output + except TypeError as error: + raise ConfigError("config.yaml contains an invalid mapping key") from error + if duplicate: + raise ConfigError(f"config.yaml contains duplicate field {key!r}") + output[key] = loader.construct_object(value_node, deep=deep) + return output + + +_UniqueKeyLoader.add_constructor( + yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, + _construct_unique_mapping, +) + + +def _mapping(value: object, path: str) -> Mapping[str, object]: + if not isinstance(value, Mapping): + raise ConfigError(f"{path} must be a mapping") + if not all(isinstance(key, str) for key in value): + raise ConfigError(f"every field in {path} must be a string") + return value + + +def _reject_unknown( + payload: Mapping[str, object], allowed: set[str], path: str +) -> None: + unknown = sorted(set(payload) - allowed) + if unknown: + rendered = ", ".join(unknown) + raise ConfigError(f"{path} contains unknown field(s): {rendered}") + + +def _required(payload: Mapping[str, object], key: str, path: str) -> object: + if key not in payload: + raise ConfigError(f"{path} is missing required field {key!r}") + return payload[key] + + +def _integer(value: object, path: str, *, positive: bool = False) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise ConfigError(f"{path} must be an integer") + minimum = 1 if positive else 0 + if value < minimum: + qualifier = "positive" if positive else "non-negative" + raise ConfigError(f"{path} must be {qualifier}") + return int(value) + + +def _number(value: object, path: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ConfigError(f"{path} must be a number") + result = float(value) + if not math.isfinite(result): + raise ConfigError(f"{path} must be finite") + return result + + +def _number_or_vector(value: object, width: int, path: str) -> NumberOrVector: + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + result = tuple( + _number(item, f"{path}[{index}]") for index, item in enumerate(value) + ) + if len(result) != width: + raise ConfigError(f"{path} must be a scalar or contain {width} values") + return result + return _number(value, path) + + +def _expanded(value: NumberOrVector, width: int) -> tuple[float, ...]: + if isinstance(value, tuple): + return value + return (value,) * width + + +def _serialize_number_or_vector(value: NumberOrVector) -> float | list[float]: + return list(value) if isinstance(value, tuple) else value + + +def _parse_dims(value: object, input_dim: int, path: str) -> tuple[int, ...]: + if isinstance(value, Mapping): + payload = _mapping(value, path) + _reject_unknown(payload, {"start", "stop"}, path) + start = _integer(_required(payload, "start", path), f"{path}.start") + stop = _integer(_required(payload, "stop", path), f"{path}.stop") + if stop <= start: + raise ConfigError(f"{path}.stop must be greater than {path}.start") + dims = tuple(range(start, stop)) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + dims = tuple( + _integer(item, f"{path}[{index}]") for index, item in enumerate(value) + ) + if not dims: + raise ConfigError(f"{path} must not be empty") + else: + raise ConfigError(f"{path} must be a list or a start/stop range") + + if len(set(dims)) != len(dims): + raise ConfigError(f"{path} contains a repeated dimension") + outside = [dimension for dimension in dims if dimension >= input_dim] + if outside: + raise ConfigError( + f"{path} contains dimension {outside[0]}, outside input_dim={input_dim}" + ) + return dims + + +def _serialize_dims(dims: tuple[int, ...]) -> object: + if len(dims) > 1 and dims == tuple(range(dims[0], dims[-1] + 1)): + return {"start": dims[0], "stop": dims[-1] + 1} + return list(dims) + + +@dataclass(frozen=True) +class ConstraintConfig: + """Allowed inclusive interval for one differential order.""" + + minimum: NumberOrVector + maximum: NumberOrVector + + @classmethod + def from_dict(cls, payload: object, *, width: int, path: str) -> ConstraintConfig: + values = _mapping(payload, path) + _reject_unknown(values, {"min", "max"}, path) + minimum = _number_or_vector( + _required(values, "min", path), width, f"{path}.min" + ) + maximum = _number_or_vector( + _required(values, "max", path), width, f"{path}.max" + ) + if any( + lower > upper + for lower, upper in zip( + _expanded(minimum, width), _expanded(maximum, width), strict=True + ) + ): + raise ConfigError(f"{path}.min must not exceed {path}.max") + return cls(minimum=minimum, maximum=maximum) + + def to_dict(self) -> dict[str, object]: + return { + "min": _serialize_number_or_vector(self.minimum), + "max": _serialize_number_or_vector(self.maximum), + } + + +@dataclass(frozen=True) +class InputTransformConfig: + """Input-unit conversion and optional source-side clipping.""" + + scale: NumberOrVector = 1.0 + clip: bool = False + + @classmethod + def from_dict( + cls, payload: object | None, *, width: int, path: str + ) -> InputTransformConfig: + if payload is None: + return cls() + values = _mapping(payload, path) + _reject_unknown(values, {"scale", "clip"}, path) + scale = _number_or_vector(values.get("scale", 1.0), width, f"{path}.scale") + if any(value <= 0.0 for value in _expanded(scale, width)): + raise ConfigError(f"{path}.scale must be positive") + clip = values.get("clip", False) + if not isinstance(clip, bool): + raise ConfigError(f"{path}.clip must be true or false") + return cls(scale=scale, clip=clip) + + def to_dict(self) -> dict[str, object]: + return { + "scale": _serialize_number_or_vector(self.scale), + "clip": self.clip, + } + + +def _parse_component_geometry( + values: Mapping[str, object], *, input_dim: int, path: str +) -> tuple[tuple[int, ...], str, int, int, str | None, str | None]: + dims = _parse_dims(_required(values, "dims", path), input_dim, f"{path}.dims") + space = _required(values, "space", path) + if space not in _SPACES: + raise ConfigError(f"{path}.space must be one of {_SPACES}") + input_order = _integer( + _required(values, "input_order", path), f"{path}.input_order" + ) + primitive_order = _integer( + _required(values, "primitive_order", path), f"{path}.primitive_order" + ) + representation = values.get("representation") + multiplication = values.get("multiplication") + if space == "so3": + if len(dims) != 3: + raise ConfigError(f"{path} with space='so3' must contain exactly 3 dims") + if input_order not in {0, 1} or primitive_order not in {1, 2}: + raise ConfigError( + f"{path} with space='so3' supports input_order 0 or 1 and primitive_order 1 or 2" + ) + if representation not in _SO3_REPRESENTATIONS: + raise ConfigError( + f"{path}.representation must be one of {_SO3_REPRESENTATIONS}" + ) + if multiplication not in _SO3_MULTIPLICATIONS: + raise ConfigError( + f"{path}.multiplication must be one of {_SO3_MULTIPLICATIONS}" + ) + else: + if input_order not in {0, 1, 2} or primitive_order not in {0, 1, 2}: + raise ConfigError( + f"{path} with space='additive' supports input_order and primitive_order 0, 1, or 2" + ) + if representation is not None or multiplication is not None: + raise ConfigError( + f"{path}.representation and multiplication are only valid for space='so3'" + ) + return ( + dims, + str(space), + input_order, + primitive_order, + representation, + multiplication, + ) + + +def _parse_component_constraints( + values: Mapping[str, object], *, width: int, path: str +) -> dict[str, ConstraintConfig]: + raw = _mapping(_required(values, "constraints", path), f"{path}.constraints") + if not raw: + raise ConfigError(f"{path}.constraints must not be empty") + constraints: dict[str, ConstraintConfig] = {} + for key, constraint in raw.items(): + if _ORDER_KEY.fullmatch(key) is None: + raise ConfigError( + f"{path}.constraints field {key!r} must be named order_0, order_1, ..." + ) + constraints[key] = ConstraintConfig.from_dict( + constraint, width=width, path=f"{path}.constraints.{key}" + ) + return constraints + + +@dataclass(frozen=True) +class ComponentConfig: + """One orthogonal physical component of the action vector. + + ``dims`` defines the only global slots this component owns. ``input_order`` + describes the dataset boundary, ``primitive_order`` describes the decoder + recursion, constraints are expressed after ``input_transform.scale``, and + ``distortion_limit`` is coordinate-wise for additive spaces but one + geodesic angle for SO(3). + """ + + dims: tuple[int, ...] + space: str + input_order: int + primitive_order: int + input_transform: InputTransformConfig + constraints: Mapping[str, ConstraintConfig] + distortion_limit: NumberOrVector + primitive_selection: str + representation: str | None = None + multiplication: str | None = None + + @classmethod + def from_dict( + cls, + payload: object, + *, + input_dim: int, + path: str, + ) -> ComponentConfig: + values = _mapping(payload, path) + allowed = { + "dims", + "space", + "input_order", + "primitive_order", + "input_transform", + "constraints", + "distortion_limit", + "primitive_selection", + "representation", + "multiplication", + } + _reject_unknown(values, allowed, path) + dims, space, input_order, primitive_order, representation, multiplication = ( + _parse_component_geometry(values, input_dim=input_dim, path=path) + ) + width = len(dims) + constraints = _parse_component_constraints(values, width=width, path=path) + + input_transform = InputTransformConfig.from_dict( + values.get("input_transform"), + width=width, + path=f"{path}.input_transform", + ) + if input_transform.clip and f"order_{input_order}" not in constraints: + raise ConfigError( + f"{path}.input_transform.clip requires constraints.order_{input_order}" + ) + + distortion_limit = _number_or_vector( + _required(values, "distortion_limit", path), + width, + f"{path}.distortion_limit", + ) + if any(value <= 0.0 for value in _expanded(distortion_limit, width)): + raise ConfigError(f"{path}.distortion_limit must be positive") + if space == "so3" and isinstance(distortion_limit, tuple): + raise ConfigError( + f"{path}.distortion_limit must be one scalar geodesic angle for SO(3)" + ) + + primitive_selection = _required(values, "primitive_selection", path) + if primitive_selection not in _PRIMITIVE_SELECTIONS: + raise ConfigError(f"{path}.primitive_selection must be nearest or smooth") + if ( + primitive_selection == "smooth" + and space == "additive" + and primitive_order not in {1, 2} + ): + raise ConfigError( + f"{path}.primitive_selection='smooth' requires first- or second-order additive primitives" + ) + if primitive_order == 2 and "order_2" not in constraints: + raise ConfigError( + f"{path} second-order primitives require constraints.order_2" + ) + return cls( + dims=dims, + space=space, + input_order=input_order, + primitive_order=primitive_order, + input_transform=input_transform, + constraints=MappingProxyType(constraints), + distortion_limit=distortion_limit, + primitive_selection=str(primitive_selection), + representation=None if representation is None else str(representation), + multiplication=None if multiplication is None else str(multiplication), + ) + + def to_dict(self) -> dict[str, object]: + payload: dict[str, object] = { + "dims": _serialize_dims(self.dims), + "space": self.space, + "input_order": self.input_order, + "primitive_order": self.primitive_order, + "input_transform": self.input_transform.to_dict(), + "constraints": { + key: constraint.to_dict() + for key, constraint in self.constraints.items() + }, + "distortion_limit": _serialize_number_or_vector(self.distortion_limit), + "primitive_selection": self.primitive_selection, + } + if self.space == "so3": + payload["representation"] = self.representation + payload["multiplication"] = self.multiplication + return payload + + +@dataclass(frozen=True) +class PhysicalConfig: + """Complete, non-overlapping component partition of the input's last axis.""" + + input_dim: int + components: Mapping[str, ComponentConfig] + + @classmethod + def from_dict(cls, payload: object, *, path: str = "physical") -> PhysicalConfig: + values = _mapping(payload, path) + _reject_unknown(values, {"input_dim", "components"}, path) + input_dim = _integer( + _required(values, "input_dim", path), f"{path}.input_dim", positive=True + ) + raw_components = _mapping( + _required(values, "components", path), f"{path}.components" + ) + if not raw_components: + raise ConfigError(f"{path}.components must not be empty") + + components: dict[str, ComponentConfig] = {} + for name, component_payload in raw_components.items(): + if not name.strip(): + raise ConfigError(f"{path}.components contains an empty component name") + component = ComponentConfig.from_dict( + component_payload, + input_dim=input_dim, + path=f"{path}.components.{name}", + ) + components[name] = component + + validate_dimension_partition( + ((name, component.dims) for name, component in components.items()), + input_dim, + label="physical dimension", + error_type=ConfigError, + ) + return cls(input_dim=input_dim, components=MappingProxyType(components)) + + def to_dict(self) -> dict[str, object]: + return { + "input_dim": self.input_dim, + "components": { + name: component.to_dict() for name, component in self.components.items() + }, + } + + +TOKENIZATION_SCHEMES = ("set_bpe", "bpe") + + +@dataclass(frozen=True) +class SetBPEConfig: + """Runtime settings of the tokenization stage. + + ``scheme`` selects between rooted spatial-temporal Set-BPE (``"set_bpe"``, + the production default) and plain BPE over the row-major flattened cell + stream (``"bpe"``), the FAST-style ablation baseline. Both schemes share + the primitive atoms, the artifact layout, and the vocabulary budget. + + ``lookahead`` caps a Set-BPE temporal token at that many control steps and + a plain-BPE token at ``lookahead * physical.input_dim`` primitive cells. + + ``allow_cross_component_merge`` is frozen into the artifact. ``False`` + makes each arm/body component a hard spatial and temporal BPE boundary, + while all components still contribute pooled frequencies to the same + joint or EEF vocabulary. + """ + + lookahead: int + allow_cross_component_merge: bool = False + scheme: str = "set_bpe" + + @classmethod + def from_dict(cls, payload: object, *, path: str = "set_bpe") -> SetBPEConfig: + values = _mapping(payload, path) + _reject_unknown( + values, {"lookahead", "allow_cross_component_merge", "scheme"}, path + ) + lookahead = _integer( + _required(values, "lookahead", path), f"{path}.lookahead", positive=True + ) + allow_cross = values.get("allow_cross_component_merge", False) + if not isinstance(allow_cross, bool): + raise ConfigError( + f"{path}.allow_cross_component_merge must be true or false" + ) + scheme = values.get("scheme", "set_bpe") + if scheme not in TOKENIZATION_SCHEMES: + raise ConfigError( + f"{path}.scheme must be one of {', '.join(TOKENIZATION_SCHEMES)}" + ) + return cls( + lookahead=lookahead, allow_cross_component_merge=allow_cross, scheme=scheme + ) + + def to_dict(self) -> dict[str, object]: + payload: dict[str, object] = { + "lookahead": self.lookahead, + "allow_cross_component_merge": self.allow_cross_component_merge, + } + # Omitted for the default so every already-published artifact keeps its + # fingerprint; the field only appears on ablation artifacts. + if self.scheme != "set_bpe": + payload["scheme"] = self.scheme + return payload + + +@dataclass(frozen=True) +class PrimitiveTrainingConfig: + """Serialized config-v1 metadata retained until artifacts are versioned again. + + Fitting forces exact zero on every primitive axis regardless of this list. + Keeping strict parsing here preserves the currently published artifact + fingerprint; the next schema version can remove the field and republish. + """ + + zero_primitive_components: tuple[str, ...] + + @classmethod + def from_dict( + cls, + payload: object, + *, + physical: PhysicalConfig, + path: str = "train.primitives", + ) -> PrimitiveTrainingConfig: + values = _mapping(payload, path) + _reject_unknown(values, {"zero_primitive_components"}, path) + raw_names = _required(values, "zero_primitive_components", path) + if not isinstance(raw_names, Sequence) or isinstance( + raw_names, (str, bytes, bytearray) + ): + raise ConfigError(f"{path}.zero_primitive_components must be a list") + names = tuple(raw_names) + if not all(isinstance(name, str) and name for name in names): + raise ConfigError( + f"{path}.zero_primitive_components must contain non-empty component names" + ) + if len(set(names)) != len(names): + raise ConfigError(f"{path}.zero_primitive_components contains a duplicate") + unknown = sorted(set(names) - set(physical.components)) + if unknown: + raise ConfigError( + f"{path}.zero_primitive_components contains unknown component {unknown[0]!r}" + ) + + selected = set(names) + for name, component in physical.components.items(): + if component.space == "additive": + if component.primitive_order not in {0, 1, 2}: + raise ConfigError( + f"{path} cannot train additive primitive_order=" + f"{component.primitive_order} for component {name!r}" + ) + domain_order = component.primitive_order + if name in selected and component.primitive_order == 0: + raise ConfigError( + f"{path}.zero_primitive_components cannot contain absolute " + f"component {name!r}" + ) + else: + domain_order = component.primitive_order + if f"order_{domain_order}" not in component.constraints: + raise ConfigError( + f"physical component {name!r} needs constraints.order_{domain_order} " + "to train its primitives" + ) + return cls(zero_primitive_components=names) + + def to_dict(self) -> dict[str, object]: + return {"zero_primitive_components": list(self.zero_primitive_components)} + + +@dataclass(frozen=True) +class BPETrainingConfig: + """Spatial and final vocabulary budgets used for whole-trajectory BPE. + + ``spatial_vocab_size`` splits the budget between the two Set-BPE stages and + is unused by the plain-BPE scheme, which has a single merge stage. + """ + + spatial_vocab_size: int + vocab_size: int + + @classmethod + def from_dict( + cls, payload: object, *, path: str = "train.bpe" + ) -> BPETrainingConfig: + values = _mapping(payload, path) + _reject_unknown(values, {"spatial_vocab_size", "vocab_size"}, path) + spatial_vocab_size = _integer( + _required(values, "spatial_vocab_size", path), + f"{path}.spatial_vocab_size", + positive=True, + ) + vocab_size = _integer( + _required(values, "vocab_size", path), f"{path}.vocab_size", positive=True + ) + if spatial_vocab_size > vocab_size: + raise ConfigError(f"{path}.spatial_vocab_size must not exceed vocab_size") + return cls(spatial_vocab_size=spatial_vocab_size, vocab_size=vocab_size) + + def to_dict(self) -> dict[str, object]: + return { + "spatial_vocab_size": self.spatial_vocab_size, + "vocab_size": self.vocab_size, + } + + +@dataclass(frozen=True) +class QualityControlConfig: + """Pretraining-only policy for extreme first-order discontinuities. + + ``constraints.order_1`` remains the primitive domain and therefore the + region with the configured distortion guarantee. Values outside that + domain are still encodable through saturated primitives and causal error + feedback. Only steps beyond ``hard_order_1_scale`` times the primitive + domain are treated as data discontinuities and split before BPE fitting. + """ + + hard_order_1_scale: float + min_segment_length: int + + @classmethod + def from_dict( + cls, payload: object, *, path: str = "train.quality_control" + ) -> QualityControlConfig: + values = _mapping(payload, path) + _reject_unknown(values, {"hard_order_1_scale", "min_segment_length"}, path) + scale = _number( + _required(values, "hard_order_1_scale", path), + f"{path}.hard_order_1_scale", + ) + if scale <= 1.0: + raise ConfigError(f"{path}.hard_order_1_scale must be greater than 1") + return cls( + hard_order_1_scale=scale, + min_segment_length=_integer( + _required(values, "min_segment_length", path), + f"{path}.min_segment_length", + positive=True, + ), + ) + + def to_dict(self) -> dict[str, object]: + return { + "hard_order_1_scale": self.hard_order_1_scale, + "min_segment_length": self.min_segment_length, + } + + +@dataclass(frozen=True) +class TrainConfig: + """Non-derivable fit choices and BPE budgets. + + Physical domains and distortion limits are intentionally absent: primitive + construction derives them from :class:`PhysicalConfig`, preventing two + conflicting sources of truth. + """ + + primitives: PrimitiveTrainingConfig + bpe: BPETrainingConfig + quality_control: QualityControlConfig | None = None + + @classmethod + def from_dict( + cls, + payload: object, + *, + physical: PhysicalConfig, + path: str = "train", + ) -> TrainConfig: + values = _mapping(payload, path) + _reject_unknown(values, {"primitives", "bpe", "quality_control"}, path) + raw_quality_control = values.get("quality_control") + return cls( + primitives=PrimitiveTrainingConfig.from_dict( + _required(values, "primitives", path), physical=physical + ), + bpe=BPETrainingConfig.from_dict(_required(values, "bpe", path)), + quality_control=( + None + if raw_quality_control is None + else QualityControlConfig.from_dict(raw_quality_control) + ), + ) + + def to_dict(self) -> dict[str, object]: + payload = {"primitives": self.primitives.to_dict(), "bpe": self.bpe.to_dict()} + if self.quality_control is not None: + payload["quality_control"] = self.quality_control.to_dict() + return payload + + +@dataclass(frozen=True) +class ActionCodec2Config: + """Strict YAML configuration for an ActionCodec2 processor artifact. + + ``physical`` and ``set_bpe`` are required at deployment. ``train`` is an + optional reproducibility recipe and is never consulted by encode/decode. + This is a processor configuration, not a neural ``PretrainedConfig``; + ``config.json`` exists only as Hugging Face AutoProcessor discovery metadata. + """ + + physical: PhysicalConfig + set_bpe: SetBPEConfig + train: TrainConfig | None = None + model_type: str = MODEL_TYPE + config_version: int = CONFIG_VERSION + + @classmethod + def from_dict(cls, payload: object) -> ActionCodec2Config: + values = _mapping(payload, "config") + _reject_unknown( + values, + {"model_type", "config_version", "physical", "set_bpe", "train"}, + "config", + ) + model_type = _required(values, "model_type", "config") + if model_type != MODEL_TYPE: + raise ConfigError(f"config.model_type must be {MODEL_TYPE!r}") + config_version = _integer( + _required(values, "config_version", "config"), "config.config_version" + ) + if config_version != CONFIG_VERSION: + raise ConfigError( + f"unsupported config_version {config_version}; expected {CONFIG_VERSION}" + ) + physical = PhysicalConfig.from_dict(_required(values, "physical", "config")) + raw_train = values.get("train") + return cls( + physical=physical, + set_bpe=SetBPEConfig.from_dict(_required(values, "set_bpe", "config")), + train=( + None + if raw_train is None + else TrainConfig.from_dict(raw_train, physical=physical) + ), + model_type=str(model_type), + config_version=config_version, + ) + + @classmethod + def from_yaml(cls, path: str | Path) -> ActionCodec2Config: + source = Path(path) + try: + payload = yaml.load( + source.read_text(encoding="utf-8"), Loader=_UniqueKeyLoader + ) + except OSError: + raise + except yaml.YAMLError as error: + raise ConfigError(f"cannot parse {source}: {error}") from error + return cls.from_dict(payload) + + @classmethod + def from_pretrained(cls, directory: str | Path) -> ActionCodec2Config: + """Load the strict numerical configuration from ``config.yaml``.""" + return cls.from_yaml(Path(directory) / CONFIG_NAME) + + def to_dict(self) -> dict[str, object]: + payload: dict[str, object] = { + "model_type": self.model_type, + "config_version": self.config_version, + "physical": self.physical.to_dict(), + "set_bpe": self.set_bpe.to_dict(), + } + if self.train is not None: + payload["train"] = self.train.to_dict() + return payload + + def save_yaml(self, path: str | Path) -> None: + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text( + yaml.safe_dump( + self.to_dict(), + sort_keys=False, + allow_unicode=True, + default_flow_style=False, + ), + encoding="utf-8", + ) + + def save_pretrained(self, directory: str | Path) -> None: + """Write ``config.yaml`` without embedding primitive or BPE arrays.""" + self.save_yaml(Path(directory) / CONFIG_NAME) diff --git a/runtime/physical/physical_codec.py b/runtime/physical/physical_codec.py new file mode 100644 index 0000000000000000000000000000000000000000..2bafe39483611ecda30d3e875c3cd12d815bca68 --- /dev/null +++ b/runtime/physical/physical_codec.py @@ -0,0 +1,950 @@ +"""Configuration-driven physical quantization. + +This module owns only the lossy physical stage. Components receive transformed +actions and emit one primitive cell per global action dimension; Set-BPE never +appears here. Batch dimensions are vectorized, while the short time recursion +remains explicit because error feedback is inherently causal. +""" + +from __future__ import annotations + +import itertools +import os +from typing import Protocol + +import numpy as np + +from ._so3 import conjugate_quaternion as _conjugate +from ._so3 import euler_xyz_to_quaternion as _euler_xyz_to_quaternion +from ._so3 import multiply_quaternions as _multiply +from ._so3 import quaternion_to_euler_xyz as _quaternion_to_euler_xyz +from ._so3 import quaternion_to_rotvec as _quaternion_to_rotvec +from ._so3 import rotvec_to_quaternion as _rotvec_to_quaternion +from ..common._validation import broadcast_vector +from .configuration import ActionCodec2Config, ComponentConfig, ConfigError +from ..backend.native_backend import require_native, resolve_backend +from .primitives import PrimitiveStore + +_FLOAT32_SLACK = np.float32(8 * np.finfo(np.float32).eps) +_IDENTITY = np.asarray((0.0, 0.0, 0.0, 1.0), dtype=np.float64) + + +def _constraint( + component: ComponentConfig, order: int +) -> tuple[np.ndarray, np.ndarray] | None: + value = component.constraints.get(f"order_{order}") + if value is None: + return None + return ( + np.asarray(value.minimum, dtype=np.float32), + np.asarray(value.maximum, dtype=np.float32), + ) + + +def _scale(component: ComponentConfig) -> np.ndarray: + return broadcast_vector( + component.input_transform.scale, + len(component.dims), + dtype=np.float32, + name="input_transform.scale", + ) + + +def _transform(component: ComponentConfig, actions: np.ndarray) -> np.ndarray: + """Apply the component's unit conversion and encoder-only clipping.""" + values = np.asarray(actions, dtype=np.float32) * _scale(component) + if component.input_transform.clip: + bounds = _constraint(component, component.input_order) + if bounds is None: # defensive: configuration validation normally catches this + raise ConfigError( + f"clipping requires constraints.order_{component.input_order}" + ) + values = np.clip(values, bounds[0], bounds[1]) + return np.asarray(values, dtype=np.float32) + + +def _integrate_additive(values: np.ndarray, order: int) -> np.ndarray: + """Integrate along the penultimate (time) axis from a zero boundary.""" + result = np.asarray(values, dtype=np.float32) + for _ in range(order): + result = np.cumsum(result, axis=-2, dtype=np.float32) + return result + + +def _difference_additive(values: np.ndarray, order: int) -> np.ndarray: + """Take zero-anchored finite differences along the time axis.""" + result = np.asarray(values, dtype=np.float32) + for _ in range(order): + previous = np.concatenate( + (np.zeros_like(result[..., :1, :]), result[..., :-1, :]), axis=-2 + ) + result = result - previous + return result + + +class _ComponentRuntime(Protocol): + """Internal interface shared by additive and SO(3) components.""" + + def encode_batch( + self, actions: np.ndarray, previous_first_order: np.ndarray | None = None + ) -> np.ndarray: ... + + def decode_batch( + self, cells: np.ndarray, previous_first_order: np.ndarray | None = None + ) -> np.ndarray: ... + + +class _AdditiveRuntime: + """Causal Euclidean error-feedback quantizer for one component.""" + + def __init__( + self, component: ComponentConfig, primitives: tuple[np.ndarray, ...] + ) -> None: + self.config = component + self.primitives = tuple( + np.asarray(value, dtype=np.float32) for value in primitives + ) + self.width = len(component.dims) + self.order = component.primitive_order + self._second_order_axes = ( + tuple(axis.astype(np.float64) for axis in self.primitives) + if self.order == 2 + else () + ) + self._second_order_sorted = tuple( + bool(np.all(axis[1:] > axis[:-1])) for axis in self._second_order_axes + ) + self._second_order_zero_bounds = tuple( + ( + int(np.searchsorted(axis, 0)), + int(np.searchsorted(axis, 0, side="right")) - 1, + ) + if ordered + else (0, 0) + for axis, ordered in zip( + self._second_order_axes, self._second_order_sorted, strict=True + ) + ) + raw_position_bounds = _constraint(component, 0) + self.position_bounds = ( + None + if raw_position_bounds is None + else ( + np.broadcast_to(raw_position_bounds[0], (self.width,)).copy(), + np.broadcast_to(raw_position_bounds[1], (self.width,)).copy(), + ) + ) + self.limit = broadcast_vector( + component.distortion_limit, + self.width, + dtype=np.float32, + name="distortion_limit", + ) + table_width = max(len(values) for values in self.primitives) + table = np.full((self.width, table_width), np.inf, dtype=np.float32) + for coordinate, values in enumerate(self.primitives): + table[coordinate, : len(values)] = values + self._native_table = np.ascontiguousarray(table) + self._native_bins = np.asarray( + [len(values) for values in self.primitives], dtype=np.int32 + ) + self._native_binary = np.zeros(self.width, dtype=np.uint8) + + def _pick_batch(self, request: np.ndarray, direction: np.ndarray) -> np.ndarray: + """Select cells for an arbitrary leading batch shape.""" + flat_request = np.asarray(request, dtype=np.float32).reshape(-1, self.width) + flat_direction = np.asarray(direction, dtype=np.float32).reshape(-1, self.width) + choice = np.empty(flat_request.shape, dtype=np.int64) + for coordinate, primitives in enumerate(self.primitives): + distance = np.abs(flat_request[:, coordinate, None] - primitives[None, :]) + nearest = np.argmin(distance, axis=1) + if self.config.primitive_selection == "nearest": + choice[:, coordinate] = nearest + continue + feasible = distance <= self.limit[coordinate] + _FLOAT32_SLACK + preferred = feasible & ( + np.sign(primitives)[None, :] * flat_direction[:, coordinate, None] + >= 0.0 + ) + candidates = np.where( + np.any(preferred, axis=1, keepdims=True), preferred, feasible + ) + smooth = np.argmin(np.where(candidates, distance, np.inf), axis=1) + choice[:, coordinate] = np.where(np.any(feasible, axis=1), smooth, nearest) + return choice.reshape(request.shape) + + def _advance_batch(self, state: np.ndarray, primitive: np.ndarray) -> None: + """Advance all integrator states in-place by one timestep.""" + higher = primitive + for level in range(self.order - 1, -1, -1): + state[:, level] += higher + higher = state[:, level] + if self.order == 1 and self.position_bounds is not None: + state[:, 0] = np.clip( + state[:, 0], self.position_bounds[0], self.position_bounds[1] + ) + + def encode_batch( + self, actions: np.ndarray, previous_first_order: np.ndarray | None = None + ) -> np.ndarray: + """Quantize ``[B,T,d]`` actions, vectorizing all trajectories per step.""" + if self.order == 2: + return self._encode_second_order(actions, previous_first_order) + transformed = _transform(self.config, actions) + target = _integrate_additive(transformed, self.config.input_order) + batch, steps, _ = target.shape + if self.order == 0: + return self._pick_batch(target, np.zeros_like(target)) + if self.position_bounds is not None and resolve_backend("auto") == "native": + starts = np.zeros((batch, self.width), dtype=np.float32) + cells = require_native().physical_encode_batch( + self._native_table, + self._native_bins, + self._native_binary, + np.ascontiguousarray(self.position_bounds[0], dtype=np.float32), + np.ascontiguousarray(self.position_bounds[1], dtype=np.float32), + np.ascontiguousarray(target, dtype=np.float32), + starts, + starts, + self.order, + min(batch, os.cpu_count() or 1), + np.ascontiguousarray(self.limit, dtype=np.float32), + self.config.primitive_selection == "smooth", + ) + return np.asarray(cells, dtype=np.int64) + + cells = np.empty(target.shape, dtype=np.int64) + state = np.zeros((batch, self.order, self.width), dtype=np.float32) + direction = np.zeros((batch, self.width), dtype=np.float32) + zero = np.zeros((batch, self.width), dtype=np.float32) + for step in range(steps): + baseline = state.copy() + self._advance_batch(baseline, zero) + request = target[:, step] - baseline[:, 0] + choice = self._pick_batch(request, direction) + primitive = np.stack( + [ + self.primitives[index][choice[:, index]] + for index in range(self.width) + ], + axis=1, + ) + self._advance_batch(state, primitive) + moved = primitive != 0.0 + direction = np.where(moved, np.sign(primitive), direction) + cells[:, step] = choice + return cells + + def decode_batch( + self, cells: np.ndarray, previous_first_order: np.ndarray | None = None + ) -> np.ndarray: + """Decode ``[B,T,d]`` cells to actions in the component's input units.""" + if self.order == 2: + return self._decode_second_order(cells, previous_first_order) + indices = np.asarray(cells, dtype=np.int64) + if self.order == 0: + positions = np.stack( + [ + self.primitives[index][indices[:, :, index]] + for index in range(self.width) + ], + axis=2, + ) + elif self.position_bounds is not None and resolve_backend("auto") == "native": + batch = len(indices) + starts = np.zeros((batch, self.width), dtype=np.float32) + positions = require_native().physical_decode_batch( + self._native_table, + self._native_bins, + self._native_binary, + np.ascontiguousarray(self.position_bounds[0], dtype=np.float32), + np.ascontiguousarray(self.position_bounds[1], dtype=np.float32), + np.ascontiguousarray(indices, dtype=np.int32), + starts, + starts, + self.order, + min(batch, os.cpu_count() or 1), + ) + else: + picked = np.stack( + [ + self.primitives[index][indices[:, :, index]] + for index in range(self.width) + ], + axis=2, + ) + batch, steps, _ = indices.shape + positions = np.empty(indices.shape, dtype=np.float32) + state = np.zeros((batch, self.order, self.width), dtype=np.float32) + for step in range(steps): + self._advance_batch(state, picked[:, step]) + positions[:, step] = state[:, 0] + actions = _difference_additive(positions, self.config.input_order) + return np.asarray(actions / _scale(self.config), dtype=np.float32) + + def _encode_second_order( + self, actions: np.ndarray, previous_first_order: np.ndarray + ) -> np.ndarray: + """Track positions from a supplied ``[B,d]`` reconstructed displacement. + + Double integration uses float64 to avoid horizon-dependent float32 + cancellation. Source clipping never changes the shared boundary. + """ + transformed = _transform(self.config, actions).astype(np.float64) + velocity = previous_first_order.astype(np.float64) * _scale(self.config) + target = transformed + if self.config.input_order == 2: + target = np.cumsum(target, axis=1) + velocity[:, None] + if self.config.input_order >= 1: + target = np.cumsum(target, axis=1) + if resolve_backend("auto") == "native" and hasattr( + require_native(), "additive_second_order_encode_batch" + ): + return np.asarray( + require_native().additive_second_order_encode_batch( + self._native_table, + self._native_bins, + np.ascontiguousarray(target), + np.ascontiguousarray(velocity), + self.limit, + self.config.primitive_selection == "smooth", + min( + len(actions), + max( + 1, + int(os.environ.get("ACTIONCODEC2_PHYSICAL_THREADS", "16")), + ), + ), + ), + dtype=np.int64, + ) + position = np.zeros_like(velocity) + direction = np.zeros_like(velocity) + cells = np.empty(actions.shape, dtype=np.int64) + for step in range(actions.shape[1]): + request = target[:, step] - (position + velocity) + # Keep the feedback calculation in double precision, including + # distances near the boundary of the admissible error interval. + choice = self._pick_second_order(request, direction) + primitive = np.stack( + [axis[choice[:, j]] for j, axis in enumerate(self.primitives)], axis=1 + ) + velocity += primitive + position += velocity + direction = np.where(primitive != 0, np.sign(primitive), direction) + cells[:, step] = choice + return cells + + def _pick_second_order( + self, request: np.ndarray, direction: np.ndarray + ) -> np.ndarray: + choice = np.empty(request.shape, dtype=np.int64) + for j, axis in enumerate(self._second_order_axes): + # A tiny nearest-only table is faster as one NumPy reduction. + if self._second_order_sorted[j] and ( + self.config.primitive_selection == "smooth" or len(axis) > 32 + ): + right = np.searchsorted(axis, request[:, j]) + left = np.clip(right - 1, 0, len(axis) - 1) + right = np.minimum(right, len(axis) - 1) + nearest = np.where( + np.abs(request[:, j] - axis[left]) + <= np.abs(request[:, j] - axis[right]), + left, + right, + ) + # Very large requests can round several distances to the same + # float64 value. Preserve argmin's first-cell tie-break there. + tied = (nearest > 0) & ( + np.abs(request[:, j] - axis[np.maximum(nearest - 1, 0)]) + == np.abs(request[:, j] - axis[nearest]) + ) + if np.any(tied): + nearest[tied] = np.argmin( + np.abs(request[tied, j, None] - axis), axis=1 + ) + if self.config.primitive_selection == "smooth": + nonnegative, nonpositive = self._second_order_zero_bounds[j] + preferred = np.where( + direction[:, j] > 0, + np.maximum(nearest, min(nonnegative, len(axis) - 1)), + np.where( + direction[:, j] < 0, + np.minimum(nearest, max(nonpositive, 0)), + nearest, + ), + ) + admissible = (np.sign(axis[preferred]) * direction[:, j] >= 0) & ( + np.abs(request[:, j] - axis[preferred]) <= self.limit[j] + ) + nearest = np.where(admissible, preferred, nearest) + choice[:, j] = nearest + continue + # Arbitrary stored cell order retains the original ID tie-break. + distance = np.abs(request[:, j, None] - axis) + nearest = np.argmin(distance, axis=1) + if self.config.primitive_selection == "nearest": + choice[:, j] = nearest + continue + feasible = distance <= self.limit[j] + preferred = feasible & (np.sign(axis)[None] * direction[:, j, None] >= 0) + candidates = np.where( + preferred.any(axis=1, keepdims=True), preferred, feasible + ) + smooth = np.argmin(np.where(candidates, distance, np.inf), axis=1) + choice[:, j] = np.where(feasible.any(axis=1), smooth, nearest) + return choice + + def _decode_second_order( + self, cells: np.ndarray, previous_first_order: np.ndarray + ) -> np.ndarray: + """Reconstruct ``[B,T,d]`` at the configured input differential order.""" + picked = np.stack( + [axis[cells[:, :, j]] for j, axis in enumerate(self.primitives)], axis=2 + ).astype(np.float64) + if self.config.input_order == 2: + result = picked + else: + result = np.cumsum(picked, axis=1) + previous_first_order[:, None] * _scale( + self.config + ) + if self.config.input_order == 0: + result = np.cumsum(result, axis=1) + return np.asarray(result / _scale(self.config), dtype=np.float32) + + +class _SO3Runtime: + """Intrinsic rotation quantizer with left- or right-multiplicative dynamics.""" + + def __init__( + self, component: ComponentConfig, primitives: tuple[np.ndarray, ...] + ) -> None: + self.config = component + self.primitives = tuple( + np.asarray(value, dtype=np.float64) for value in primitives + ) + self.left = component.multiplication == "left" + self.limit = float(component.distortion_limit) + if component.primitive_order == 2: + bounds = component.constraints["order_2"] + corner = np.maximum(np.abs(bounds.minimum), np.abs(bounds.maximum)) + corner = np.broadcast_to(corner, (3,)) + radius = ( + np.linalg.norm(corner) + if component.representation == "rotation_vector" + else np.sum(corner) + ) + self.second_order_radius = min(np.pi, float(radius) + 3 * self.limit) + table_width = max(len(values) for values in self.primitives) + table = np.full((3, table_width), np.inf, dtype=np.float64) + for coordinate, values in enumerate(self.primitives): + table[coordinate, : len(values)] = values + self._native_table = table + self._native_bins = np.asarray( + [len(values) for values in self.primitives], dtype=np.int32 + ) + + def _from_representation(self, values: np.ndarray) -> np.ndarray: + shape = values.shape[:-1] + flat = np.asarray(values).reshape(-1, 3) + converted = ( + _rotvec_to_quaternion(flat) + if self.config.representation == "rotation_vector" + else _euler_xyz_to_quaternion(flat) + ) + return converted.reshape(*shape, 4) + + def _to_representation(self, values: np.ndarray) -> np.ndarray: + shape = values.shape[:-1] + flat = np.asarray(values).reshape(-1, 4) + converted = ( + _quaternion_to_rotvec(flat) + if self.config.representation == "rotation_vector" + else _quaternion_to_euler_xyz(flat) + ) + return converted.reshape(*shape, 3) + + def _residual(self, target: np.ndarray, reconstruction: np.ndarray) -> np.ndarray: + return ( + _multiply(target, _conjugate(reconstruction)) + if self.left + else _multiply(_conjugate(reconstruction), target) + ) + + def _apply(self, state: np.ndarray, increment: np.ndarray) -> np.ndarray: + return _multiply(increment, state) if self.left else _multiply(state, increment) + + def _nearest_batch(self, request: np.ndarray) -> np.ndarray: + return np.stack( + [ + np.argmin(np.abs(request[:, index, None] - values[None, :]), axis=1) + for index, values in enumerate(self.primitives) + ], + axis=1, + ).astype(np.int64, copy=False) + + def _smooth( + self, + request: np.ndarray, + reconstruction: np.ndarray, + target: np.ndarray, + direction: np.ndarray, + nearest: np.ndarray, + ) -> np.ndarray: + """Apply the exact admissible/non-reversing SO(3) selection rule. + + On the principal rotation-vector ball, the logarithm's co-Lipschitz + bound implies that a geodesically admissible primitive lies within + ``2*limit`` in every log coordinate. This prunes the Cartesian lattice + without the former arbitrary candidate-count cap. Every retained tuple + is then checked with the actual geodesic error. + """ + axes = [] + for coordinate, values in enumerate(self.primitives): + candidates = np.flatnonzero( + np.abs(values - request[coordinate]) <= 2.0 * self.limit + 1e-12 + ) + if nearest[coordinate] not in candidates: + candidates = np.append(candidates, nearest[coordinate]) + axes.append(candidates) + + best_any: tuple[float, tuple[int, int, int]] | None = None + best_preferred: tuple[float, tuple[int, int, int]] | None = None + for raw_choice in itertools.product(*axes): + choice = tuple(int(value) for value in raw_choice) + rotvec = np.asarray( + [self.primitives[index][choice[index]] for index in range(3)] + ) + rebuilt = self._apply(reconstruction, _rotvec_to_quaternion(rotvec)) + error = float( + np.linalg.norm(_quaternion_to_rotvec(self._residual(target, rebuilt))) + ) + if error > self.limit + 1e-12: + continue + candidate = (error, choice) + if best_any is None or candidate < best_any: + best_any = candidate + # Orthogonal moves can land a few ULP below zero depending on the + # BLAS reduction order. Treat that numerical tie as non-reversing. + if np.dot(rotvec, direction) >= -1e-15 and ( + best_preferred is None or candidate < best_preferred + ): + best_preferred = candidate + selected = best_preferred or best_any + return nearest if selected is None else np.asarray(selected[1], dtype=np.int64) + + def encode_batch( + self, actions: np.ndarray, previous_first_order: np.ndarray | None = None + ) -> np.ndarray: + """Quantize ``[B,T,3]`` rotations with a vectorized batch state.""" + if self.config.primitive_order == 2: + return self._encode_second_order(actions, previous_first_order) + transformed = _transform(self.config, actions).astype(np.float64) + if resolve_backend("auto") == "native": + cells = require_native().so3_encode_batch( + self._native_table, + self._native_bins, + np.ascontiguousarray(transformed, dtype=np.float32), + self.config.input_order, + self.config.representation == "rotation_vector", + self.left, + self.limit, + self.config.primitive_selection == "smooth", + min(len(transformed), os.cpu_count() or 1), + ) + return np.asarray(cells, dtype=np.int64) + converted = self._from_representation(transformed) + batch, steps = transformed.shape[:2] + true = np.broadcast_to(_IDENTITY, (batch, 4)).copy() + reconstruction = true.copy() + direction = np.zeros((batch, 3), dtype=np.float64) + cells = np.empty((batch, steps, 3), dtype=np.int64) + for step in range(steps): + value = converted[:, step] + true = value if self.config.input_order == 0 else self._apply(true, value) + residual = _quaternion_to_rotvec(self._residual(true, reconstruction)) + choice = self._nearest_batch(residual) + if self.config.primitive_selection == "smooth": + choice = np.stack( + [ + self._smooth( + residual[index], + reconstruction[index], + true[index], + direction[index], + choice[index], + ) + for index in range(batch) + ] + ) + primitive = np.stack( + [self.primitives[index][choice[:, index]] for index in range(3)], axis=1 + ) + reconstruction = self._apply( + reconstruction, _rotvec_to_quaternion(primitive) + ) + moved = np.linalg.norm(primitive, axis=1) > 0.0 + direction[moved] = primitive[moved] + cells[:, step] = choice + return cells + + def decode_batch( + self, cells: np.ndarray, previous_first_order: np.ndarray | None = None + ) -> np.ndarray: + """Decode ``[B,T,3]`` cells in the configured boundary representation.""" + if self.config.primitive_order == 2: + return self._decode_second_order(cells, previous_first_order) + indices = np.asarray(cells, dtype=np.int64) + rotvec = np.stack( + [self.primitives[index][indices[:, :, index]] for index in range(3)], axis=2 + ) + if self.config.input_order == 1: + represented = ( + rotvec + if self.config.representation == "rotation_vector" + else self._to_representation( + _rotvec_to_quaternion(rotvec.astype(np.float32)) + ) + ) + actions = np.asarray(represented / _scale(self.config), dtype=np.float32) + # Existing artifacts expose first-order actions through a zero- + # anchored position adapter. The float32 round trip is observable + # and therefore intentionally preserved for exact compatibility. + return _difference_additive(np.cumsum(actions, axis=1, dtype=np.float32), 1) + + batch, steps = indices.shape[:2] + reconstruction = np.broadcast_to(_IDENTITY, (batch, 4)).copy() + quaternion = np.empty((batch, steps, 4), dtype=np.float64) + primitive_quaternion = _rotvec_to_quaternion(rotvec) + for step in range(steps): + reconstruction = self._apply(reconstruction, primitive_quaternion[:, step]) + quaternion[:, step] = reconstruction + represented = self._to_representation(quaternion) + return np.asarray(represented / _scale(self.config), dtype=np.float32) + + def _encode_second_order( + self, actions: np.ndarray, previous_first_order: np.ndarray + ) -> np.ndarray: + """Quantize increment changes on SO(3), respecting multiplication order.""" + converted = self._from_representation( + _transform(self.config, actions).astype(np.float64) + ) + increment = self._from_representation( + previous_first_order.astype(np.float64) * _scale(self.config) + ) + if resolve_backend("auto") == "native" and hasattr( + require_native(), "so3_second_order_encode_batch" + ): + return np.asarray( + require_native().so3_second_order_encode_batch( + self._native_table, + self._native_bins, + np.ascontiguousarray(converted), + np.ascontiguousarray(increment), + self.config.input_order, + self.left, + self.limit, + self.second_order_radius, + self.config.primitive_selection == "smooth", + min( + len(actions), + max( + 1, + int(os.environ.get("ACTIONCODEC2_PHYSICAL_THREADS", "16")), + ), + ), + ), + dtype=np.int64, + ) + batch, steps = actions.shape[:2] + true = np.broadcast_to(_IDENTITY, (batch, 4)).copy() + reconstruction = true.copy() + direction = np.zeros((batch, 3), dtype=np.float64) + cells = np.empty((batch, steps, 3), dtype=np.int64) + for step in range(steps): + true = ( + converted[:, step] + if self.config.input_order == 0 + else self._apply(true, converted[:, step]) + ) + desired = self._residual(true, reconstruction) + correction = self._residual(desired, increment) + request = _quaternion_to_rotvec(correction) + choice = self._nearest_batch(request) + if self.config.primitive_selection == "smooth": + choice = self._smooth_second_order_batch( + request, correction, direction, choice + ) + primitive = np.stack( + [axis[choice[:, j]] for j, axis in enumerate(self.primitives)], axis=1 + ) + primitive = self._second_order_primitive(primitive) + increment = self._apply(increment, _rotvec_to_quaternion(primitive)) + reconstruction = self._apply(reconstruction, increment) + moved = np.linalg.norm(primitive, axis=1) > 0 + direction[moved] = primitive[moved] + cells[:, step] = choice + return cells + + def _second_order_primitive(self, rotvec: np.ndarray) -> np.ndarray: + """Project the lattice to the declared increment-change ball. + + Euclidean projection cannot increase distance to any log in the ball; + Exp is nonexpansive. Thus the Cartesian epsilon cover remains valid, + while every executable primitive respects c_R + 3*epsilon_R. + """ + norm = np.linalg.norm(rotvec, axis=-1, keepdims=True) + return rotvec * np.minimum( + 1.0, self.second_order_radius / np.maximum(norm, np.finfo(float).tiny) + ) + + def _smooth_second_order_batch( + self, + request: np.ndarray, + correction: np.ndarray, + direction: np.ndarray, + nearest: np.ndarray, + ) -> np.ndarray: + """Evaluate local Cartesian candidates in batches with intrinsic errors. + + Search windows use the same geometric pruning as ``_smooth``. Blocks + cap scratch memory even for unusually dense user-supplied grids. + """ + if len(request) > 64: + return np.concatenate( + [ + self._smooth_second_order_batch( + request[i : i + 64], + correction[i : i + 64], + direction[i : i + 64], + nearest[i : i + 64], + ) + for i in range(0, len(request), 64) + ] + ) + axes, counts = [], [] + for j, axis in enumerate(self.primitives): + mask = np.abs(axis[None] - request[:, j, None]) <= 2 * self.limit + 1e-12 + mask[np.arange(len(request)), nearest[:, j]] = True + count = mask.sum(axis=1) + rows, columns = np.nonzero(mask) + slots = np.arange(len(rows)) - np.repeat(np.cumsum(count) - count, count) + candidates = np.zeros((len(request), int(count.max())), dtype=np.int64) + candidates[rows, slots] = columns + axes.append(candidates) + counts.append(count) + counts = np.stack(counts, axis=1) + widths = counts.max(axis=0) + total = int(np.prod(widths)) + best = nearest.copy() + best_error = np.full(len(request), np.inf) + best_preferred = np.zeros(len(request), dtype=bool) + for offset in range(0, total, 4096): + lattice = np.stack( + np.unravel_index( + np.arange(offset, min(total, offset + 4096)), tuple(widths) + ), + axis=1, + ) + valid = np.all(lattice[None] < counts[:, None], axis=2) + choices = np.stack( + [ + axis[ + np.arange(len(request))[:, None], + np.minimum(lattice[None, :, j], counts[:, j, None] - 1), + ] + for j, axis in enumerate(axes) + ], + axis=2, + ) + rotvec = np.stack( + [axis[choices[:, :, j]] for j, axis in enumerate(self.primitives)], + axis=2, + ) + rotvec = self._second_order_primitive(rotvec) + primitive = _rotvec_to_quaternion(rotvec) + residual = _multiply(_conjugate(primitive), correction[:, None]) + error = np.linalg.norm(_quaternion_to_rotvec(residual), axis=2) + valid &= error <= self.limit + 1e-12 + preferred = valid & (np.sum(rotvec * direction[:, None], axis=2) >= -1e-15) + has_preferred = preferred.any(axis=1) + eligible = np.where(has_preferred[:, None], preferred, valid) + scores = np.where(eligible, error, np.inf) + index = np.argmin(scores, axis=1) + score = scores[np.arange(len(request)), index] + improve = np.isfinite(score) & ( + (has_preferred & ~best_preferred) + | ((has_preferred == best_preferred) & (score < best_error)) + ) + best[improve] = choices[np.arange(len(request)), index][improve] + best_error[improve] = score[improve] + best_preferred[improve] = has_preferred[improve] + return best + + def _decode_second_order( + self, cells: np.ndarray, previous_first_order: np.ndarray + ) -> np.ndarray: + """Integrate rotation changes from a shared ``[B,3]`` first-order state.""" + rotvec = np.stack( + [axis[cells[:, :, j]] for j, axis in enumerate(self.primitives)], axis=2 + ) + rotvec = self._second_order_primitive(rotvec) + changes = _rotvec_to_quaternion(rotvec) + increment = self._from_representation( + previous_first_order.astype(np.float64) * _scale(self.config) + ) + reconstruction = np.broadcast_to(_IDENTITY, increment.shape).copy() + output = np.empty((*cells.shape[:2], 4), dtype=np.float64) + for step in range(cells.shape[1]): + increment = self._apply(increment, changes[:, step]) + reconstruction = self._apply(reconstruction, increment) + output[:, step] = ( + increment if self.config.input_order == 1 else reconstruction + ) + return np.asarray( + self._to_representation(output) / _scale(self.config), dtype=np.float32 + ) + + +class PhysicalCodec: + """Compose independent physical components in global action-dimension order. + + The public contract is ``[B,T,D] -> [B,T,D]``. Components may be declared + in any YAML order and may use arbitrary non-overlapping dimension lists; + scattering always restores global dimension ``j`` to slot ``j``. + """ + + def __init__(self, config: ActionCodec2Config, primitives: PrimitiveStore) -> None: + if primitives.dimension != config.physical.input_dim: + raise ConfigError( + f"primitive dimension {primitives.dimension} does not match input_dim " + f"{config.physical.input_dim}" + ) + self.config = config + self.primitives = primitives + runtimes: list[tuple[ComponentConfig, _ComponentRuntime]] = [] + for component in config.physical.components.values(): + component_primitives = tuple( + primitives.values[index] for index in component.dims + ) + runtime: _ComponentRuntime = ( + _AdditiveRuntime(component, component_primitives) + if component.space == "additive" + else _SO3Runtime(component, component_primitives) + ) + runtimes.append((component, runtime)) + self._runtimes = tuple(runtimes) + self.bins = primitives.bins + self.requires_previous_first_order = any( + c.primitive_order == 2 for c, _ in runtimes + ) + + def _check_previous_first_order( + self, value: np.ndarray | None, batch: int + ) -> np.ndarray | None: + if value is None: + if self.requires_previous_first_order: + raise ValueError( + "second-order tokenization requires previous_first_order; no zero initialization is inferred" + ) + return None + array = np.asarray(value, dtype=np.float32) + shape = (batch, self.config.physical.input_dim) + if array.shape != shape or not np.all(np.isfinite(array)): + raise ValueError(f"previous_first_order must be finite with shape {shape}") + return array + + def _check_actions(self, actions: np.ndarray) -> np.ndarray: + values = np.asarray(actions, dtype=np.float32) + expected = self.config.physical.input_dim + if values.ndim != 3 or values.shape[2] != expected or values.shape[1] < 1: + raise ValueError(f"actions must have non-empty shape [B, T, {expected}]") + if not np.all(np.isfinite(values)): + raise ValueError("actions must contain only finite values") + return values + + def final_first_order( + self, cells: np.ndarray, *, previous_first_order: np.ndarray + ) -> np.ndarray: + """Return ``[B,D]`` state after these cells, without position differencing. + + Only second-order component slots carry state; other slots are zero. + Pass a cell prefix when only part of the generated chunk was executed. + """ + values = self._check_cells(cells) + previous = self._check_previous_first_order(previous_first_order, len(values)) + result = np.zeros( + (len(values), self.config.physical.input_dim), dtype=np.float32 + ) + for component, runtime in self._runtimes: + if component.primitive_order != 2: + continue + local = np.take(values, component.dims, axis=2) + boundary = np.take(previous, component.dims, axis=1) + picked = np.stack( + [axis[local[:, :, j]] for j, axis in enumerate(runtime.primitives)], + axis=2, + ) + if component.space == "additive": + final = boundary.astype(np.float64) + np.sum( + picked.astype(np.float64), axis=1 + ) / _scale(component) + else: + increment = runtime._from_representation( + boundary.astype(np.float64) * _scale(component) + ) + changes = _rotvec_to_quaternion(runtime._second_order_primitive(picked)) + for step in range(local.shape[1]): + increment = runtime._apply(increment, changes[:, step]) + final = runtime._to_representation(increment) / _scale(component) + result[:, component.dims] = final + return result + + def _check_cells(self, cells: np.ndarray) -> np.ndarray: + values = np.asarray(cells, dtype=np.int64) + expected = self.config.physical.input_dim + if values.ndim != 3 or values.shape[2] != expected or values.shape[1] < 1: + raise ValueError(f"cells must have non-empty shape [B, T, {expected}]") + counts = np.asarray(self.bins, dtype=np.int64) + if values.size and ( + np.any(values < 0) or np.any(values >= counts[None, None, :]) + ): + raise ValueError("a cell index is outside its dimension's primitive set") + return values + + def encode_batch( + self, actions: np.ndarray, *, previous_first_order: np.ndarray | None = None + ) -> np.ndarray: + """Quantize ``[B,T,D]`` with a required ``[B,D]`` boundary for order 2. + + The boundary is the preceding reconstructed first-order quantity, in + input units and representation, before scaling (never source-clipped). + Position/orientation coordinates are local to zero/identity per call. + """ + values = self._check_actions(actions) + previous = self._check_previous_first_order(previous_first_order, len(values)) + cells = np.empty(values.shape, dtype=np.int64) + for component, runtime in self._runtimes: + component_cells = runtime.encode_batch( + np.take(values, component.dims, axis=2), + None if previous is None else np.take(previous, component.dims, axis=1), + ) + for local, global_dimension in enumerate(component.dims): + cells[:, :, global_dimension] = component_cells[:, :, local] + return cells + + def decode_batch( + self, cells: np.ndarray, *, previous_first_order: np.ndarray | None = None + ) -> np.ndarray: + """Decode ``[B,T,D]`` using the same ``[B,D]`` boundary as encoding.""" + values = self._check_cells(cells) + previous = self._check_previous_first_order(previous_first_order, len(values)) + actions = np.empty(values.shape, dtype=np.float32) + for component, runtime in self._runtimes: + component_actions = runtime.decode_batch( + np.take(values, component.dims, axis=2), + None if previous is None else np.take(previous, component.dims, axis=1), + ) + for local, global_dimension in enumerate(component.dims): + actions[:, :, global_dimension] = component_actions[:, :, local] + return actions diff --git a/runtime/physical/physical_processor.py b/runtime/physical/physical_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..0fcec6b1bf79ddabd90a07462c77adf7a0d1e96c --- /dev/null +++ b/runtime/physical/physical_processor.py @@ -0,0 +1,384 @@ +"""Public ActionCodec2 processor and constrained-generation grammars. + +The processor is scheme-agnostic: ``config.set_bpe.scheme`` picks the +serializer and the matching grammar, and everything else -- physical +quantization, batching, second-order boundaries, artifact layout -- is shared. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Sequence +from numbers import Integral +from pathlib import Path +from typing import Any + +import numpy as np +from transformers.processing_utils import ProcessorMixin + +from ..common._validation import active_dimensions +from ..tokenization.bpe import ( + BPE_NAME, + BPESerializer, + SetBPESerializer, + serializer_class, +) +from ..routing.cache import BoundedCache +from .configuration import CONFIG_NAME, MODEL_TYPE, ActionCodec2Config +from ..tokenization.grammar import ActionCodec2BPEGrammar as ActionCodec2BPEGrammar +from ..tokenization.grammar import ActionCodec2Grammar as ActionCodec2Grammar +from ..integration.hf import ( + resolve_pretrained_directory, + save_huggingface_metadata, + upload_pretrained_directory, +) +from .physical_codec import PhysicalCodec +from .primitives import PRIMITIVES_NAME, PrimitiveStore + + +class ActionCodec2Processor(ProcessorMixin): + """Hugging Face-style processor for continuous robot action trajectories. + + ActionCodec2 is an ``AutoProcessor``-class object, not an ``AutoModel``: it + has no neural ``forward`` method or trainable weights. Its fitted state is a + physical primitive vocabulary plus lossless Set-BPE rules. Encoding maps + ``[B,T,D]`` continuous actions to ragged token-ID streams; decoding performs + the inverse token/cell path and the configured lossy physical reconstruction. + + Args: + config: Physical quantization and tokenization scheme configuration. + primitives: Fitted primitive values for every physical coordinate. + serializer: Fitted BPE rules matching the configured scheme, primitive + bin counts, and lookahead. + """ + + model_type = MODEL_TYPE + name = "actioncodec2" + is_exact = False + set_bpe_is_lossless = True + model_input_names = ["input_ids"] + attributes = [] + + def __init__( + self, + config: ActionCodec2Config, + primitives: PrimitiveStore, + serializer: SetBPESerializer | BPESerializer, + ) -> None: + if serializer.scheme != config.set_bpe.scheme: + raise ValueError("serializer scheme does not match config.yaml") + if tuple(serializer.bins) != tuple(primitives.bins): + raise ValueError("BPE bins do not match primitives.npz") + if serializer.lookahead != config.set_bpe.lookahead: + raise ValueError("BPE lookahead does not match config.yaml") + # ProcessorMixin is designed around modality sub-processors. This + # processor is atomic, so it deliberately owns and serializes its three + # artifact objects directly instead of invoking ProcessorMixin.__init__. + self.config = config + self.primitives = primitives + self.serializer = serializer + self._grammar_cache = BoundedCache(8) + self.physical = PhysicalCodec(config, primitives) + self.action_dim = config.physical.input_dim + self.bins = primitives.bins + self.vocab_size = serializer.vocab_size + self.lookahead = serializer.lookahead + self.requires_previous_first_order = self.physical.requires_previous_first_order + zero_cells = [] + for dimension, axis in enumerate(primitives.values): + matches = np.flatnonzero(np.asarray(axis) == 0) + if not len(matches): + raise ValueError( + "ActionCodec2 artifacts require an exact zero primitive in every axis; " + f"dimension {dimension} has none" + ) + zero_cells.append(int(matches[0])) + self._zero_cells = np.asarray(zero_cells, dtype=np.int64) + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str | Path, + **kwargs: Any, + ) -> ActionCodec2Processor: + """Load a local directory or Hugging Face Hub snapshot. + + Standard Hub keywords such as ``revision``, ``cache_dir``, ``token``, + and ``local_files_only`` are supported. Numerical files remain the + strict physical artifact rather than being embedded in JSON metadata. + """ + directory = resolve_pretrained_directory( + pretrained_model_name_or_path, **kwargs + ) + config = ActionCodec2Config.from_pretrained(directory) + primitives = PrimitiveStore.from_pretrained( + directory, input_dim=config.physical.input_dim + ) + serializer = serializer_class(config).from_pretrained( + directory, + bins=primitives.bins, + lookahead=config.set_bpe.lookahead, + ) + return cls(config, primitives, serializer) + + def save_pretrained( + self, + save_directory: str | Path, + push_to_hub: bool = False, + **kwargs: Any, + ) -> list[str]: + """Save numerical artifacts plus standard AutoProcessor metadata.""" + if not push_to_hub and kwargs: + raise TypeError( + f"unused save_pretrained keyword(s): {', '.join(sorted(kwargs))}" + ) + directory = Path(save_directory) + self.config.save_pretrained(directory) + self.primitives.save_pretrained(directory) + self.serializer.save_pretrained(directory) + metadata = save_huggingface_metadata(directory) + files = [ + directory / CONFIG_NAME, + directory / PRIMITIVES_NAME, + directory / BPE_NAME, + *metadata, + ] + if push_to_hub: + upload_pretrained_directory(directory, **kwargs) + return [str(path) for path in files] + + def to_dict(self, legacy_serialization: bool = True) -> dict[str, object]: + """Return metadata without embedding numerical artifact arrays. + + ``legacy_serialization`` is retained because Transformers passes it to + processor serializers in older releases; ActionCodec2 has one format. + """ + del legacy_serialization + return { + "processor_class": self.__class__.__name__, + "model_type": self.model_type, + "action_dim": self.action_dim, + "vocab_size": self.vocab_size, + } + + @property + def artifact_fingerprint(self) -> str: + """Content hash over config, primitive axes, and ordered BPE rules.""" + digest = hashlib.sha256() + digest.update( + json.dumps( + self.config.to_dict(), sort_keys=True, separators=(",", ":") + ).encode() + ) + for axis in self.primitives.values: + values = np.ascontiguousarray(axis) + digest.update(str(values.dtype).encode()) + digest.update(np.asarray(values.shape, dtype=np.int64).tobytes()) + digest.update(values.tobytes()) + digest.update( + json.dumps( + self.serializer.to_dict(), sort_keys=True, separators=(",", ":") + ).encode() + ) + return digest.hexdigest() + + def quantize_batch( + self, actions: np.ndarray, *, previous_first_order: np.ndarray | None = None + ) -> np.ndarray: + """Quantize ``[B,T,D]`` using a ``[B,D]`` first-order boundary if required.""" + return self.physical.encode_batch( + actions, previous_first_order=previous_first_order + ) + + def _dense_previous_first_order( + self, previous_first_order: np.ndarray | None, dims: tuple[int, ...], batch: int + ) -> np.ndarray | None: + """Expand a compact ``[B,D_active]`` boundary, validating active slots.""" + active = set(dims) + for component in self.config.physical.components.values(): + present = active.intersection(component.dims) + if ( + component.primitive_order == 2 + and component.space == "so3" + and present + and present != set(component.dims) + ): + raise ValueError( + "second-order SO(3) requires all three active component dimensions" + ) + if previous_first_order is None: + if any( + c.primitive_order == 2 and active.intersection(c.dims) + for c in self.config.physical.components.values() + ): + raise ValueError( + f"second-order tokenization requires previous_first_order with shape {(batch, len(dims))} in canonical input units" + ) + return ( + np.zeros((batch, self.action_dim), dtype=np.float32) + if self.requires_previous_first_order + else None + ) + values = np.asarray(previous_first_order, dtype=np.float32) + if values.shape != (batch, len(dims)) or not np.all(np.isfinite(values)): + raise ValueError( + f"previous_first_order must be finite with shape {(batch, len(dims))}" + ) + dense = np.zeros((batch, self.action_dim), dtype=np.float32) + dense[:, dims] = values + return dense + + def tokens_from_cells( + self, + cells: np.ndarray, + *, + active_dims: Sequence[int] | None = None, + ) -> list[int]: + """Apply the lossless tokenizer to one compact or dense cell trajectory.""" + return self.serializer.encode(cells, active_dims=active_dims) + + def cells_from_tokens( + self, + tokens: Sequence[int], + *, + active_dims: Sequence[int] | None = None, + ) -> np.ndarray: + """Losslessly recover dense cells with ``-1`` in absent dimensions.""" + return self.serializer.decode(tokens, active_dims=active_dims) + + def _fill_inactive_cells(self, cells: np.ndarray) -> np.ndarray: + """Replace semantic absence with each physical axis's exact zero cell.""" + values = np.asarray(cells, dtype=np.int64) + return np.where(values >= 0, values, self._zero_cells) + + def encode_batch( + self, + actions: np.ndarray, + *, + active_dims: Sequence[int] | None = None, + previous_first_order: np.ndarray | None = None, + ) -> list[list[int]]: + """Encode a dense or compact ``[B,T,D]`` batch into ragged streams. + + ``active_dims`` maps compact coordinates into this processor's physical + profile. Supplying the whole batch at once preserves native physical + quantization across the batch dimension. + ``previous_first_order[B,D_active]`` is required for active second-order + components and follows the same compact mapping, in input units. + """ + values = np.asarray(actions, dtype=np.float32) + if ( + values.ndim != 3 + or len(values) < 1 + or values.shape[1] < 1 + or values.shape[2] < 1 + ): + raise ValueError("actions must have non-empty shape [B, T, D_active]") + dims = active_dimensions(active_dims, self.action_dim) + if len(dims) != values.shape[2]: + raise ValueError("active_dims must have one entry per action coordinate") + full_width = dims == tuple(range(self.action_dim)) + if full_width: + dense = values + else: + dense = np.zeros( + (len(values), values.shape[1], self.action_dim), dtype=np.float32 + ) + dense[:, :, dims] = values + previous = self._dense_previous_first_order( + previous_first_order, dims, len(values) + ) + cells = self.quantize_batch(dense, previous_first_order=previous) + if not full_width: + cells = np.take(cells, dims, axis=2) + return [self.tokens_from_cells(row, active_dims=dims) for row in cells] + + def decode_batch( + self, + streams: Sequence[Sequence[int]], + *, + active_dims: Sequence[int] | None = None, + previous_first_order: np.ndarray | None = None, + ) -> np.ndarray: + """Decode equal-length streams into one batched physical array. + + With ``active_dims``, the last output dimension is compact and follows + that mapping. Inactive physical-profile slots never appear at the + semantic boundary. + ``previous_first_order[B,D_active]`` must match the encoding boundary + for active second-order components. + """ + dims = active_dimensions(active_dims, self.action_dim) + cells = [self.cells_from_tokens(stream, active_dims=dims) for stream in streams] + if not cells: + raise ValueError("streams must not be empty") + lengths = {len(value) for value in cells} + if len(lengths) != 1: + raise ValueError( + "all token streams in a batch must decode to the same length" + ) + stacked_cells = np.stack(cells) + dense_cells = self._fill_inactive_cells(stacked_cells) + previous = self._dense_previous_first_order( + previous_first_order, dims, len(cells) + ) + decoded = self.physical.decode_batch(dense_cells, previous_first_order=previous) + return np.take(decoded, dims, axis=2) + + def final_first_order( + self, + streams: Sequence[Sequence[int]], + *, + previous_first_order: np.ndarray, + active_dims: Sequence[int] | None = None, + executed_steps: int | None = None, + ) -> np.ndarray: + """Return the next compact ``[B,D_active]`` reconstructed boundary. + + ``executed_steps`` selects the executed prefix in codec timesteps; + omitted means the full chunk. Tokens may span the prefix boundary. + """ + dims = active_dimensions(active_dims, self.action_dim) + cells = [self.cells_from_tokens(stream, active_dims=dims) for stream in streams] + if not cells or len({len(row) for row in cells}) != 1: + raise ValueError( + "streams must be non-empty and have equal decoded horizons" + ) + if executed_steps is not None: + if ( + isinstance(executed_steps, bool) + or not isinstance(executed_steps, Integral) + or not 1 <= executed_steps <= len(cells[0]) + ): + raise ValueError("executed_steps must be within the decoded horizon") + cells = [row[:executed_steps] for row in cells] + previous = self._dense_previous_first_order( + previous_first_order, dims, len(cells) + ) + result = self.physical.final_first_order( + self._fill_inactive_cells(np.stack(cells)), previous_first_order=previous + ) + return np.take(result, dims, axis=1) + + def grammar( + self, + horizon: int, + *, + active_dims: Sequence[int] | None = None, + ) -> ActionCodec2Grammar | ActionCodec2BPEGrammar: + """Build a constrained grammar for one horizon and route support.""" + if isinstance(horizon, bool) or not isinstance(horizon, int) or horizon < 1: + raise ValueError("horizon must be a positive integer") + dims = active_dimensions(active_dims, self.action_dim) + key = (horizon, dims) + grammar = self._grammar_cache.get(key) + if grammar is None: + grammar_class = ( + ActionCodec2BPEGrammar + if isinstance(self.serializer, BPESerializer) + else ActionCodec2Grammar + ) + grammar = grammar_class(self.serializer, horizon=horizon, active_dims=dims) + self._grammar_cache[key] = grammar + return grammar diff --git a/runtime/physical/primitives.py b/runtime/physical/primitives.py new file mode 100644 index 0000000000000000000000000000000000000000..560cd330a83632b2e76c1fe179adcbbc0e8dbd0f --- /dev/null +++ b/runtime/physical/primitives.py @@ -0,0 +1,83 @@ +"""Numerical primitive parameters stored in ``primitives.npz``.""" + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path + +import numpy as np + +from .configuration import ConfigError + +PRIMITIVES_NAME = "primitives.npz" + + +def _key(dimension: int) -> str: + return f"dimension_{dimension:04d}" + + +class PrimitiveStore: + """Immutable numerical primitive axes in global action-dimension order. + + Array ``values[j]`` is the decoder vocabulary selected by cell slot ``j``. + NPZ key names are internal and deliberately absent from ``config.yaml``. + """ + + def __init__(self, values: Sequence[np.ndarray | Sequence[float]]) -> None: + if not values: + raise ConfigError("primitives must contain at least one dimension") + arrays = [] + for dimension, raw in enumerate(values): + array = np.asarray(raw) + if array.dtype.hasobject: + raise ConfigError( + f"primitive dimension {dimension} must not use object dtype" + ) + if array.ndim != 1 or not len(array): + raise ConfigError( + f"primitive dimension {dimension} must be a non-empty vector" + ) + real_numeric = np.issubdtype(array.dtype, np.integer) or np.issubdtype( + array.dtype, np.floating + ) + if not real_numeric or not np.all(np.isfinite(array)): + raise ConfigError( + f"primitive dimension {dimension} must contain finite numbers" + ) + arrays.append(np.array(array, copy=True)) + self.values = tuple(arrays) + self.bins = tuple(len(array) for array in arrays) + + @property + def dimension(self) -> int: + return len(self.values) + + @classmethod + def from_pretrained( + cls, directory: str | Path, *, input_dim: int + ) -> PrimitiveStore: + """Load exactly one safe numeric vector per expected global dimension.""" + path = Path(directory) / PRIMITIVES_NAME + with np.load(path, allow_pickle=False) as payload: + expected = {_key(index) for index in range(input_dim)} + actual = set(payload.files) + missing = sorted(expected - actual) + unknown = sorted(actual - expected) + if missing: + raise ConfigError( + f"{PRIMITIVES_NAME} is missing array(s): {', '.join(missing)}" + ) + if unknown: + raise ConfigError( + f"{PRIMITIVES_NAME} contains unknown array(s): {', '.join(unknown)}" + ) + values = [payload[_key(index)] for index in range(input_dim)] + return cls(values) + + def save_pretrained(self, directory: str | Path) -> None: + """Write ``primitives.npz`` without pickled objects or path metadata.""" + path = Path(directory) / PRIMITIVES_NAME + path.parent.mkdir(parents=True, exist_ok=True) + np.savez( + path, **{_key(index): value for index, value in enumerate(self.values)} + ) diff --git a/runtime/physical/semantic.py b/runtime/physical/semantic.py new file mode 100644 index 0000000000000000000000000000000000000000..b6878e05e77c41d15e54756dfac23eaaf7634f17 --- /dev/null +++ b/runtime/physical/semantic.py @@ -0,0 +1,452 @@ +"""Canonical physical preprocessing for routed ActionCodec2 profiles. + +Dataset presets may select columns, convert units, and resample observations, +but they may not define primitives or BPE rules. This boundary keeps robot +storage conventions outside the shared physical vocabularies. +""" + +from __future__ import annotations + +import math + +import numpy as np +from scipy.interpolate import CubicSpline, PchipInterpolator +from scipy.spatial.transform import Rotation, RotationSpline + +from ._so3 import euler_xyz_to_quaternion as _euler_xyz_to_quaternion +from ._so3 import multiply_quaternions as _quaternion_multiply +from ._so3 import normalize_quaternion as _normalize_quaternion +from ._so3 import quaternion_to_rotvec as _quaternion_to_rotvec +from ._so3 import rotvec_to_quaternion as _rotvec_to_quaternion + +_INTERPOLATION_METHODS = {"linear", "pchip", "cubic"} + + +def _validate_rate(value: float, name: str) -> float: + rate = float(value) + if not math.isfinite(rate) or rate <= 0.0: + raise ValueError(f"{name} must be positive and finite") + return rate + + +def _quaternion_slerp( + left: np.ndarray, right: np.ndarray, weight: np.ndarray +) -> np.ndarray: + """Shortest-path SLERP over matching ``(N, 4)`` xyzw arrays.""" + first = _normalize_quaternion(np.asarray(left, dtype=np.float64)) + second = _normalize_quaternion(np.asarray(right, dtype=np.float64)) + dot = np.sum(first * second, axis=1) + flip = dot < 0.0 + second[flip] *= -1.0 + dot = np.clip(np.abs(dot), 0.0, 1.0) + output = np.empty_like(first) + near = dot > 0.9995 + output[near] = first[near] + weight[near, None] * (second[near] - first[near]) + regular = ~near + if np.any(regular): + angle = np.arccos(dot[regular]) + denominator = np.sin(angle) + w = weight[regular] + output[regular] = ( + np.sin((1.0 - w) * angle)[:, None] / denominator[:, None] * first[regular] + + np.sin(w * angle)[:, None] / denominator[:, None] * second[regular] + ) + return _normalize_quaternion(output) + + +def resampled_horizon(length: int, *, source_hz: float, target_hz: float) -> int: + """Return the target control-step count for a regular source horizon. + + A ``[T,D]`` action/state target represents ``T`` controller intervals, so + its duration is ``T / source_hz``. This differs intentionally from a + point-sampled signal whose first and last rows span only ``T-1`` periods. + The interval convention maps duration to a rounded number of target-rate + controller intervals, which is the runtime action-chunk contract. + """ + if ( + isinstance(length, bool) + or not isinstance(length, (int, np.integer)) + or length < 1 + ): + raise ValueError("a trajectory must contain at least one sample") + source_rate = _validate_rate(source_hz, "source_hz") + target_rate = _validate_rate(target_hz, "target_hz") + scaled = int(math.floor(int(length) * target_rate / source_rate + 0.5)) + return max(1, scaled) + + +def _target_times( + length: int, + source_hz: float, + target_hz: float, + *, + target_length: int | None, +) -> np.ndarray: + count = ( + resampled_horizon(length, source_hz=source_hz, target_hz=target_hz) + if target_length is None + else target_length + ) + if isinstance(count, bool) or not isinstance(count, (int, np.integer)) or count < 1: + raise ValueError("target_length must be a positive integer") + return np.arange(int(count), dtype=np.float64) / target_hz + + +def _guarded_natural_cubic( + source_time: np.ndarray, + values: np.ndarray, + target_time: np.ndarray, +) -> np.ndarray: + """Use natural cubic unless an axis materially leaves its observed envelope. + + Natural cubic is C2-continuous but can overshoot sparse control knots. + Falling back only the offending additive axes to PCHIP preserves the + smoother curve wherever it is supported by observations while preventing + interpolation from inventing a large new coordinate-wise extremum. + """ + sampled = np.asarray( + CubicSpline(source_time, values, axis=0, bc_type="natural")(target_time), + dtype=np.float64, + ) + lower = np.min(values, axis=0) + upper = np.max(values, axis=0) + # Small cubic ringing can be the mechanism that removes velocity corners; + # treat only a material (>5% of the observed axis span) new extremum as an + # unsafe invention. Profile order-0/order-1 checks remain the hard gate. + tolerance = 1e-9 + 0.05 * (upper - lower) + fallback_axes = np.any( + (sampled < lower[None] - tolerance[None]) + | (sampled > upper[None] + tolerance[None]), + axis=0, + ) + if np.any(fallback_axes): + sampled[:, fallback_axes] = PchipInterpolator( + source_time, values[:, fallback_axes], axis=0 + )(target_time) + return sampled + + +def resample_uniform( + values: np.ndarray, + *, + source_hz: float, + target_hz: float, + kind: str = "linear", + target_length: int | None = None, + method: str = "linear", +) -> np.ndarray: + """Resample one regular control-horizon trajectory onto a target clock. + + Args: + values (np.ndarray): ``(T, D)`` continuous/discrete state samples. + source_hz: Source frequency in Hz. + target_hz: Target frequency in Hz. + kind: ``linear``, ``hold``, ``pose_xyzw``, or ``pose_euler_xyz``. + Pose output is always ``(x, y, z, qx, qy, qz, qw)`` and uses + SLERP for orientation. + target_length: Explicit target horizon. By default the interval + duration ``len(values) / source_hz`` is preserved. + method: ``linear`` uses linear translation plus SLERP. ``pchip`` uses + shape-preserving cubic translation. ``cubic`` uses a natural cubic + translation spline, with per-axis PCHIP fallback if the spline + materially leaves the source trajectory's observed value envelope. + Both cubic methods use an intrinsic SO(3) ``RotationSpline`` for + orientation. + + Returns: + np.ndarray: ``(T_target, D)`` regularly sampled targets. When + upsampling, times after the final source knot hold its value until + the final source control interval ends. + """ + source_rate = _validate_rate(source_hz, "source_hz") + target_rate = _validate_rate(target_hz, "target_hz") + array = np.asarray(values, dtype=np.float64) + if array.ndim != 2 or len(array) < 1 or not np.all(np.isfinite(array)): + raise ValueError("values must be a non-empty finite [T, D] array") + if kind not in {"linear", "hold", "pose_xyzw", "pose_euler_xyz"}: + raise ValueError("unsupported resampling kind") + if method not in _INTERPOLATION_METHODS: + raise ValueError(f"method must be one of {sorted(_INTERPOLATION_METHODS)}") + if kind == "pose_xyzw" and array.shape[1] != 7: + raise ValueError("pose_xyzw values must have shape [T, 7]") + if kind == "pose_euler_xyz" and array.shape[1] != 6: + raise ValueError("pose_euler_xyz values must have shape [T, 6]") + standard_length = resampled_horizon( + len(array), source_hz=source_rate, target_hz=target_rate + ) + if source_rate == target_rate and target_length in {None, standard_length}: + if kind == "pose_euler_xyz": + array = np.concatenate( + (array[:, :3], _euler_xyz_to_quaternion(array[:, 3:6])), axis=1 + ) + return np.asarray(array, dtype=np.float32) + + target_time = _target_times( + len(array), source_rate, target_rate, target_length=target_length + ) + source_time = np.arange(len(array), dtype=np.float64) / source_rate + interpolation_time = np.minimum(target_time, source_time[-1]) + source_position = target_time * source_rate + left = np.minimum(np.floor(source_position).astype(np.int64), len(array) - 1) + right = np.minimum(left + 1, len(array) - 1) + weight = source_position - left + if kind == "hold": + result = array[left] + else: + pose = kind in {"pose_xyzw", "pose_euler_xyz"} + continuous = array[:, :3] if pose else array + if len(array) == 1: + sampled = np.repeat(continuous, len(target_time), axis=0) + elif method == "pchip": + sampled = PchipInterpolator(source_time, continuous, axis=0)( + interpolation_time + ) + elif method == "cubic": + sampled = _guarded_natural_cubic( + source_time, continuous, interpolation_time + ) + else: + sampled = array[left, : continuous.shape[1]] + weight[:, None] * ( + array[right, : continuous.shape[1]] - array[left, : continuous.shape[1]] + ) + if not pose: + result = sampled + else: + quaternion = ( + array[:, 3:7] + if kind == "pose_xyzw" + else _euler_xyz_to_quaternion(array[:, 3:6]) + ) + if len(array) == 1: + rotation = np.repeat(quaternion, len(target_time), axis=0) + elif method == "linear": + rotation = _quaternion_slerp( + quaternion[left], quaternion[right], weight + ) + else: + rotation = RotationSpline(source_time, Rotation.from_quat(quaternion))( + interpolation_time + ).as_quat() + result = np.concatenate((sampled, rotation), axis=1) + return np.asarray(result, dtype=np.float32) + + +def eef_pose_absolute(poses: np.ndarray) -> np.ndarray: + """Convert canonical ``xyz + xyzw`` state to ``xyz + rotation_vector``.""" + values = np.asarray(poses, dtype=np.float64) + if values.ndim != 2 or values.shape[1] != 7 or not np.all(np.isfinite(values)): + raise ValueError("poses must be a finite [T, 7] array") + return np.asarray( + np.concatenate((values[:, :3], _quaternion_to_rotvec(values[:, 3:7])), axis=1), + dtype=np.float32, + ) + + +def eef_pose_components(poses: np.ndarray, representation: str) -> np.ndarray: + """Convert canonical ``xyz + xyzw`` poses to one registered source layout.""" + values = np.asarray(poses, dtype=np.float64) + if values.ndim != 2 or values.shape[1] != 7 or not np.all(np.isfinite(values)): + raise ValueError("poses must be a finite [T, 7] array") + quaternion = _normalize_quaternion(values[:, 3:7]) + if representation == "rotation_vector": + rotation = _quaternion_to_rotvec(quaternion) + elif representation == "euler_xyz": + rotation = Rotation.from_quat(quaternion).as_euler("xyz") + elif representation == "quaternion_xyzw": + rotation = quaternion + else: + raise ValueError("unsupported EEF rotation representation") + return np.asarray( + np.concatenate((values[:, :3], rotation), axis=1), dtype=np.float32 + ) + + +def eef_pose_from_components( + position: np.ndarray, + rotation: np.ndarray, + representation: str, +) -> np.ndarray: + """Build canonical ``xyz + xyzw`` poses from an explicit source layout. + + Args: + position (np.ndarray): ``(T, 3)`` translation in metres. + rotation (np.ndarray): ``(T, 3)`` rotation vectors/XYZ Euler angles or + ``(T, 4)`` xyzw quaternions. + representation: ``rotation_vector``, ``euler_xyz``, or + ``quaternion_xyzw``. + + Returns: + np.ndarray: Finite normalized poses with shape ``(T, 7)``. + """ + translation = np.asarray(position, dtype=np.float64) + orientation = np.asarray(rotation, dtype=np.float64) + if translation.ndim != 2 or translation.shape[1] != 3 or len(translation) < 1: + raise ValueError("position must have non-empty shape [T, 3]") + if len(orientation) != len(translation) or not np.all(np.isfinite(translation)): + raise ValueError("position and rotation must have matching finite trajectories") + if representation == "rotation_vector": + if orientation.ndim != 2 or orientation.shape[1] != 3: + raise ValueError("rotation_vector must have shape [T, 3]") + quaternion = _rotvec_to_quaternion(orientation) + elif representation == "euler_xyz": + if orientation.ndim != 2 or orientation.shape[1] != 3: + raise ValueError("euler_xyz must have shape [T, 3]") + quaternion = _euler_xyz_to_quaternion(orientation) + elif representation == "quaternion_xyzw": + if orientation.ndim != 2 or orientation.shape[1] != 4: + raise ValueError("quaternion_xyzw must have shape [T, 4]") + quaternion = _normalize_quaternion(orientation) + else: + raise ValueError("unsupported EEF rotation representation") + if not np.all(np.isfinite(orientation)): + raise ValueError("position and rotation must be finite") + return np.asarray( + np.concatenate((translation, quaternion), axis=1), dtype=np.float32 + ) + + +def eef_pose_local_absolute( + poses: np.ndarray, *, initial_pose: np.ndarray | None = None +) -> np.ndarray: + """Express measured poses relative to an explicit or inferred initial pose. + + Translation is relative to the initial position. Orientation is the + right-multiplicative relative pose ``q_initial^-1 * q_t``, matching the + shared EEF profile. When ``initial_pose`` is omitted, the first observation + is the initial pose and row zero is exactly zero/identity. This is a local + position trajectory, not a per-step delta sequence. + """ + values = np.asarray(poses, dtype=np.float64) + if ( + values.ndim != 2 + or values.shape[1] != 7 + or len(values) < 1 + or not np.all(np.isfinite(values)) + ): + raise ValueError("poses must be a finite non-empty [T, 7] array") + if initial_pose is None: + initial = values[:1] + else: + initial = np.asarray(initial_pose, dtype=np.float64) + if initial.shape == (7,): + initial = initial[None] + if initial.shape != (1, 7) or not np.all(np.isfinite(initial)): + raise ValueError("initial_pose must be a finite [7] pose") + quaternion = _normalize_quaternion(values[:, 3:7]) + initial_quaternion = _normalize_quaternion(initial[:, 3:7]) + inverse_start = initial_quaternion.copy() + inverse_start[:, :3] *= -1.0 + relative = _quaternion_multiply( + np.broadcast_to(inverse_start, quaternion.shape), quaternion + ) + return np.asarray( + np.concatenate( + (values[:, :3] - initial[:, :3], _quaternion_to_rotvec(relative)), axis=1 + ), + dtype=np.float32, + ) + + +def eef_pose_deltas(poses: np.ndarray) -> np.ndarray: + """Differentiate a local ``xyz + xyzw`` path from zero/identity. + + Translation uses fixed-axis differences. Rotation uses right-multiplicative + increments ``q_previous^-1 * q_t``. The returned ``(T, 6)`` array is the + exact inverse convention of :func:`integrate_eef_deltas` before floating + point roundoff. + """ + values = np.asarray(poses, dtype=np.float64) + if ( + values.ndim != 2 + or values.shape[1] != 7 + or len(values) < 1 + or not np.all(np.isfinite(values)) + ): + raise ValueError("poses must be a finite non-empty [T, 7] array") + quaternion = _normalize_quaternion(values[:, 3:7]) + previous_position = np.concatenate((np.zeros((1, 3)), values[:-1, :3]), axis=0) + identity = np.asarray([[0.0, 0.0, 0.0, 1.0]], dtype=np.float64) + previous_quaternion = np.concatenate((identity, quaternion[:-1]), axis=0) + inverse_previous = previous_quaternion.copy() + inverse_previous[:, :3] *= -1.0 + increment = _quaternion_multiply(inverse_previous, quaternion) + return np.asarray( + np.concatenate( + (values[:, :3] - previous_position, _quaternion_to_rotvec(increment)), + axis=1, + ), + dtype=np.float32, + ) + + +def eef_pose_from_local_absolute( + poses: np.ndarray, *, initial_pose: np.ndarray +) -> np.ndarray: + """Attach an external initial pose to a local ``xyz + xyzw`` path.""" + values = np.asarray(poses, dtype=np.float64) + initial = np.asarray(initial_pose, dtype=np.float64) + if ( + values.ndim != 2 + or values.shape[1] != 7 + or len(values) < 1 + or not np.all(np.isfinite(values)) + ): + raise ValueError("poses must be a finite non-empty [T, 7] array") + if initial.shape == (7,): + initial = initial[None] + if initial.shape != (1, 7) or not np.all(np.isfinite(initial)): + raise ValueError("initial_pose must be a finite [7] pose") + local_quaternion = _normalize_quaternion(values[:, 3:7]) + initial_quaternion = _normalize_quaternion(initial[:, 3:7]) + quaternion = _quaternion_multiply( + np.broadcast_to(initial_quaternion, local_quaternion.shape), local_quaternion + ) + return np.asarray( + np.concatenate((values[:, :3] + initial[:, :3], quaternion), axis=1), + dtype=np.float32, + ) + + +def integrate_eef_deltas(deltas: np.ndarray) -> np.ndarray: + """Integrate ``[dx,dy,dz,body-rotvec]`` from a zero/identity pose. + + Args: + deltas (np.ndarray): ``(T, 6)`` translation and right-multiplicative + rotation increments in canonical units. + + Returns: + np.ndarray: ``(T, 7)`` absolute ``xyz + xyzw`` poses. Row ``t`` is the + pose after applying delta ``t``; the implicit boundary before row zero + is position zero and identity orientation. + """ + values = np.asarray(deltas, dtype=np.float64) + if values.ndim != 2 or values.shape[1] != 6 or not np.all(np.isfinite(values)): + raise ValueError("deltas must be a finite [T, 6] array") + if len(values) < 1: + raise ValueError("deltas must contain at least one step") + translation = np.cumsum(values[:, :3], axis=0) + increments = _rotvec_to_quaternion(values[:, 3:6]) + orientation = np.empty((len(values), 4), dtype=np.float64) + current = np.asarray([[0.0, 0.0, 0.0, 1.0]], dtype=np.float64) + for index, increment in enumerate(increments): + current = _quaternion_multiply(current, increment[None]) + orientation[index] = current[0] + return np.asarray( + np.concatenate((translation, orientation), axis=1), dtype=np.float32 + ) + + +def binarize_gripper_actions( + values: np.ndarray, + *, + threshold: float, + open_is_high: bool, +) -> np.ndarray: + """Map one controller gripper channel to canonical ``0=closed, 1=open``.""" + array = np.asarray(values, dtype=np.float32) + if array.ndim != 2 or array.shape[1] != 1 or not np.all(np.isfinite(array)): + raise ValueError("gripper action must be a finite [T, 1] array") + if not math.isfinite(threshold): + raise ValueError("gripper threshold must be finite") + opened = array >= threshold if open_is_high else array <= threshold + return opened.astype(np.float32) diff --git a/runtime/physical/trajectory.py b/runtime/physical/trajectory.py new file mode 100644 index 0000000000000000000000000000000000000000..6c9ef55a2d9188f780a2d609bacb26fbe316575c --- /dev/null +++ b/runtime/physical/trajectory.py @@ -0,0 +1,151 @@ +"""Variable-width trajectory contract shared by statistics and fitting.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass + +import numpy as np + +from ..common._validation import active_dimensions +from .configuration import ActionCodec2Config + +__all__ = ["ActiveTrajectory", "TrajectoryLike", "prepare_trajectory"] + + +@dataclass(frozen=True) +class ActiveTrajectory: + """One complete trajectory over an explicit subset of codebook dimensions. + + Args: + values (np.ndarray): ``(T, D_active)`` finite canonical values. ``T`` and + ``D_active`` may differ between trajectories. + active_dims: Strictly increasing global codebook dimension ids. This is + required when compact columns are not the global prefix, for example + ``[j1..j6, gripper]`` in a seven-joint-plus-gripper codebook. + previous_first_order (np.ndarray | None): ``(D_active,)`` preceding + reconstructed first-order values in canonical input units. Required + by second-order fitting; never inferred from a zero boundary. + """ + + values: np.ndarray + active_dims: tuple[int, ...] + previous_first_order: np.ndarray | None = None + + def __init__( + self, + values: np.ndarray, + active_dims: Sequence[int], + *, + previous_first_order: np.ndarray | None = None, + ) -> None: + object.__setattr__(self, "values", np.asarray(values)) + object.__setattr__(self, "active_dims", tuple(active_dims)) + previous = ( + None + if previous_first_order is None + else np.asarray(previous_first_order, dtype=np.float32) + ) + if previous is not None and ( + previous.shape != (len(self.active_dims),) + or not np.all(np.isfinite(previous)) + ): + raise ValueError( + "previous_first_order must be finite with shape [D_active]" + ) + object.__setattr__(self, "previous_first_order", previous) + + @classmethod + def from_padded( + cls, + values: np.ndarray, + *, + active_dims: Sequence[int], + previous_first_order: np.ndarray | None = None, + ) -> ActiveTrajectory: + """Drop engineering-only dense padding before the semantic boundary. + + Values in inactive columns are deliberately not inspected. They may be + zero, NaN, or any collator sentinel because they never reach physical + quantization, statistics, or Set-BPE. + """ + array = np.asarray(values) + dims = tuple(int(value) for value in active_dims) + if array.ndim != 2 or not dims or min(dims) < 0 or max(dims) >= array.shape[1]: + raise ValueError( + "padded values and active_dims do not define a valid trajectory" + ) + previous = ( + None + if previous_first_order is None + else np.take(previous_first_order, dims) + ) + return cls(np.take(array, dims, axis=1), dims, previous_first_order=previous) + + +def trajectory_boundary( + trajectory: TrajectoryLike, config: ActionCodec2Config, active_dims: tuple[int, ...] +) -> np.ndarray | None: + """Expand explicit compact state to ``[D]`` for fitting and validation.""" + previous = getattr(trajectory, "previous_first_order", None) + active = set(active_dims) + required = any( + c.primitive_order == 2 and active.intersection(c.dims) + for c in config.physical.components.values() + ) + if previous is None and required: + raise ValueError( + "second-order trajectories require explicit previous_first_order" + ) + if previous is None and not any( + c.primitive_order == 2 for c in config.physical.components.values() + ): + return None + dense = np.zeros(config.physical.input_dim, dtype=np.float32) + if previous is not None: + dense[list(active_dims)] = previous + return dense + + +TrajectoryLike = np.ndarray | ActiveTrajectory + + +def prepare_trajectory( + trajectory: TrajectoryLike, + *, + dimension: int, + index: int, + dtype: np.dtype | type = np.float32, +) -> tuple[np.ndarray, tuple[int, ...]]: + """Validate a trajectory and return compact values plus global dimension ids. + + A plain ``(T,D)`` array uses prefix dimensions ``0..D-1``. Callers must use + :class:`ActiveTrajectory` for a non-prefix layout. No padding value is part + of this public contract. + """ + if isinstance(trajectory, ActiveTrajectory): + raw = trajectory.values + active_dims = trajectory.active_dims + else: + raw = trajectory + array = np.asarray(raw) + active_dims = tuple(range(array.shape[1])) if array.ndim == 2 else () + + values = np.asarray(raw, dtype=dtype) + if values.ndim != 2 or len(values) < 1 or values.shape[1] < 1: + raise ValueError( + f"trajectory {index} must have non-empty shape [T, D_active], got {values.shape}" + ) + if len(active_dims) != values.shape[1]: + raise ValueError( + f"trajectory {index} active_dims has {len(active_dims)} entries for " + f"width {values.shape[1]}" + ) + if any(isinstance(item, bool) or not isinstance(item, int) for item in active_dims): + raise ValueError(f"trajectory {index} active_dims must contain integers") + active_dims = active_dimensions( + active_dims, dimension, subject=f"trajectory {index} active_dims" + ) + if not np.all(np.isfinite(values)): + raise ValueError(f"trajectory {index} contains a non-finite value") + return np.ascontiguousarray(values), active_dims diff --git a/runtime/processing_actioncodec2.py b/runtime/processing_actioncodec2.py new file mode 100644 index 0000000000000000000000000000000000000000..10ced2a08d5d7afbf61f56d785da3f5f7f03dc84 --- /dev/null +++ b/runtime/processing_actioncodec2.py @@ -0,0 +1,553 @@ +"""Hugging Face processor over the shared ActionCodec2 frontend implementation.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from transformers.feature_extraction_utils import BatchFeature +from transformers.processing_utils import ProcessorMixin + +from .routing.defaults import CodecBuilder, ProfileBundleBuilder, builtin_action_spaces +from .integration.hf import ( + PROCESSOR_CONFIG_NAME, + ROUTER_ALLOW_PATTERNS, + ActionCodec2FrontendHFConfig, + resolve_pretrained_directory, + save_huggingface_metadata, + save_model_card, + upload_pretrained_directory, +) +from .routing.presets import PresetConfig +from .routing.routing import ActionCodec2Router +from .frontend.tokenizer import ActionCodec2Tokenizer + + +def _saved_metadata(directory: Path) -> dict[str, Any]: + metadata = directory / PROCESSOR_CONFIG_NAME + return ( + json.loads(metadata.read_text(encoding="utf-8")) if metadata.is_file() else {} + ) + + +def _fitted_status(profiles) -> str: + status = ", ".join(sorted(profiles)) or "none; call codec.fit(episodes)" + return f"Fitted profiles: {status}\n" + + +class ActionCodec2(ActionCodec2Tokenizer, ProcessorMixin): + """Action-space-bound processor with lossless token serialization. + + ``encode`` returns ragged token-ID rows; ``decode`` reconstructs a CPU + ``torch.float32`` tensor in the original ``(B,T,D)`` column layout. + Calling the processor wraps those same rows in a ``BatchFeature`` under + ``input_ids``. Physical quantization remains lossy as configured. + + Released first- and second-order bundles both register 13 layouts: single + and dual EEF arms, single and dual arms with 6 or 7 joint coordinates, each + with absolute/delta variants and binary grippers, plus a six-continuous- + coordinate absolute layout without a separate binary gripper column. + The registered names and column meanings come from the loaded artifact. + ``print(codec)`` lists all of them; ``print(cls.describe_pretrained(path))`` + inspects them before binding. ``print_action_spaces(action_dim=D)`` only + filters candidates, because shape alone cannot identify physical meaning. + + Inputs accept ``(T,D)`` or ``(B,T,D)`` with any positive T; unequal-length + episodes must be encoded separately. Absolute layouts require the same + ``current_state(B,D)`` for encode/decode. Second-order physical profiles + additionally require ``previous_first_order: {component: (B,D_active)}`` + at the action space's configured codec rate, independently of the + absolute/delta choice. + Boundary arrays retain B=1 for a single episode. Physical order does not + change the caller's action layout into an acceleration tensor. + + Args: + backend: Optional fitted router. Omit to construct a fresh codec for fit(). + action_space: Registered name defining the input tensor's semantics. + primitive_order: Physical primitive order, 1 or 2, for fresh construction. + bpe_scheme: Optional 'set_bpe' or 'bpe' override for fresh construction. + profile_configs: Optional joint/EEF physical configs, mappings or YAML paths. + Defaults ship inside the package; no repository checkout is required. + """ + + model_type = ActionCodec2FrontendHFConfig.model_type + model_input_names = ["input_ids"] + # Numerical artifacts are owned by the router, not modality sub-processors. + attributes = [] + + def __init__( + self, + backend: ActionCodec2Router | None = None, + *, + action_space: str, + primitive_order: int = 1, + bpe_scheme: str | None = None, + profile_configs: Mapping[str, object] | None = None, + ) -> None: + fresh = backend is None + if fresh: + backend = CodecBuilder( + primitive_order=primitive_order, + bpe_scheme=bpe_scheme, + profile_configs=profile_configs, + ).build() + elif ( + primitive_order != 1 + or bpe_scheme is not None + or profile_configs is not None + ): + raise ValueError( + "profile construction options cannot accompany an existing backend" + ) + super().__init__(backend, action_space=action_space) + ProcessorMixin.__init__(self) + self._fitted_profiles = frozenset() if fresh else frozenset(backend.profiles) + self.training_report = None + self._source_directory = Path(__file__).parent + + def __dir__(self): + inapplicable = { + "action_tokenizer", + "attributes", + "audio_tokenizer", + "chat_template", + "check_argument_for_proper_class", + "encode_kwargs", + "feature_extractor_class", + "from_args_and_dict", + "get_possibly_dynamic_module", + "get_processor_dict", + "model_input_names", + "model_type", + "optional_attributes", + "optional_call_args", + "apply_chat_template", + "post_process_image_text_to_text", + "pad_action_id", + "bos_blk_id", + "parts_meta", + "BAR_SENTINEL_TOKENS", + "push_to_hub", + "register_for_auto_class", + "tokenizer_class", + "to_dict", + "to_json_file", + "to_json_string", + "valid_processor_kwargs", + "validate_init_kwargs", + } + return [name for name in super().__dir__() if name not in inapplicable] + + def __repr__(self) -> str: + return ( + f"ActionCodec2(action_space={self.action_space.name!r}, " + f"D={self.action_dim}, vocab={self.action_tokenizer.vocab_size})" + ) + + @staticmethod + def builtin_action_spaces(): + """Return the immutable named layout definitions shipped in the package.""" + return builtin_action_spaces() + + @classmethod + def from_profiles(cls, joint, eef, *, action_space: str, token_budgets=None): + """Bind an artifact assembled from existing joint/EEF profile directories. + + Args: + joint: Fitted joint profile directory. + eef: Fitted EEF profile directory. + action_space: Explicit built-in tensor layout to bind. + token_budgets: Optional joint/EEF overrides; defaults to recorded + training budgets, or fitted vocabulary sizes if unavailable. + """ + router = ProfileBundleBuilder(joint, eef, token_budgets=token_budgets).build() + return cls(router, action_space=action_space) + + def logits_processor( + self, + horizon: int, + *, + fps: float, + prompt_length: int, + eos_token_id: int, + token_offset: int = 0, + pad_token_id=None, + ): + """Create a Transformers LogitsProcessor for a fixed action horizon. + + ``prompt_length`` excludes leading model prompt IDs from the grammar; + ``token_offset`` explicitly maps codec IDs to a contiguous model range. + EOS/padding IDs must be outside that range. No model vocabulary is changed. + """ + from .tokenization.generation import ActionCodec2LogitsProcessor + + return ActionCodec2LogitsProcessor( + self, + horizon, + fps=fps, + prompt_length=prompt_length, + eos_token_id=eos_token_id, + token_offset=token_offset, + pad_token_id=pad_token_id, + ) + + def batch_decode( + self, + token_rows, + *, + fps=None, + current_state=None, + previous_first_order=None, + errors="raise", + ): + """Decode independent token rows, including different valid horizons. + + Args: + token_rows: Iterable of token-ID rows. Each must cover a complete + trajectory; intended generation horizon must be checked separately + with grammar.is_complete before decode. + fps: Output sampling rate, shared by all rows. + current_state: Optional (B,D) preceding states in original layout. + previous_first_order: Optional {component: (B,D_active)} boundaries. + errors: 'raise' stops at the first invalid row; 'return' returns the + ValueError at that position so other valid rows remain usable. + + Returns: + List of CPU float32 tensors shaped (1,T_i,D), or ValueError entries + when errors='return'. No truncated row is padded or executed partially. + """ + self._require_fitted() + if errors not in ("raise", "return"): + raise ValueError("errors must be 'raise' or 'return'") + rows = list(token_rows) + current = ( + None + if current_state is None + else np.asarray( + current_state.detach().cpu() + if isinstance(current_state, torch.Tensor) + else current_state + ) + ) + if current is not None and current.shape != (len(rows), self.action_dim): + raise ValueError( + f"current_state must have shape {(len(rows), self.action_dim)}" + ) + previous = self._numpy_boundary(previous_first_order) + if previous is not None and any( + np.ndim(v) != 2 or len(v) != len(rows) for v in previous.values() + ): + raise ValueError( + "previous_first_order arrays must have batch dimension equal to token rows" + ) + result = [] + for index, row in enumerate(rows): + try: + decoded = self.decode( + [row], + fps=fps, + current_state=None + if current is None + else current[index : index + 1], + previous_first_order=None + if previous is None + else {k: v[index : index + 1] for k, v in previous.items()}, + ) + except ValueError as error: + error = ValueError(f"token row {index}: {error}") + if errors == "raise": + raise error from None + decoded = error + result.append(decoded) + return result + + @property + def fitted_profiles(self) -> frozenset[str]: + """Names of profiles fitted on data; persisted by save_pretrained().""" + return self._fitted_profiles + + @property + def is_fitted(self) -> bool: + """Whether every profile used by the selected action space is fitted.""" + return all( + c.space in self.fitted_profiles for c in self.action_space.components + ) + + def _require_fitted(self): + missing = {c.space for c in self.action_space.components} - self.fitted_profiles + if missing: + raise RuntimeError( + f"Unfitted profiles {sorted(missing)}; call codec.fit(episodes) first" + ) + + def _describe_action_spaces(self, action_dim: int | None = None) -> str: + return _fitted_status(self.fitted_profiles) + super()._describe_action_spaces( + action_dim + ) + + def for_action_space(self, name: str) -> ActionCodec2: + """Share current weights with a new binding; subsequent fit replaces only that instance.""" + codec = super().for_action_space(name) + codec._fitted_profiles = self.fitted_profiles + codec.training_report = self.training_report + codec._source_directory = self._source_directory + return codec + + def fit( + self, + episodes, + *, + fps: float | None = None, + current_state=None, + previous_first_order=None, + bootstrap_context: bool = False, + backend: str = "auto", + threads: int = 0, + batch_size: int = 256, + progress: bool = False, + max_cell_bytes: int | None = None, + work_dir: str | Path | None = None, + ) -> ActionCodec2: + """Fit physical vocabularies from actions in declared source layouts. + + Args: + episodes: Array/tensor (T,D), (B,T,D), or iterable of variable-length + episodes. Records may contain ``actions``, ``action_space`` (a + registered name), ``fps``, ``current_state(B,D)``, and + ``previous_first_order: {component: (B,D_active)}``. A record's + metadata must not duplicate the corresponding keyword argument. + fps: Source frequency override shared by all input episodes. + current_state: Preceding source state (B,D), required for absolute + inputs. Use per-record values when episodes have different states. + previous_first_order: Preceding canonical codec-rate increments by component, + with shape (B,D_active), required for second-order profiles. + bootstrap_context: Use two measured canonical frames as second-order + context, training on the remaining suffix. Mutually exclusive with + previous_first_order. Inference still requires explicit boundaries. + backend: Exact BPE implementation: 'auto', 'native', or 'python'. + threads: Native worker count; 0 uses the backend default. + batch_size: Maximum canonical trajectories per quantization batch. + progress: Display BPE training progress. + max_cell_bytes: Hard bound on retained quantized corpus bytes per profile. + work_dir: Parent scratch directory for sequential source staging files, + automatically removed on success or failure. Defaults to system temp. + + Returns: + This instance, with per-profile results in ``training_report``. Only + profiles present in the data are fitted; other weights are preserved. + Vocabulary replacement is atomic on success. Refit changes token IDs: + regenerate previously encoded datasets before using the new vocabulary. + """ + # Training is intentionally lazy: a loaded Hub artifact only needs the + # inference runtime. The full fitting pipeline remains available from + # the source package and is imported when the user calls fit(). + from .training.fitting import ActionCodec2Fitter + + router, report = ActionCodec2Fitter( + self, bootstrap_context=bootstrap_context + ).fit( + episodes, + work_dir=work_dir, + metadata=dict( + fps=fps, + current_state=current_state, + previous_first_order=previous_first_order, + ), + backend=backend, + threads=threads, + batch_size=batch_size, + progress=progress, + max_cell_bytes=max_cell_bytes, + ) + fitted = self.fitted_profiles | report["profiles"].keys() + replacement = type(self)(router, action_space=self.action_space.name) + replacement._fitted_profiles = frozenset(fitted) + replacement.training_report = report + replacement._source_directory = self._source_directory + replacement._warned_default_fps = self._warned_default_fps + self.__dict__.update(replacement.__dict__) + return self + + def _encode_action_indices(self, action, encode_kwargs=None): + self._require_fitted() + return super()._encode_action_indices(action, encode_kwargs) + + def _decode_action_indices(self, action_indices, **kwargs): + self._require_fitted() + return super()._decode_action_indices(action_indices, **kwargs) + + def grammar(self, horizon: int, *, fps: float | None = None): + self._require_fitted() + return super().grammar(horizon, fps=fps) + + def final_first_order( + self, action_token_ids, *, previous_first_order, executed_steps=None + ): + self._require_fitted() + return super().final_first_order( + action_token_ids, + previous_first_order=previous_first_order, + executed_steps=executed_steps, + ) + + @classmethod + def describe_pretrained( + cls, + pretrained_model_name_or_path: str | Path, + *, + action_dim: int | None = None, + **kwargs: Any, + ) -> str: + """Inspect a local/Hub artifact before choosing an action space. + + Args: + pretrained_model_name_or_path: Saved processor/router directory or + Hugging Face Hub repository ID. Loads its fitted router for inspection. + action_dim: Optional D from ``actions.shape[-1]``. Only filters the + report; tensor shape cannot establish physical semantics. + **kwargs: Hub options such as ``revision``, ``subfolder``, + ``cache_dir``, ``token``, and ``local_files_only``. + + Returns: + Human-readable layout guide, suitable for ``print(...)``. No binding + is required or inferred, and no action data is encoded. + """ + directory = resolve_pretrained_directory( + pretrained_model_name_or_path, + allow_patterns=ROUTER_ALLOW_PATTERNS, + **kwargs, + ) + router = ActionCodec2Router.from_pretrained(directory) + metadata = _saved_metadata(directory) + return _fitted_status( + metadata.get("fitted_profiles", router.profiles) + ) + router.describe_action_spaces( + selected=metadata.get("action_space"), action_dim=action_dim + ) + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str | Path, + *, + action_space: str | Path | PresetConfig | Mapping[str, object] | None = None, + **kwargs: Any, + ) -> ActionCodec2: + """Load a local/Hub artifact and restore or override its action space. + + Args: + pretrained_model_name_or_path: Saved processor or router directory, + or a Hugging Face Hub repository ID. + action_space: Registered name, preset YAML path, config, or mapping. + Required for an unbound router; otherwise defaults to the saved + binding. Custom definitions are saved in the router manifest. + **kwargs: Hub options including ``revision``, ``subfolder``, + ``cache_dir``, ``token``, and ``local_files_only``. + """ + directory = resolve_pretrained_directory( + pretrained_model_name_or_path, + allow_patterns=ROUTER_ALLOW_PATTERNS, + **kwargs, + ) + if action_space is None: + action_space = _saved_metadata(directory).get("action_space") + if action_space is None: + raise ValueError( + "action_space is required when loading an unbound router artifact. " + "First inspect its layouts with print(ActionCodec2.describe_pretrained(path)), " + "then pass action_space='NAME'; shape alone cannot identify physical meaning." + ) + codec = super().from_pretrained(directory, action_space=action_space) + metadata = _saved_metadata(directory) + fitted = metadata.get("fitted_profiles", list(codec.action_tokenizer.profiles)) + if not isinstance(fitted, list) or any( + not isinstance(name, str) or name not in codec.action_tokenizer.profiles + for name in fitted + ): + raise ValueError("invalid fitted_profiles in processor_config.json") + codec._fitted_profiles = frozenset(fitted) + codec._source_directory = directory + report_path = directory / "fit_report.json" + if report_path.is_file(): + codec.training_report = json.loads(report_path.read_text(encoding="utf-8")) + return codec + + def to_dict(self, legacy_serialization: bool = True) -> dict[str, object]: + """Return lightweight discovery metadata. + + ``legacy_serialization`` is accepted for ProcessorMixin signature + compatibility. ActionCodec2 has one serialization format for both values. + """ + del legacy_serialization + return { + "processor_class": type(self).__name__, + "model_type": self.model_type, + "action_space": self.action_space.name, + "fitted_profiles": sorted(self.fitted_profiles), + } + + def save_pretrained( + self, save_directory: str | Path, push_to_hub: bool = False, **kwargs: Any + ) -> list[str]: + """Save the complete router and the selected action-space binding.""" + if not push_to_hub and kwargs: + raise TypeError( + f"unused save_pretrained keyword(s): {', '.join(sorted(kwargs))}" + ) + from .integration.exporting import ArtifactExporter + + directory = Path(save_directory) + save_model_card(directory, str(self)) + files = self.action_tokenizer.save_pretrained(directory, include_code=False) + files.extend( + ArtifactExporter( + directory, source_directory=self._source_directory + ).export() + ) + save_huggingface_metadata( + directory, + **self.to_dict(), + auto_map={"AutoProcessor": "processing_actioncodec2.ActionCodec2"}, + ) + report_path = directory / "fit_report.json" + if self.training_report is not None: + report_path.write_text( + json.dumps(self.training_report, indent=2) + "\n", encoding="utf-8" + ) + files.append(str(report_path)) + elif report_path.exists(): + report_path.unlink() + if push_to_hub: + upload_pretrained_directory(directory, **kwargs) + return files + + def __call__( + self, + action: torch.Tensor | np.ndarray, + *, + fps: float | None = None, + current_state: torch.Tensor | np.ndarray | None = None, + previous_first_order: Mapping[str, torch.Tensor | np.ndarray] | None = None, + encode_kwargs: dict | None = None, + return_tensors: str | None = None, + ) -> BatchFeature: + """Prepare ``input_ids`` from ``(B,T,D)`` or a single ``(T,D)`` action. + + ``current_state(B,D)`` is required for absolute actions. Second-order + artifacts require ``previous_first_order`` mapping component names to + ``(B,D_active)`` canonical increments, as in :meth:`encode`. + ``return_tensors`` may be ``"pt"`` or ``"np"`` for equal-length rows; + omitted, rows remain ragged lists without adding padding tokens. + """ + rows = self.encode( + action, + fps=fps, + current_state=current_state, + previous_first_order=previous_first_order, + encode_kwargs=encode_kwargs, + ) + return BatchFeature({"input_ids": rows}, tensor_type=return_tensors) diff --git a/runtime/routing/__init__.py b/runtime/routing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7cb6b89d66b1e8d01e229067c0c6a1adf3976690 --- /dev/null +++ b/runtime/routing/__init__.py @@ -0,0 +1 @@ +"""Internal routing components for ActionCodec2.""" diff --git a/runtime/routing/cache.py b/runtime/routing/cache.py new file mode 100644 index 0000000000000000000000000000000000000000..d375b03c2fa5982882764c5abf071272ce43aa43 --- /dev/null +++ b/runtime/routing/cache.py @@ -0,0 +1,45 @@ +"""Small per-instance LRU caches; no global references to fitted codecs.""" + +from collections import OrderedDict +from collections.abc import MutableMapping +from threading import RLock + + +class BoundedCache(MutableMapping): + """Evict least-recently-used entries above a fixed capacity. + + Args: + capacity: Maximum retained entries (positive integer). + """ + + def __init__(self, capacity): + if type(capacity) is not int or capacity < 1: + raise ValueError("cache capacity must be a positive integer") + self.capacity = capacity + self._items = OrderedDict() + self._lock = RLock() + + def __getitem__(self, key): + with self._lock: + value = self._items[key] + self._items.move_to_end(key) + return value + + def __setitem__(self, key, value): + with self._lock: + self._items[key] = value + self._items.move_to_end(key) + while len(self._items) > self.capacity: + self._items.popitem(last=False) + + def __delitem__(self, key): + with self._lock: + del self._items[key] + + def __iter__(self): + with self._lock: + return iter(tuple(self._items)) + + def __len__(self): + with self._lock: + return len(self._items) diff --git a/runtime/routing/defaults.py b/runtime/routing/defaults.py new file mode 100644 index 0000000000000000000000000000000000000000..696afff7f0e36331ab9c88eacf9e387c1a0072e4 --- /dev/null +++ b/runtime/routing/defaults.py @@ -0,0 +1,622 @@ +"""Package-owned configuration templates for fresh ActionCodec2 instances.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from ..tokenization.bpe import BPESerializer, SetBPESerializer +from ..physical.configuration import ActionCodec2Config +from ..physical.physical_processor import ActionCodec2Processor +from .presets import PresetConfig +from .routing import ActionCodec2Router + + +class CodecBuilder: + """Build unfitted physical profiles from packaged defaults or explicit configs. + + Args: + primitive_order: First- or second-order default physical primitives. + bpe_scheme: Optional scheme override for both physical profiles. + profile_configs: Optional joint/EEF config objects, mappings, or YAML paths. + """ + + def __init__(self, *, primitive_order=1, bpe_scheme=None, profile_configs=None): + if type(primitive_order) is not int or primitive_order not in (1, 2): + raise ValueError("primitive_order must be 1 or 2") + if bpe_scheme not in (None, "set_bpe", "bpe"): + raise ValueError("bpe_scheme must be 'set_bpe' or 'bpe'") + if profile_configs is not None and not isinstance(profile_configs, Mapping): + raise TypeError( + "profile_configs must map 'joint'/'eef' to physical configurations" + ) + self.profile_configs = dict(profile_configs or {}) + if set(self.profile_configs) - {"joint", "eef"}: + raise ValueError("profile_configs keys must be 'joint' or 'eef'") + self.primitive_order = primitive_order + self.bpe_scheme = bpe_scheme + + def _config(self, name: str) -> ActionCodec2Config: + suffix = "" if self.primitive_order == 1 else "_second_order_pretrain_v2" + raw = self.profile_configs.get(name, _PROFILE_DEFAULTS[f"{name}{suffix}"]) + if isinstance(raw, ActionCodec2Config): + config = raw + elif isinstance(raw, Mapping): + config = ActionCodec2Config.from_dict(raw) + else: + config = ActionCodec2Config.from_yaml(raw) + if config.train is None: + raise ValueError(f"{name} profile config requires a train section") + if self.bpe_scheme is not None: + payload = config.to_dict() + payload["set_bpe"]["scheme"] = self.bpe_scheme + if self.bpe_scheme == "bpe": + payload["train"]["bpe"]["spatial_vocab_size"] = payload["train"]["bpe"][ + "vocab_size" + ] + config = ActionCodec2Config.from_dict(payload) + return config + + def build(self) -> ActionCodec2Router: + """Create primitive-only profiles; the public facade tracks fitted state.""" + processors = {} + budgets = {} + for name in ("joint", "eef"): + config = self._config(name) + # Keep training-only modules out of the inference artifact. They + # are imported only when a fresh codec is constructed for fitting. + from ..training.training import build_primitives + + primitives = build_primitives(config) + if config.set_bpe.scheme == "bpe": + serializer = BPESerializer( + bins=primitives.bins, lookahead=config.set_bpe.lookahead, merges=() + ) + else: + serializer = SetBPESerializer( + bins=primitives.bins, + lookahead=config.set_bpe.lookahead, + spatial_merges=(), + temporal_merges=(), + ) + processors[name] = ActionCodec2Processor(config, primitives, serializer) + budgets[name] = config.train.bpe.vocab_size + presets = list(builtin_action_spaces().values()) + return ActionCodec2Router.from_processors( + processors, token_budgets=budgets, presets=presets + ) + + +def builtin_action_spaces(): + """Return a read-only name-to-preset mapping shipped in the installed package.""" + from types import MappingProxyType + + presets = [PresetConfig.from_dict(value) for value in _PRESET_DEFAULTS.values()] + return MappingProxyType({preset.name: preset for preset in presets}) + + +class ProfileBundleBuilder: + """Assemble two existing physical artifacts using their recorded vocabulary budgets. + + Args: + joint: Local directory containing the fitted joint profile. + eef: Local directory containing the fitted EEF profile. + token_budgets: Optional explicit joint/EEF budget overrides. Defaults to + recorded training budgets, or the fitted vocabulary size if absent. + """ + + def __init__(self, joint, eef, *, token_budgets=None): + self.paths = {"joint": joint, "eef": eef} + self.budgets = dict(token_budgets or {}) + if self.budgets.keys() - self.paths.keys(): + raise ValueError("token_budgets accepts only 'joint' and 'eef'") + + def build(self): + processors = { + name: ActionCodec2Processor.from_pretrained(path) + for name, path in self.paths.items() + } + budgets = { + name: self.budgets.get( + name, + processor.config.train.bpe.vocab_size + if processor.config.train is not None + else processor.vocab_size, + ) + for name, processor in processors.items() + } + return ActionCodec2Router.from_processors( + processors, + token_budgets=budgets, + presets=list(builtin_action_spaces().values()), + ) + + +_PROFILE_DEFAULTS = { + "eef": { + "model_type": "actioncodec2_unified", + "config_version": 1, + "physical": { + "input_dim": 7, + "components": { + "translation_m": { + "dims": {"start": 0, "stop": 3}, + "space": "additive", + "input_order": 0, + "primitive_order": 1, + "input_transform": {"scale": 1.0, "clip": False}, + "constraints": { + "order_0": {"min": -2.0, "max": 2.0}, + "order_1": {"min": -0.08, "max": 0.08}, + }, + "distortion_limit": 0.005, + "primitive_selection": "smooth", + }, + "orientation_so3": { + "dims": {"start": 3, "stop": 6}, + "space": "so3", + "input_order": 0, + "primitive_order": 1, + "input_transform": {"scale": 1.0, "clip": False}, + "constraints": { + "order_0": { + "min": -3.141592653589793, + "max": 3.141592653589793, + }, + "order_1": {"min": -0.25, "max": 0.25}, + }, + "distortion_limit": 0.02, + "primitive_selection": "smooth", + "representation": "rotation_vector", + "multiplication": "right", + }, + "gripper_open_command": { + "dims": [6], + "space": "additive", + "input_order": 0, + "primitive_order": 0, + "input_transform": {"scale": 1.0, "clip": True}, + "constraints": {"order_0": {"min": 0.0, "max": 1.0}}, + "distortion_limit": 0.5, + "primitive_selection": "nearest", + }, + }, + }, + "set_bpe": {"lookahead": 8, "allow_cross_component_merge": False}, + "train": { + "primitives": { + "zero_primitive_components": ["translation_m", "orientation_so3"] + }, + "bpe": {"spatial_vocab_size": 2048, "vocab_size": 4096}, + "quality_control": {"hard_order_1_scale": 3.0, "min_segment_length": 2}, + }, + }, + "eef_second_order_pretrain_v2": { + "model_type": "actioncodec2_unified", + "config_version": 1, + "physical": { + "input_dim": 7, + "components": { + "translation_m": { + "dims": {"start": 0, "stop": 3}, + "space": "additive", + "input_order": 0, + "primitive_order": 2, + "input_transform": {"scale": 1.0, "clip": False}, + "constraints": { + "order_0": {"min": -2.0, "max": 2.0}, + "order_2": {"min": -0.025, "max": 0.025}, + "order_1": {"min": -0.08, "max": 0.08}, + }, + "distortion_limit": 0.005, + "primitive_selection": "smooth", + }, + "orientation_so3": { + "dims": {"start": 3, "stop": 6}, + "space": "so3", + "input_order": 0, + "primitive_order": 2, + "input_transform": {"scale": 1.0, "clip": False}, + "constraints": { + "order_0": { + "min": -3.141592653589793, + "max": 3.141592653589793, + }, + "order_2": {"min": -0.08, "max": 0.08}, + "order_1": {"min": -0.25, "max": 0.25}, + }, + "distortion_limit": 0.02, + "primitive_selection": "smooth", + "representation": "rotation_vector", + "multiplication": "right", + }, + "gripper_open_command": { + "dims": [6], + "space": "additive", + "input_order": 0, + "primitive_order": 0, + "input_transform": {"scale": 1.0, "clip": True}, + "constraints": {"order_0": {"min": 0.0, "max": 1.0}}, + "distortion_limit": 0.5, + "primitive_selection": "nearest", + }, + }, + }, + "set_bpe": {"lookahead": 8, "allow_cross_component_merge": False}, + "train": { + "primitives": { + "zero_primitive_components": ["translation_m", "orientation_so3"] + }, + "bpe": {"spatial_vocab_size": 2048, "vocab_size": 4096}, + "quality_control": {"hard_order_1_scale": 3.0, "min_segment_length": 2}, + }, + }, + "joint": { + "model_type": "actioncodec2_unified", + "config_version": 1, + "physical": { + "input_dim": 8, + "components": { + "generalized_joint_position": { + "dims": {"start": 0, "stop": 7}, + "space": "additive", + "input_order": 0, + "primitive_order": 1, + "input_transform": {"scale": 1.0, "clip": False}, + "constraints": { + "order_0": { + "min": -6.283185307179586, + "max": 6.283185307179586, + }, + "order_1": {"min": -0.35, "max": 0.35}, + }, + "distortion_limit": [ + 0.004, + 0.004, + 0.006, + 0.0075, + 0.015, + 0.02, + 0.02, + ], + "primitive_selection": "smooth", + }, + "gripper_open_command": { + "dims": [7], + "space": "additive", + "input_order": 0, + "primitive_order": 0, + "input_transform": {"scale": 1.0, "clip": True}, + "constraints": {"order_0": {"min": 0.0, "max": 1.0}}, + "distortion_limit": 0.5, + "primitive_selection": "nearest", + }, + }, + }, + "set_bpe": {"lookahead": 8, "allow_cross_component_merge": False}, + "train": { + "primitives": {"zero_primitive_components": ["generalized_joint_position"]}, + "bpe": {"spatial_vocab_size": 2048, "vocab_size": 4096}, + "quality_control": {"hard_order_1_scale": 3.0, "min_segment_length": 2}, + }, + }, + "joint_second_order_pretrain_v2": { + "model_type": "actioncodec2_unified", + "config_version": 1, + "physical": { + "input_dim": 8, + "components": { + "generalized_joint_position": { + "dims": {"start": 0, "stop": 7}, + "space": "additive", + "input_order": 0, + "primitive_order": 2, + "input_transform": {"scale": 1.0, "clip": False}, + "constraints": { + "order_0": { + "min": -6.283185307179586, + "max": 6.283185307179586, + }, + "order_2": {"min": -0.25, "max": 0.25}, + "order_1": {"min": -0.35, "max": 0.35}, + }, + "distortion_limit": [ + 0.004, + 0.004, + 0.006, + 0.0075, + 0.015, + 0.02, + 0.02, + ], + "primitive_selection": "smooth", + }, + "gripper_open_command": { + "dims": [7], + "space": "additive", + "input_order": 0, + "primitive_order": 0, + "input_transform": {"scale": 1.0, "clip": True}, + "constraints": {"order_0": {"min": 0.0, "max": 1.0}}, + "distortion_limit": 0.5, + "primitive_selection": "nearest", + }, + }, + }, + "set_bpe": {"lookahead": 8, "allow_cross_component_merge": False}, + "train": { + "primitives": {"zero_primitive_components": ["generalized_joint_position"]}, + "bpe": {"spatial_vocab_size": 2048, "vocab_size": 4096}, + "quality_control": {"hard_order_1_scale": 3.0, "min_segment_length": 2}, + }, + }, +} + +_PRESET_DEFAULTS = { + "dual_eef_abs": { + "name": "dual_eef_abs", + "input_dim": 14, + "input": "absolute", + "source_hz": 15.0, + "target_hz": 15.0, + "interpolation": "cubic", + "components": { + "left_arm": { + "space": "eef", + "position_dims": [0, 1, 2], + "rotation_dims": [3, 4, 5], + "rotation_format": "rotation_vector", + "gripper_dim": 6, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + }, + "right_arm": { + "space": "eef", + "position_dims": [7, 8, 9], + "rotation_dims": [10, 11, 12], + "rotation_format": "rotation_vector", + "gripper_dim": 13, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + }, + }, + }, + "dual_eef_delta": { + "name": "dual_eef_delta", + "input_dim": 14, + "input": "delta", + "source_hz": 15.0, + "target_hz": 15.0, + "interpolation": "cubic", + "components": { + "left_arm": { + "space": "eef", + "position_dims": [0, 1, 2], + "rotation_dims": [3, 4, 5], + "rotation_format": "rotation_vector", + "gripper_dim": 6, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + }, + "right_arm": { + "space": "eef", + "position_dims": [7, 8, 9], + "rotation_dims": [10, 11, 12], + "rotation_format": "rotation_vector", + "gripper_dim": 13, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + }, + }, + }, + "dual_joint6_abs": { + "name": "dual_joint6_abs", + "input_dim": 14, + "input": "absolute", + "source_hz": 15.0, + "target_hz": 15.0, + "interpolation": "cubic", + "components": { + "left_arm": { + "space": "joint", + "joint_dims": [0, 1, 2, 3, 4, 5], + "gripper_dim": 6, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + }, + "right_arm": { + "space": "joint", + "joint_dims": [7, 8, 9, 10, 11, 12], + "gripper_dim": 13, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + }, + }, + }, + "dual_joint6_delta": { + "name": "dual_joint6_delta", + "input_dim": 14, + "input": "delta", + "source_hz": 15.0, + "target_hz": 15.0, + "interpolation": "cubic", + "components": { + "left_arm": { + "space": "joint", + "joint_dims": [0, 1, 2, 3, 4, 5], + "gripper_dim": 6, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + }, + "right_arm": { + "space": "joint", + "joint_dims": [7, 8, 9, 10, 11, 12], + "gripper_dim": 13, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + }, + }, + }, + "dual_joint7_abs": { + "name": "dual_joint7_abs", + "input_dim": 16, + "input": "absolute", + "source_hz": 15.0, + "target_hz": 15.0, + "interpolation": "cubic", + "components": { + "left_arm": { + "space": "joint", + "joint_dims": [0, 1, 2, 3, 4, 5, 6], + "gripper_dim": 7, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + }, + "right_arm": { + "space": "joint", + "joint_dims": [8, 9, 10, 11, 12, 13, 14], + "gripper_dim": 15, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + }, + }, + }, + "dual_joint7_delta": { + "name": "dual_joint7_delta", + "input_dim": 16, + "input": "delta", + "source_hz": 15.0, + "target_hz": 15.0, + "interpolation": "cubic", + "components": { + "left_arm": { + "space": "joint", + "joint_dims": [0, 1, 2, 3, 4, 5, 6], + "gripper_dim": 7, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + }, + "right_arm": { + "space": "joint", + "joint_dims": [8, 9, 10, 11, 12, 13, 14], + "gripper_dim": 15, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + }, + }, + }, + "single_eef_abs": { + "name": "single_eef_abs", + "input_dim": 7, + "input": "absolute", + "source_hz": 15.0, + "target_hz": 15.0, + "interpolation": "cubic", + "components": { + "arm": { + "space": "eef", + "position_dims": [0, 1, 2], + "rotation_dims": [3, 4, 5], + "rotation_format": "rotation_vector", + "gripper_dim": 6, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + } + }, + }, + "single_eef_delta": { + "name": "single_eef_delta", + "input_dim": 7, + "input": "delta", + "source_hz": 15.0, + "target_hz": 15.0, + "interpolation": "cubic", + "components": { + "arm": { + "space": "eef", + "position_dims": [0, 1, 2], + "rotation_dims": [3, 4, 5], + "rotation_format": "rotation_vector", + "gripper_dim": 6, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + } + }, + }, + "single_joint6_abs": { + "name": "single_joint6_abs", + "input_dim": 7, + "input": "absolute", + "source_hz": 15.0, + "target_hz": 15.0, + "interpolation": "cubic", + "components": { + "arm": { + "space": "joint", + "joint_dims": [0, 1, 2, 3, 4, 5], + "gripper_dim": 6, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + } + }, + }, + "single_joint6_abs_continuous_gripper": { + "name": "single_joint6_abs_continuous_gripper", + "input_dim": 6, + "input": "absolute", + "source_hz": 15.0, + "target_hz": 15.0, + "interpolation": "cubic", + "components": {"arm": {"space": "joint", "joint_dims": [0, 1, 2, 3, 4, 5]}}, + }, + "single_joint6_delta": { + "name": "single_joint6_delta", + "input_dim": 7, + "input": "delta", + "source_hz": 15.0, + "target_hz": 15.0, + "interpolation": "cubic", + "components": { + "arm": { + "space": "joint", + "joint_dims": [0, 1, 2, 3, 4, 5], + "gripper_dim": 6, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + } + }, + }, + "single_joint7_abs": { + "name": "single_joint7_abs", + "input_dim": 8, + "input": "absolute", + "source_hz": 15.0, + "target_hz": 15.0, + "interpolation": "cubic", + "components": { + "arm": { + "space": "joint", + "joint_dims": [0, 1, 2, 3, 4, 5, 6], + "gripper_dim": 7, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + } + }, + }, + "single_joint7_delta": { + "name": "single_joint7_delta", + "input_dim": 8, + "input": "delta", + "source_hz": 15.0, + "target_hz": 15.0, + "interpolation": "cubic", + "components": { + "arm": { + "space": "joint", + "joint_dims": [0, 1, 2, 3, 4, 5, 6], + "gripper_dim": 7, + "gripper_threshold": 0.8, + "gripper_open_is_high": True, + } + }, + }, +} diff --git a/runtime/routing/presets.py b/runtime/routing/presets.py new file mode 100644 index 0000000000000000000000000000000000000000..e116c09370c529c67244fbbcc169af727f0fba2c --- /dev/null +++ b/runtime/routing/presets.py @@ -0,0 +1,872 @@ +"""User-facing action layouts compiled into shared joint/EEF trajectories.""" + +from __future__ import annotations + +import re +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import yaml + +from ..common._validation import optional_finite_array, validate_dimension_partition +from ..physical.configuration import ( + ConfigError, + _integer, + _mapping, + _number, + _reject_unknown, + _required, + _UniqueKeyLoader, +) +from ..physical.semantic import ( + binarize_gripper_actions, + eef_pose_absolute, + eef_pose_components, + eef_pose_deltas, + eef_pose_from_components, + eef_pose_from_local_absolute, + eef_pose_local_absolute, + integrate_eef_deltas, + resample_uniform, +) +from ..physical.trajectory import ActiveTrajectory + +_NAME = re.compile(r"[a-z][a-z0-9_]*\Z") +_SPACES = {"eef", "joint"} +_INPUTS = {"absolute", "delta"} +_ROTATIONS = {"rotation_vector", "euler_xyz", "quaternion_xyzw"} +_INTERPOLATIONS = {"linear", "pchip", "cubic"} +_GRIPPER_DIM = {"eef": 6, "joint": 7} + + +def _format_dims(dims: Sequence[int]) -> str: + values = tuple(dims) + if len(values) == 1: + return str(values[0]) + if values == tuple(range(values[0], values[-1] + 1)): + return f"{values[0]}..{values[-1]}" + return ",".join(map(str, values)) + + +def _dims(value: object, path: str) -> tuple[int, ...]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes, bytearray)): + raise ConfigError(f"{path} must be a list of dimensions") + output = tuple( + _integer(item, f"{path}[{index}]") for index, item in enumerate(value) + ) + if not output: + raise ConfigError(f"{path} must not be empty") + if len(set(output)) != len(output): + raise ConfigError(f"{path} contains a repeated dimension") + return output + + +@dataclass(frozen=True) +class ComponentSpec: + """One independently encoded arm/body component in an input tensor. + + ``dims`` always refer to source tensor columns. Component names separate + BPE boundaries and output ordering, but never select primitive/BPE state. + The ``space`` field alone routes values to the shared joint or EEF profile. + + Args: + name: Component label, such as ``arm``, ``left_arm``, or ``right_arm``; + also the key for component-wise first-order boundaries. + space: ``joint`` for generalized joint coordinates, or ``eef`` for + end-effector position and orientation. + joint_dims: Ordered source columns for 1..7 joint coordinates. Revolute + coordinates use radians; other quantities need a canonical mapping. + position_dims: Three source columns for EEF x, y, z translation in meters. + rotation_dims: Source rotation columns, in the order of ``rotation_format``. + rotation_format: ``rotation_vector`` (axis times angle in radians), + ``euler_xyz`` (extrinsic XYZ angles in radians), or + ``quaternion_xyzw`` (unit quaternion). EEF delta layouts require + rotation vectors interpreted as right-multiplicative increments. + gripper_dim: Optional binary controller-command column. Omitted means + no separate binary gripper channel; it does not add a source column. + gripper_threshold: Inclusive threshold for open. Defaults to 0.8 when + high means open, or 0.2 when low means open. + gripper_open_is_high: If true, values >= threshold mean open; otherwise + values <= threshold mean open. Decode uses 1=open and 0=closed in + both absolute and delta layouts. + """ + + name: str + space: str + joint_dims: tuple[int, ...] | None = None + position_dims: tuple[int, ...] | None = None + rotation_dims: tuple[int, ...] | None = None + rotation_format: str = "rotation_vector" + gripper_dim: int | None = None + gripper_threshold: float | None = None + gripper_open_is_high: bool = True + + def __post_init__(self) -> None: + if not isinstance(self.name, str) or _NAME.fullmatch(self.name) is None: + raise ConfigError(f"component name must match {_NAME.pattern!r}") + if self.space not in _SPACES: + raise ConfigError(f"component {self.name!r}.space must be joint or eef") + if self.rotation_format not in _ROTATIONS: + raise ConfigError( + f"component {self.name!r}.rotation_format must be one of {sorted(_ROTATIONS)}" + ) + for label, dims in ( + ("joint_dims", self.joint_dims), + ("position_dims", self.position_dims), + ("rotation_dims", self.rotation_dims), + ): + if dims is not None: + _dims(dims, f"component {self.name}.{label}") + self._validate_motion_dims() + raw_dims = self.source_dims + if self.gripper_dim is not None: + _integer(self.gripper_dim, f"component {self.name}.gripper_dim") + if self.gripper_dim in raw_dims: + raise ConfigError( + f"component {self.name!r}.gripper_dim overlaps motion dims" + ) + if not isinstance(self.gripper_open_is_high, bool): + raise ConfigError( + f"component {self.name!r}.gripper_open_is_high must be boolean" + ) + threshold = self.gripper_threshold + if threshold is None: + threshold = 0.8 if self.gripper_open_is_high else 0.2 + object.__setattr__(self, "gripper_threshold", threshold) + _number(threshold, f"component {self.name}.gripper_threshold") + + def _validate_motion_dims(self) -> None: + if self.space == "joint": + if self.joint_dims is None or not 1 <= len(self.joint_dims) <= 7: + raise ConfigError( + f"joint component {self.name!r} requires 1..7 joint_dims" + ) + if self.position_dims is not None or self.rotation_dims is not None: + raise ConfigError( + f"joint component {self.name!r} cannot define EEF dims" + ) + return + if self.joint_dims is not None: + raise ConfigError(f"EEF component {self.name!r} cannot define joint_dims") + if self.position_dims is None or len(self.position_dims) != 3: + raise ConfigError( + f"EEF component {self.name!r} requires three position_dims" + ) + expected = 4 if self.rotation_format == "quaternion_xyzw" else 3 + if self.rotation_dims is None or len(self.rotation_dims) != expected: + raise ConfigError( + f"EEF component {self.name!r} requires {expected} rotation_dims for " + f"{self.rotation_format}" + ) + if set(self.position_dims).intersection(self.rotation_dims): + raise ConfigError(f"EEF component {self.name!r} reuses a pose dimension") + + @property + def source_dims(self) -> tuple[int, ...]: + """All source motion columns in canonical component order.""" + if self.space == "joint": + assert self.joint_dims is not None + return self.joint_dims + assert self.position_dims is not None and self.rotation_dims is not None + return (*self.position_dims, *self.rotation_dims) + + @property + def all_source_dims(self) -> tuple[int, ...]: + """Motion plus optional binary-gripper source columns.""" + return ( + self.source_dims + if self.gripper_dim is None + else (*self.source_dims, self.gripper_dim) + ) + + @property + def active_dims(self) -> tuple[int, ...]: + """Shared-profile slots emitted by this component.""" + motion_width = len(self.joint_dims or ()) if self.space == "joint" else 6 + dims = tuple(range(motion_width)) + return dims if self.gripper_dim is None else (*dims, _GRIPPER_DIM[self.space]) + + @classmethod + def from_dict(cls, name: str, payload: object, *, path: str) -> ComponentSpec: + values = _mapping(payload, path) + allowed = { + "space", + "joint_dims", + "position_dims", + "rotation_dims", + "rotation_format", + "gripper_dim", + "gripper_threshold", + "gripper_open_is_high", + } + _reject_unknown(values, allowed, path) + gripper_dim = values.get("gripper_dim") + if gripper_dim is not None: + gripper_dim = _integer(gripper_dim, f"{path}.gripper_dim") + return cls( + name=name, + space=str(_required(values, "space", path)), + joint_dims=None + if "joint_dims" not in values + else _dims(values["joint_dims"], f"{path}.joint_dims"), + position_dims=None + if "position_dims" not in values + else _dims(values["position_dims"], f"{path}.position_dims"), + rotation_dims=None + if "rotation_dims" not in values + else _dims(values["rotation_dims"], f"{path}.rotation_dims"), + rotation_format=str(values.get("rotation_format", "rotation_vector")), + gripper_dim=gripper_dim, + gripper_threshold=( + None + if values.get("gripper_threshold") is None + else _number(values["gripper_threshold"], f"{path}.gripper_threshold") + ), + gripper_open_is_high=values.get("gripper_open_is_high", True), + ) + + def to_dict(self) -> dict[str, object]: + payload: dict[str, object] = {"space": self.space} + if self.space == "joint": + payload["joint_dims"] = list(self.joint_dims or ()) + else: + payload.update( + { + "position_dims": list(self.position_dims or ()), + "rotation_dims": list(self.rotation_dims or ()), + "rotation_format": self.rotation_format, + } + ) + if self.gripper_dim is not None: + payload.update( + { + "gripper_dim": self.gripper_dim, + "gripper_threshold": self.gripper_threshold, + "gripper_open_is_high": self.gripper_open_is_high, + } + ) + return payload + + +@dataclass(frozen=True) +class PresetConfig: + """A dataset- and embodiment-independent physical ``[B,T,D]`` contract. + + Names describe physical structure and control semantics. Dataset field + selection, normalization, and benchmark identity belong outside presets. + + Repeating components supports multi-arm layouts; their source columns must + partition all D columns exactly once. The same layout can be used with + first- or second-order physical processors; primitive order is not selected + by this preset. See ``describe()`` for the complete source-column contract. + + Args: + name: Registered layout name describing physical structure and semantics. + input_dim: Total D in a source tensor ``(B,T,D)``; includes any grippers. + input: ``absolute`` targets or per-step ``delta`` motion increments. + The public frontend requires ``current_state(B,D)`` for absolute + inputs and omits it for delta inputs. Grippers remain commands. + components: Ordered joint/EEF component definitions with disjoint source + columns, including their optional binary gripper columns. + source_hz: Default source sampling rate, overridden by runtime ``fps``. + target_hz: Physical codec clock used by the token stream. + interpolation: Motion resampling method: ``linear``, ``pchip`` or + ``cubic``. Binary gripper commands use hold interpolation. + """ + + name: str + input_dim: int + input: str + components: tuple[ComponentSpec, ...] + source_hz: float = 15.0 + target_hz: float = 15.0 + interpolation: str = "cubic" + + def __post_init__(self) -> None: + if not isinstance(self.name, str) or _NAME.fullmatch(self.name) is None: + raise ConfigError(f"preset.name must match {_NAME.pattern!r}") + _integer(self.input_dim, "preset.input_dim", positive=True) + if self.input not in _INPUTS: + raise ConfigError("preset.input must be absolute or delta") + if not self.components: + raise ConfigError("preset.components must not be empty") + if not all( + isinstance(component, ComponentSpec) for component in self.components + ): + raise ConfigError("preset.components must contain ComponentSpec values") + names = [component.name for component in self.components] + if len(set(names)) != len(names): + raise ConfigError("preset.components contains a duplicate name") + self._validate_source_partition() + for rate, path in ( + (self.source_hz, "source_hz"), + (self.target_hz, "target_hz"), + ): + checked = _number(rate, f"preset.{path}") + if checked <= 0.0: + raise ConfigError(f"preset.{path} must be positive and finite") + if self.interpolation not in _INTERPOLATIONS: + raise ConfigError( + f"preset.interpolation must be one of {sorted(_INTERPOLATIONS)}" + ) + if self.input == "delta" and any( + component.space == "eef" and component.rotation_format != "rotation_vector" + for component in self.components + ): + raise ConfigError( + "EEF delta components require rotation_format=rotation_vector" + ) + + def _validate_source_partition(self) -> None: + validate_dimension_partition( + ( + (component.name, component.all_source_dims) + for component in self.components + ), + self.input_dim, + label="input dim", + error_type=ConfigError, + ) + + @classmethod + def from_dict(cls, payload: object, *, path: str = "preset") -> PresetConfig: + values = _mapping(payload, path) + allowed = { + "name", + "input_dim", + "input", + "source_hz", + "target_hz", + "interpolation", + "components", + } + _reject_unknown(values, allowed, path) + raw_components = _mapping( + _required(values, "components", path), f"{path}.components" + ) + components = tuple( + ComponentSpec.from_dict(name, raw, path=f"{path}.components.{name}") + for name, raw in raw_components.items() + ) + return cls( + name=str(_required(values, "name", path)), + input_dim=_integer( + _required(values, "input_dim", path), f"{path}.input_dim", positive=True + ), + input=str(_required(values, "input", path)), + components=components, + source_hz=_number(values.get("source_hz", 15.0), f"{path}.source_hz"), + target_hz=_number(values.get("target_hz", 15.0), f"{path}.target_hz"), + interpolation=str(values.get("interpolation", "cubic")), + ) + + @classmethod + def from_yaml(cls, path: str | Path) -> PresetConfig: + """Load one strict user preset without executing arbitrary code.""" + source = Path(path) + try: + payload = yaml.load( + source.read_text(encoding="utf-8"), Loader=_UniqueKeyLoader + ) + except yaml.YAMLError as error: + raise ConfigError(f"cannot parse {source}: {error}") from error + return cls.from_dict(payload, path=str(source)) + + def save_yaml(self, path: str | Path) -> Path: + """Save the complete reproducible preset and return its path.""" + target = Path(path) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text( + yaml.safe_dump(self.to_dict(), sort_keys=False), encoding="utf-8" + ) + return target + + def to_dict(self) -> dict[str, object]: + payload = { + "name": self.name, + "input_dim": self.input_dim, + "input": self.input, + "source_hz": self.source_hz, + "target_hz": self.target_hz, + "interpolation": self.interpolation, + "components": { + component.name: component.to_dict() for component in self.components + }, + } + return payload + + def summary(self) -> str: + """Return a compact action-space label for discovery UIs.""" + spaces = "+".join( + dict.fromkeys(component.space for component in self.components) + ) + return f"{self.name} (D={self.input_dim}, {self.input}, {spaces}, {self.source_hz:g} Hz input)" + + def describe(self) -> str: + """Describe every source tensor column and its physical interpretation.""" + motion = "position" if self.input == "absolute" else "per-step delta" + lines = [self.summary()] + for component in self.components: + if component.space == "joint": + assert component.joint_dims is not None + lines.append( + f" dims {_format_dims(component.joint_dims)}: {component.name} " + f"generalized-joint {motion} [rad or registered canonical unit]" + ) + else: + assert component.position_dims is not None + assert component.rotation_dims is not None + lines.append( + f" dims {_format_dims(component.position_dims)}: {component.name} " + f"EEF translation {motion} (x, y, z in listed order) [m]" + ) + rotation_motion = ( + "orientation" if self.input == "absolute" else "increment" + ) + rotation_unit = ( + "unit xyzw quaternion" + if component.rotation_format == "quaternion_xyzw" + else "rad" + ) + lines.append( + f" dims {_format_dims(component.rotation_dims)}: {component.name} EEF " + f"{rotation_motion} [{component.rotation_format}, {rotation_unit}]" + ) + if self.input == "delta": + lines.append( + " rotation_vector = axis * angle [rad], right-multiplicative " + "body-frame increment: R_next = R_previous @ Exp(rotvec)" + ) + elif component.rotation_format == "rotation_vector": + lines.append( + " rotation_vector = axis * angle [rad], not Euler angles" + ) + elif component.rotation_format == "euler_xyz": + lines.append(" euler_xyz = extrinsic x, y, z Euler angles [rad]") + if component.gripper_dim is not None: + comparison = ">=" if component.gripper_open_is_high else "<=" + lines.append( + f" dim {component.gripper_dim}: {component.name} gripper command " + f"(open if value {comparison} {component.gripper_threshold:g}, otherwise closed; " + "decoded as absolute command: 1=open, 0=closed)" + ) + else: + lines.append( + f" {component.name}: no binary gripper channel; all listed columns are motion" + ) + if self.input == "absolute": + lines.append( + " encode/decode require the same current_state (B,D), immediately before row 0" + ) + else: + lines.append( + " encode/decode: omit current_state; deltas are per-step increments, not velocities" + ) + lines.append( + f" resampling: {self.source_hz:g} Hz default -> {self.target_hz:g} Hz codec " + f"clock ({self.interpolation})" + ) + return "\n".join(lines) + + def canonicalize( + self, + values: np.ndarray, + *, + source_hz: float | None = None, + current_state: np.ndarray | None = None, + ) -> dict[str, list[ActiveTrajectory]]: + """Convert a batch into local paths for the two shared profiles. + + Args: + values (np.ndarray): ``(B,T,D)`` tensor interpreted by this preset. + source_hz: Sampling rate of ``values``. + current_state (np.ndarray | None): Optional ``(B,D)`` state before + row zero. For order-0 inputs it defines the first transition; + when omitted, each input's first row is its anchor. Order-1 + inputs are already self-contained and do not use the value. + """ + array = np.asarray(values, dtype=np.float32) + if array.ndim != 3 or len(array) < 1 or array.shape[1] < 1: + raise ValueError("preset input must have non-empty shape [B, T, D]") + if array.shape[2] != self.input_dim: + raise ValueError( + f"preset {self.name!r} expects input_dim={self.input_dim}, got {array.shape[2]}" + ) + if not np.all(np.isfinite(array)): + raise ValueError("preset input contains a non-finite value") + states: list[np.ndarray | None] + if current_state is None: + states = [None] * len(array) + else: + state_array = optional_finite_array( + current_state, shape=(len(array), self.input_dim) + ) + assert state_array is not None + states = list(state_array) + return { + component.name: [ + self._canonicalize_component( + episode, + component, + source_hz=self.source_hz if source_hz is None else source_hz, + current_state=state, + ) + for episode, state in zip(array, states, strict=True) + ] + for component in self.components + } + + def canonicalize_episode( + self, + values: np.ndarray, + *, + source_hz: float | None = None, + current_state: np.ndarray | None = None, + ) -> dict[str, ActiveTrajectory]: + """Canonicalize one variable-length episode and optional ``(D,)`` anchor.""" + episode = np.asarray(values, dtype=np.float32) + if episode.ndim != 2 or len(episode) < 1 or episode.shape[1] != self.input_dim: + raise ValueError( + f"preset {self.name!r} expects a non-empty [T, {self.input_dim}] episode" + ) + if not np.all(np.isfinite(episode)): + raise ValueError("preset input contains a non-finite value") + state = optional_finite_array(current_state, shape=(self.input_dim,)) + return { + component.name: self._canonicalize_component( + episode, + component, + source_hz=self.source_hz if source_hz is None else source_hz, + current_state=state, + ) + for component in self.components + } + + def profile_trajectories( + self, values: np.ndarray, *, current_state: np.ndarray | None = None + ) -> dict[str, list[ActiveTrajectory]]: + """Canonicalize one episode and group its components by joint/EEF profile.""" + canonical = self.canonicalize_episode(values, current_state=current_state) + output: dict[str, list[ActiveTrajectory]] = {} + for component in self.components: + output.setdefault(component.space, []).append(canonical[component.name]) + return output + + def _canonicalize_component( + self, + episode: np.ndarray, + component: ComponentSpec, + *, + source_hz: float, + current_state: np.ndarray | None, + ) -> ActiveTrajectory: + if component.space == "joint": + motion = self._joint_motion( + episode, component, source_hz=source_hz, current_state=current_state + ) + else: + motion = self._eef_motion( + episode, component, source_hz=source_hz, current_state=current_state + ) + if component.gripper_dim is None: + return ActiveTrajectory(motion, component.active_dims) + gripper = resample_uniform( + episode[:, component.gripper_dim : component.gripper_dim + 1], + source_hz=source_hz, + target_hz=self.target_hz, + kind="hold", + method=self.interpolation, + ) + binary = binarize_gripper_actions( + gripper, + threshold=component.gripper_threshold, + open_is_high=component.gripper_open_is_high, + ) + if len(binary) != len(motion): + raise RuntimeError("motion and gripper produced different target horizons") + return ActiveTrajectory( + np.concatenate((motion, binary), axis=1), component.active_dims + ) + + def _joint_motion( + self, + episode: np.ndarray, + component: ComponentSpec, + *, + source_hz: float, + current_state: np.ndarray | None, + ) -> np.ndarray: + assert component.joint_dims is not None + source = np.take(episode, component.joint_dims, axis=1) + if self.input == "delta": + source = np.cumsum(source, axis=0, dtype=np.float32) + anchor = np.zeros((1, source.shape[1]), dtype=np.float32) + else: + raw_anchor = ( + source[:1] + if current_state is None + else np.take(current_state, component.joint_dims)[None] + ) + unwrapped = np.unwrap(np.concatenate((raw_anchor, source), axis=0), axis=0) + anchor = np.asarray(unwrapped[:1], dtype=np.float32) + source = np.asarray(unwrapped[1:], dtype=np.float32) + sampled = resample_uniform( + source, + source_hz=source_hz, + target_hz=self.target_hz, + kind="linear", + method=self.interpolation, + ) + return np.asarray(sampled - anchor, dtype=np.float32) + + def _eef_motion( + self, + episode: np.ndarray, + component: ComponentSpec, + *, + source_hz: float, + current_state: np.ndarray | None, + ) -> np.ndarray: + assert ( + component.position_dims is not None and component.rotation_dims is not None + ) + position = np.take(episode, component.position_dims, axis=1) + rotation = np.take(episode, component.rotation_dims, axis=1) + if self.input == "delta": + pose = integrate_eef_deltas(np.concatenate((position, rotation), axis=1)) + pose = resample_uniform( + pose, + source_hz=source_hz, + target_hz=self.target_hz, + kind="pose_xyzw", + method=self.interpolation, + ) + return eef_pose_absolute(pose) + pose = eef_pose_from_components(position, rotation, component.rotation_format) + if current_state is None: + initial_pose = pose[0] + else: + initial_pose = eef_pose_from_components( + np.take(current_state, component.position_dims)[None], + np.take(current_state, component.rotation_dims)[None], + component.rotation_format, + )[0] + pose = resample_uniform( + pose, + source_hz=source_hz, + target_hz=self.target_hz, + kind="pose_xyzw", + method=self.interpolation, + ) + return eef_pose_local_absolute(pose, initial_pose=initial_pose) + + +def format_action_spaces( + presets: Mapping[str, PresetConfig], + *, + selected: str | None = None, + action_dim: int | None = None, + second_order_dims: Mapping[str, Sequence[int]] | None = None, +) -> str: + """Explain registered tensor layouts without inferring semantics from shape. + + Args: + presets: Registered names mapped to their tensor-layout definitions. + selected: Currently bound action-space name, if any. + action_dim: Optional source-width filter; matching layouts are candidates. + second_order_dims: Physical profile names mapped to dimensions requiring + an explicit first-order boundary in the loaded artifact. ``None`` + means the artifact's boundary requirements are unknown. + """ + if action_dim is not None: + action_dim = _integer(action_dim, "action_dim", positive=True) + if not presets: + return "No action spaces are registered." + candidates = sorted( + ( + preset + for preset in presets.values() + if action_dim is None or preset.input_dim == action_dim + ), + key=lambda preset: (preset.input_dim, preset.name), + ) + lines = [ + "ActionCodec2 — action layout guide", + f"Selected action_space: {selected or 'none (choose a layout before encoding)'}", + "", + "Input: NumPy arrays or Torch tensors, (T,D) or (B,T,D).", + " T = timesteps, D = columns per action, B = batch size; column indices start at 0.", + " Any positive T is accepted. A batch needs equal T; encode unequal-length episodes separately.", + " decode returns (B,T,D), including B=1 for a single episode; values are quantized, not lossless.", + "", + "How to choose:", + " joint = robot joint coordinates; EEF = end-effector (tool) position and orientation.", + " absolute = target coordinates; delta = per-step motion increments, not velocities.", + " Match column order, units, rotation convention, gripper commands and controller fps.", + " Shape alone cannot identify physical meaning: layouts sharing D are only candidates.", + " Consult your dataset/controller documentation if these meanings are unknown.", + " Supply physical units, not normalized [-1,1] features; convert degrees to radians.", + " Non-radian joint coordinates need an explicit mapping to the registered canonical units.", + " Pass the actual sampling rate as fps=...; defaults and codec rates appear below.", + " Resampling can change decoded T; token coverage determines the returned length.", + "", + f"Registered configurations: {len(candidates)} shown / {len(presets)} total" + + ( + "" + if action_dim is None + else f" (D={action_dim} candidates, not automatic selection)" + ), + ] + if not candidates: + widths = ", ".join( + map(str, sorted({preset.input_dim for preset in presets.values()})) + ) + lines.extend( + [ + f" No registered action space matches D={action_dim}. Available D: {widths}.", + " Confirm your column meanings and register a custom layout; do not pad/truncate to force a match.", + ] + ) + for preset in candidates: + name = preset.name + marker = " [selected]" if name == selected else "" + description = preset.describe().splitlines() + lines.extend(["", f"- {description[0]}{marker}"]) + lines.extend(f" {line.strip()}" for line in description[1:]) + if second_order_dims is not None: + boundaries = [ + f"{component.name!r}: (B,{len(component.active_dims)})" + for component in preset.components + if set(component.active_dims).intersection( + second_order_dims.get(component.space, ()) + ) + ] + lines.append( + " previous_first_order: " + + ( + "required in encode/decode, {" + ", ".join(boundaries) + "}; " + f"reconstructed increments at {preset.target_hz:g} Hz, B=1 for a single episode" + if boundaries + else "not required by this artifact for this layout" + ) + ) + lines.extend( + [ + "", + "Next steps:", + " Load: ActionCodec2.from_pretrained(path, action_space='NAME')", + " Switch: codec = codec.for_action_space('NAME')", + " Filter: codec.print_action_spaces(action_dim=D)", + " Config: codec.action_spaces['NAME'].to_dict() # exact settings for a listed layout", + " Custom: codec.register_action_space(preset) # explicit column definitions, then switch", + " Custom layouts can reorder columns; each joint component supports 1..7 coordinates.", + " EEF absolute rotations support rotation_vector, euler_xyz or quaternion_xyzw (unit x,y,z,w).", + " EEF delta rotations require rotation_vector with the convention shown above.", + " For absolute layouts, current_state is (B,D), even when actions are (T,D).", + ] + ) + return "\n".join(lines) + + +def resample_canonical_trajectory( + trajectory: ActiveTrajectory, + component: ComponentSpec, + *, + source_hz: float, + target_hz: float, + target_length: int | None = None, + method: str = "linear", +) -> ActiveTrajectory: + """Rate-convert one decoded canonical component without changing slots.""" + values = np.asarray(trajectory.values, dtype=np.float32) + motion_width = len(component.joint_dims or ()) if component.space == "joint" else 6 + motion = values[:, :motion_width] + if component.space == "eef": + pose = eef_pose_from_components( + motion[:, :3], motion[:, 3:6], "rotation_vector" + ) + sampled_pose = resample_uniform( + pose, + source_hz=source_hz, + target_hz=target_hz, + kind="pose_xyzw", + target_length=target_length, + method=method, + ) + sampled_motion = eef_pose_absolute(sampled_pose) + else: + sampled_motion = resample_uniform( + motion, + source_hz=source_hz, + target_hz=target_hz, + kind="linear", + target_length=target_length, + method=method, + ) + if component.gripper_dim is None: + return ActiveTrajectory(sampled_motion, trajectory.active_dims) + gripper = resample_uniform( + values[:, motion_width : motion_width + 1], + source_hz=source_hz, + target_hz=target_hz, + kind="hold", + target_length=target_length, + method=method, + ) + return ActiveTrajectory( + np.concatenate((sampled_motion, gripper), axis=1), trajectory.active_dims + ) + + +def externalize_canonical_trajectory( + trajectory: ActiveTrajectory, + component: ComponentSpec, + *, + current_state: np.ndarray | None, +) -> np.ndarray: + """Return deltas, or absolute values when an external anchor is supplied. + + The input is always a decoded local position/pose path. Gripper remains an + absolute binary command in both output modes. + """ + values = np.asarray(trajectory.values, dtype=np.float32) + motion_width = len(component.joint_dims or ()) if component.space == "joint" else 6 + motion = values[:, :motion_width] + if component.space == "joint": + if current_state is None: + previous = np.concatenate((np.zeros_like(motion[:1]), motion[:-1]), axis=0) + output_motion = motion - previous + else: + assert component.joint_dims is not None + output_motion = motion + np.take(current_state, component.joint_dims)[None] + else: + local_pose = eef_pose_from_components( + motion[:, :3], motion[:, 3:6], "rotation_vector" + ) + if current_state is None: + output_motion = eef_pose_deltas(local_pose) + else: + assert ( + component.position_dims is not None + and component.rotation_dims is not None + ) + initial_pose = eef_pose_from_components( + np.take(current_state, component.position_dims)[None], + np.take(current_state, component.rotation_dims)[None], + component.rotation_format, + )[0] + output_motion = eef_pose_components( + eef_pose_from_local_absolute(local_pose, initial_pose=initial_pose), + component.rotation_format, + ) + if component.gripper_dim is None: + return np.asarray(output_motion, dtype=np.float32) + gripper = values[:, motion_width : motion_width + 1] + if not component.gripper_open_is_high: + gripper = 1.0 - gripper + return np.asarray( + np.concatenate((output_motion, gripper), axis=1), + dtype=np.float32, + ) diff --git a/runtime/routing/routing.py b/runtime/routing/routing.py new file mode 100644 index 0000000000000000000000000000000000000000..b6c4dbb8b73d738cc2d53ebe643a598536bb0f6c --- /dev/null +++ b/runtime/routing/routing.py @@ -0,0 +1,809 @@ +"""Required-preset routing over the two shared ActionCodec2 processors. + +Each profile owns a disjoint token namespace and is named for a physical +quantity, never a robot. Presets describe tensor slots and preprocessing but +cannot alter already-fitted primitives, BPE rules, or token meanings. +""" + +from __future__ import annotations + +import re +from collections.abc import Mapping, Sequence +from dataclasses import asdict, dataclass +from pathlib import Path +from types import MappingProxyType +from typing import TYPE_CHECKING + +import numpy as np +import yaml + +from ..common._validation import optional_finite_array +from .cache import BoundedCache +from ..physical.configuration import ( + ConfigError, + _integer, + _mapping, + _reject_unknown, + _required, + _UniqueKeyLoader, +) +from ..integration.hf import ( + ROUTER_ALLOW_PATTERNS, + ROUTER_CONFIG_NAME, + ROUTER_MODEL_TYPE, + resolve_pretrained_directory, + save_huggingface_metadata, + save_model_card, +) +from ..physical.physical_processor import ActionCodec2Processor +from .presets import ( + ComponentSpec, + PresetConfig, + externalize_canonical_trajectory, + format_action_spaces, + resample_canonical_trajectory, +) +from ..physical.semantic import resampled_horizon +from ..physical.trajectory import ActiveTrajectory + +if TYPE_CHECKING: + from .processing_actioncodec2 import ActionCodec2 + +# v3 freezes stateless decode semantics: no current_state -> deltas; +# current_state -> absolute motion. v2 returned local position paths. +ROUTER_CONFIG_VERSION = 3 +_NAME = re.compile(r"[a-z][a-z0-9_]*\Z") +_PROFILES = {"eef", "joint"} +_PROFILE_WIDTHS = {"eef": 7, "joint": 8} + + +def _current_state_rows( + current_state: np.ndarray | None, *, count: int, input_dim: int +) -> list[np.ndarray | None]: + if current_state is None: + return [None] * count + values = optional_finite_array(current_state, shape=(count, input_dim)) + assert values is not None + return list(values) + + +@dataclass(frozen=True) +class ProfileSpec: + """Token namespace metadata for one shared physical processor.""" + + name: str + token_offset: int + token_budget: int + artifact_fingerprint: str + + def __post_init__(self) -> None: + if _NAME.fullmatch(self.name) is None: + raise ConfigError(f"profile.name must match {_NAME.pattern!r}") + _integer(self.token_offset, "profile.token_offset") + _integer(self.token_budget, "profile.token_budget", positive=True) + if re.fullmatch(r"[0-9a-f]{64}", self.artifact_fingerprint) is None: + raise ConfigError( + "profile.artifact_fingerprint must be a SHA-256 hex digest" + ) + + @classmethod + def from_dict(cls, payload: object, *, path: str) -> ProfileSpec: + values = _mapping(payload, path) + allowed = {"name", "token_offset", "token_budget", "artifact_fingerprint"} + _reject_unknown(values, allowed, path) + return cls( + name=str(_required(values, "name", path)), + token_offset=_integer( + _required(values, "token_offset", path), f"{path}.token_offset" + ), + token_budget=_integer( + _required(values, "token_budget", path), + f"{path}.token_budget", + positive=True, + ), + artifact_fingerprint=str(_required(values, "artifact_fingerprint", path)), + ) + + def to_dict(self) -> dict[str, object]: + return asdict(self) + + +@dataclass(frozen=True) +class _ProfileRuntime: + spec: ProfileSpec + processor: ActionCodec2Processor + + +@dataclass(frozen=True) +class _Segment: + component: ComponentSpec + profile: _ProfileRuntime + + +class RoutedActionCodec2Grammar: + """Constrained grammar over a preset's ordered, namespaced segments.""" + + def __init__( + self, router: ActionCodec2Router, *, preset: str, horizon: int + ) -> None: + horizon = _integer(horizon, "horizon", positive=True) + self.segments = router._segments(router.preset(preset)) + self.vocab_size = router.vocab_size + self._grammars = tuple( + segment.profile.processor.grammar( + horizon, active_dims=segment.component.active_dims + ) + for segment in self.segments + ) + initial = self._normalize_state(0, self._grammars[0]._initial_state()) + self._initial_state = initial + self._state_cache = BoundedCache(256) + self._state_cache[()] = initial + + def _normalize_state( + self, segment_index: int, state: tuple[int, ...] + ) -> tuple[int, tuple[int, ...]]: + while segment_index < len(self.segments) and self._grammars[ + segment_index + ]._state_is_complete(state): + segment_index += 1 + if segment_index < len(self.segments): + state = self._grammars[segment_index]._initial_state() + return segment_index, state + + def _advance_routed_state( + self, + routed_state: tuple[int, tuple[int, ...]], + token: int, + ) -> tuple[int, tuple[int, ...]] | None: + segment_index, state = routed_state + if segment_index == len(self.segments): + return None + segment = self.segments[segment_index] + local_token = token - segment.profile.spec.token_offset + if not 0 <= local_token < segment.profile.processor.vocab_size: + return None + next_state = self._grammars[segment_index]._advance(state, local_token) + if next_state is None: + return None + return self._normalize_state(segment_index, next_state) + + def _consume(self, tokens: Sequence[int]) -> tuple[int, tuple[int, ...]] | None: + checked: list[int] = [] + for raw_token in tokens: + if isinstance(raw_token, bool) or not isinstance( + raw_token, (int, np.integer) + ): + return None + checked.append(int(raw_token)) + key = tuple(checked) + if not key: + return self._initial_state + missing = object() + cached = self._state_cache.get(key, missing) + if cached is not missing: + return cached + parent = key[:-1] + state = self._state_cache.get(parent, missing) + if state is missing: + state = self._initial_state + for token in parent: + if state is None: + break + state = self._advance_routed_state(state, token) + self._state_cache[parent] = state + if state is not None: + state = self._advance_routed_state(state, key[-1]) + self._state_cache[key] = state + return state + + def next_token_mask(self, tokens: Sequence[int]) -> np.ndarray: + """Return a global namespace mask for the next preset token.""" + state = self._consume(tokens) + if state is None: + raise ValueError("token prefix is not legal for this routed preset") + segment_index, local_state = state + output = np.zeros(self.vocab_size, dtype=bool) + if segment_index == len(self.segments): + return output + segment = self.segments[segment_index] + mask = self._grammars[segment_index]._next_token_mask_from_state(local_state) + start = segment.profile.spec.token_offset + output[start : start + len(mask)] = mask + return output + + def is_complete(self, tokens: Sequence[int]) -> bool: + """Whether the prefix completes every segment exactly once.""" + state = self._consume(tokens) + return state is not None and state[0] == len(self.segments) + + +class ActionCodec2Router: + """Route heterogeneous dataset states into physical token namespaces. + + The router is fitted only from already-built physical processors. Presets + cannot alter their primitives, merge rules, or token meanings. + + Args: + profiles: Ordered namespace metadata and fitted processor pairs for + the shared joint and end-effector physical profiles. + presets: Action-space definitions mapping source tensor columns into + the physical profiles. + """ + + model_type = ROUTER_MODEL_TYPE + + def __init__( + self, + profiles: Sequence[tuple[ProfileSpec, ActionCodec2Processor]], + presets: Sequence[PresetConfig], + ) -> None: + runtimes: dict[str, _ProfileRuntime] = {} + expected_offset = 0 + for spec, processor in profiles: + if spec.name in runtimes: + raise ConfigError(f"duplicate physical profile {spec.name!r}") + if spec.token_offset != expected_offset: + raise ConfigError( + "profile token namespaces must be contiguous and ordered" + ) + if processor.vocab_size > spec.token_budget: + raise ConfigError( + f"profile {spec.name!r} uses {processor.vocab_size} tokens, " + f"exceeding budget {spec.token_budget}" + ) + if processor.artifact_fingerprint != spec.artifact_fingerprint: + raise ConfigError( + f"profile {spec.name!r} artifact fingerprint does not match its processor" + ) + if processor.config.set_bpe.allow_cross_component_merge: + raise ConfigError( + "allow_cross_component_merge=true is not supported by the v1 router; " + "a component-aware Set-BPE serializer and grammar are required" + ) + runtimes[spec.name] = _ProfileRuntime(spec, processor) + expected_offset += spec.token_budget + if not runtimes or not set(runtimes) <= _PROFILES: + raise ConfigError( + "router requires joint and/or eef physical profiles" + ) + for name, runtime in runtimes.items(): + if runtime.processor.action_dim != _PROFILE_WIDTHS[name]: + raise ConfigError( + f"profile {name!r} must have width {_PROFILE_WIDTHS[name]}, " + f"got {runtime.processor.action_dim}" + ) + self.profiles = MappingProxyType(runtimes) + self._source_directory = Path(__file__).resolve().parent.parent + self._grammar_cache = BoundedCache(8) + self._action_spaces: dict[str, PresetConfig] = {} + for preset in presets: + self.register_action_space(preset) + self.vocab_size = expected_offset + + @property + def action_spaces(self) -> Mapping[str, PresetConfig]: + """Read-only registered built-in and user-created action spaces.""" + return MappingProxyType(self._action_spaces) + + def __str__(self) -> str: + return self.describe_action_spaces() + + def describe_action_spaces( + self, *, selected: str | None = None, action_dim: int | None = None + ) -> str: + """Explain all layouts, or source-width candidates, using this artifact's state requirements.""" + second_order_dims = { + name: tuple( + dimension + for component in runtime.processor.config.physical.components.values() + if component.primitive_order == 2 + for dimension in component.dims + ) + for name, runtime in self.profiles.items() + } + return format_action_spaces( + self.action_spaces, + selected=selected, + action_dim=action_dim, + second_order_dims=second_order_dims, + ) + + @classmethod + def from_processors( + cls, + processors: Mapping[str, ActionCodec2Processor], + *, + token_budgets: Mapping[str, int], + presets: Sequence[PresetConfig], + ) -> ActionCodec2Router: + """Assign stable contiguous namespaces in canonical joint/eef order.""" + if set(processors) != set(token_budgets): + raise ConfigError("processors and token_budgets must share keys") + if set(processors) != _PROFILES: + raise ConfigError( + "processors must contain exactly the joint and eef profiles" + ) + offset = 0 + entries = [] + for name in ("joint", "eef"): + processor = processors[name] + budget = _integer( + token_budgets[name], f"token_budgets.{name}", positive=True + ) + spec = ProfileSpec( + name, + offset, + budget, + processor.artifact_fingerprint, + ) + entries.append((spec, processor)) + offset += budget + return cls(entries, presets) + + def preset(self, name: str) -> PresetConfig: + """Resolve one preset or fail with the available names.""" + try: + return self._action_spaces[name] + except KeyError as error: + raise KeyError( + f"unknown preset {name!r}; available: {sorted(self._action_spaces)}" + ) from error + + def register_action_space( + self, + preset: PresetConfig | Mapping[str, object] | str | Path, + ) -> PresetConfig: + """Register a validated action-space config, mapping, or YAML file. + + Registration changes only tensor interpretation. It cannot alter either + fitted physical vocabulary or any token ID. + """ + if isinstance(preset, Mapping): + preset = PresetConfig.from_dict(preset) + elif isinstance(preset, (str, Path)): + preset = PresetConfig.from_yaml(preset) + if not isinstance(preset, PresetConfig): + raise TypeError( + "register_action_space expects PresetConfig, a mapping, or a preset YAML path" + ) + existing = self._action_spaces.get(preset.name) + if existing is not None and existing != preset: + raise ConfigError( + f"action-space name {preset.name!r} is already registered differently" + ) + self._action_spaces[preset.name] = preset + return preset + + def _segments(self, preset: PresetConfig) -> tuple[_Segment, ...]: + return tuple( + _Segment(component, self.profiles[component.space]) + for component in preset.components + ) + + def encode_batch( + self, + inputs: np.ndarray, + *, + preset: str, + fps: float | None = None, + current_state: np.ndarray | None = None, + previous_first_order: Mapping[str, np.ndarray] | None = None, + ) -> list[list[int]]: + """Encode one regular-rate ``[B,T,D]`` tensor at ``fps``. + + Preset preprocessing and runtime encoding use the same canonical rate + adapter. The shared physical processors therefore only receive the + preset's target clock. For order-0 presets, + ``current_state[B,D]`` is the state immediately before row zero. The + low-level router permits omitting it and then anchors at row zero; the + public frontend requires it for absolute action spaces so the first + transition cannot be silently discarded. + ``previous_first_order`` maps component names to ``[B,D_active]`` + reconstructed increments in canonical units at the preset's target + clock; active second-order components require an entry. + """ + preset_config = self.preset(preset) + self._validate_boundary_keys(previous_first_order, preset_config) + canonical = preset_config.canonicalize( + inputs, source_hz=fps, current_state=current_state + ) + batch = len(next(iter(canonical.values()))) + rows: list[list[int]] = [[] for _ in range(batch)] + for segment in self._segments(preset_config): + values = canonical[segment.component.name] + streams = segment.profile.processor.encode_batch( + np.stack([trajectory.values for trajectory in values]), + active_dims=segment.component.active_dims, + previous_first_order=None + if previous_first_order is None + else previous_first_order.get(segment.component.name), + ) + offset = segment.profile.spec.token_offset + for row, stream in zip(rows, streams, strict=True): + row.extend(offset + token for token in stream) + return rows + + def route_trajectories( + self, + inputs: np.ndarray, + *, + preset: str, + current_state: np.ndarray | None = None, + fps: float | None = None, + previous_first_order: Mapping[str, np.ndarray] | None = None, + ) -> dict[str, list[ActiveTrajectory]]: + """Expose canonical complete trajectories grouped by physical profile. + + This is the fitting boundary: callers pool the returned trajectories + from many presets by profile name, balance datasets, then call + :meth:`ProfileTrainer.fit` once per profile. + ``previous_first_order`` uses the same component-to-``[B,D_active]`` + canonical boundary mapping as :meth:`encode_batch`. + """ + preset_config = self.preset(preset) + self._validate_boundary_keys(previous_first_order, preset_config) + canonical = preset_config.canonicalize( + inputs, current_state=current_state, source_hz=fps + ) + output: dict[str, list[ActiveTrajectory]] = {} + for component in preset_config.components: + trajectories = canonical[component.name] + processor = self.profiles[component.space].processor + previous = ( + None + if previous_first_order is None + else previous_first_order.get(component.name) + ) + dense = processor._dense_previous_first_order( + previous, component.active_dims, len(trajectories) + ) + rows = ( + [None] * len(trajectories) + if dense is None + else np.take(dense, component.active_dims, axis=1) + ) + output.setdefault(component.space, []).extend( + ActiveTrajectory(t.values, t.active_dims, previous_first_order=p) + for t, p in zip(trajectories, rows, strict=True) + ) + return output + + def _validate_boundary_keys( + self, + previous_first_order: Mapping[str, np.ndarray] | None, + preset: PresetConfig, + ) -> None: + if previous_first_order is None: + previous_first_order = {} + if not isinstance(previous_first_order, Mapping): + raise ValueError( + "previous_first_order must map component names to canonical [B,D_active] boundaries" + ) + unknown = set(previous_first_order) - {c.name for c in preset.components} + if unknown: + raise ValueError( + f"unknown previous_first_order component(s): {sorted(unknown)}" + ) + required = { + c.name: f"(B, {len(c.active_dims)})" + for c in preset.components + if any( + p.primitive_order == 2 and set(c.active_dims).intersection(p.dims) + for p in self.profiles[ + c.space + ].processor.config.physical.components.values() + ) + } + missing = set(required) - previous_first_order.keys() + if missing: + raise ValueError( + "second-order tokenization requires previous_first_order with arrays " + f"{required}; missing {sorted(missing)}. Use canonical {preset.target_hz:g} Hz " + "reconstructed increments. codec.first_order_template(batch_size) explicitly " + "creates zeros; use zeros only if the preceding increment is known to be zero." + ) + + def final_first_order( + self, + streams: Sequence[Sequence[int]], + *, + preset: str, + previous_first_order: Mapping[str, np.ndarray], + executed_steps: int | None = None, + ) -> dict[str, np.ndarray]: + """Return component ``[B,D_active]`` boundaries at the canonical clock. + + ``executed_steps`` counts codec steps, before output resampling. EEF + motion uses local Cartesian displacements and rotation vectors in the + profile's multiplication convention. Gripper slots carry no state. + """ + preset_config = self.preset(preset) + if previous_first_order is None: + raise ValueError("final_first_order requires previous_first_order") + self._validate_boundary_keys(previous_first_order, preset_config) + segments = self._segments(preset_config) + steps = self._batch_steps(streams, segments) + split = [self._split_stream(stream, segments, steps) for stream in streams] + output = {} + for index, segment in enumerate(segments): + processor = segment.profile.processor + if not processor.requires_previous_first_order: + continue + output[segment.component.name] = processor.final_first_order( + [row[index] for row in split], + previous_first_order=previous_first_order.get(segment.component.name), + active_dims=segment.component.active_dims, + executed_steps=executed_steps, + ) + return output + + def _split_stream( + self, tokens: Sequence[int], segments: Sequence[_Segment], steps: int + ) -> list[list[int]]: + cursor = 0 + output = [] + for segment in segments: + grammar = segment.profile.processor.grammar( + steps, active_dims=segment.component.active_dims + ) + state = grammar._initial_state() + local: list[int] = [] + while cursor < len(tokens) and not grammar._state_is_complete(state): + raw_token = tokens[cursor] + if isinstance(raw_token, bool) or not isinstance( + raw_token, (int, np.integer) + ): + raise ValueError(f"tokens[{cursor}] must be an integer") + token = int(raw_token) - segment.profile.spec.token_offset + if not 0 <= token < segment.profile.processor.vocab_size: + raise ValueError( + "token belongs to the wrong routed profile namespace" + ) + next_state = grammar._advance(state, token) + if next_state is None: + raise ValueError( + f"Action token stream has invalid coverage for component {segment.component.name!r}. Constrain generation with codec.grammar(H, fps=...).next_token_mask(prefix)." + ) + state = next_state + local.append(token) + cursor += 1 + if not grammar._state_is_complete(state): + raise ValueError( + f"Incomplete action token stream for component {segment.component.name!r}. Check codec.grammar(H, fps=...).is_complete(tokens); discard or regenerate this row." + ) + output.append(local) + if cursor != len(tokens): + raise ValueError( + "Action token stream has extra tokens after complete coverage; remove model prompt/EOS/padding before codec.decode()." + ) + return output + + def _token_cell_count(self, raw_token: object, *, index: int) -> int: + """Return the number of primitive cells represented by one routed token.""" + if isinstance(raw_token, bool) or not isinstance(raw_token, (int, np.integer)): + raise ValueError(f"tokens[{index}] must be an integer") + token = int(raw_token) + for runtime in self.profiles.values(): + local = token - runtime.spec.token_offset + if 0 <= local < runtime.processor.vocab_size: + return int(runtime.processor.serializer.token_cell_counts[local]) + raise ValueError(f"tokens[{index}] is outside the routed profile vocabularies") + + def _decoded_steps( + self, tokens: Sequence[int], segments: Sequence[_Segment] + ) -> int: + """Derive sequence length from the primitive-cell coverage in a routed stream.""" + if not tokens: + raise ValueError("token stream must not be empty") + cells_per_step = sum(len(segment.component.active_dims) for segment in segments) + represented_cells = sum( + self._token_cell_count(token, index=index) + for index, token in enumerate(tokens) + ) + steps, remainder = divmod(represented_cells, cells_per_step) + if steps < 1 or remainder: + raise ValueError( + "Incomplete action token stream: " + f"{cells_per_step - remainder} more coordinate values are needed to finish a step. " + "Constrain generation with codec.grammar(H, fps=...).next_token_mask(prefix) " + "and check is_complete(tokens) for the intended H before decoding. " + "Use codec.batch_decode(..., errors='return') to isolate invalid rows; " + "do not pad truncated tokens with zeros." + ) + return steps + + def _batch_steps( + self, + streams: Sequence[Sequence[int]], + segments: Sequence[_Segment], + ) -> int: + """Derive the common decoded length of a dense output batch.""" + if not streams: + raise ValueError("streams must not be empty") + lengths = [self._decoded_steps(stream, segments) for stream in streams] + if len(set(lengths)) != 1: + raise ValueError( + f"all token streams must decode to the same length, got {lengths}; use codec.batch_decode(...) for independent variable-length outputs" + ) + return lengths[0] + + def decode_batch( + self, + streams: Sequence[Sequence[int]], + *, + preset: str, + fps: float | None = None, + current_state: np.ndarray | None = None, + previous_first_order: Mapping[str, np.ndarray] | None = None, + ) -> dict[str, np.ndarray]: + """Decode rows to deltas, or absolute components with ``current_state``. + + Args: + streams: One namespaced token stream per batch item. + preset: Registered tensor/component interpretation. + fps: Requested external output rate; defaults to the preset's + ``source_hz``. The output row count is the canonical rate + conversion of the token-determined codec-rate path; no original + pre-downsampling row count is inferred. + current_state: Optional ``(B,D)`` preset-layout state immediately + before row zero. When supplied, decoded local paths are attached + to it and returned as order-0 values. When omitted, the return + value is per-step order-1 motion. Gripper is always absolute. + previous_first_order: Component names mapped to ``(B,D_active)`` + canonical reconstructed increments at the target clock, using + the same boundary as encode. Required for second-order motion. + """ + preset_config = self.preset(preset) + self._validate_boundary_keys(previous_first_order, preset_config) + segments = self._segments(preset_config) + internal_steps = self._batch_steps(streams, segments) + states = _current_state_rows( + current_state, count=len(streams), input_dim=preset_config.input_dim + ) + output_fps = preset_config.source_hz if fps is None else fps + output_steps = resampled_horizon( + internal_steps, + source_hz=preset_config.target_hz, + target_hz=output_fps, + ) + output: dict[str, list[np.ndarray]] = { + component.name: [] for component in preset_config.components + } + split_streams = [ + self._split_stream(stream, segments, internal_steps) for stream in streams + ] + for segment_index, segment in enumerate(segments): + decoded_batch = segment.profile.processor.decode_batch( + [split[segment_index] for split in split_streams], + active_dims=segment.component.active_dims, + previous_first_order=None + if previous_first_order is None + else previous_first_order.get(segment.component.name), + ) + for decoded, state in zip(decoded_batch, states, strict=True): + restored = resample_canonical_trajectory( + ActiveTrajectory(decoded, segment.component.active_dims), + segment.component, + source_hz=preset_config.target_hz, + target_hz=output_fps, + target_length=output_steps, + method=preset_config.interpolation, + ) + output[segment.component.name].append( + externalize_canonical_trajectory( + restored, segment.component, current_state=state + ) + ) + return {name: np.stack(values) for name, values in output.items()} + + def grammar(self, *, preset: str, horizon: int) -> RoutedActionCodec2Grammar: + """Build the global constrained grammar for one preset and target horizon.""" + horizon = _integer(horizon, "horizon", positive=True) + self.preset(preset) + key = (preset, horizon) + grammar = self._grammar_cache.get(key) + if grammar is None: + grammar = RoutedActionCodec2Grammar(self, preset=preset, horizon=horizon) + self._grammar_cache[key] = grammar + return grammar + + @property + def new_action_tokens(self) -> list[str]: + """Stable physical token names, including reserved profile budget slots.""" + return [ + f"" + for name, runtime in self.profiles.items() + for index in range(runtime.spec.token_budget) + ] + + def save_pretrained( + self, directory: str | Path, *, include_code: bool = True + ) -> list[str]: + """Save one manifest plus isolated physical processor subdirectories.""" + root = Path(directory) + files = [] + for runtime in self.profiles.values(): + files.extend( + runtime.processor.save_pretrained(root / "profiles" / runtime.spec.name) + ) + manifest = { + "model_type": self.model_type, + "config_version": ROUTER_CONFIG_VERSION, + "profiles": [runtime.spec.to_dict() for runtime in self.profiles.values()], + "presets": [preset.to_dict() for preset in self.action_spaces.values()], + } + path = root / ROUTER_CONFIG_NAME + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(yaml.safe_dump(manifest, sort_keys=False), encoding="utf-8") + metadata = save_huggingface_metadata( + root, + model_type=self.model_type, + processor_class=type(self).__name__, + **( + {"auto_map": {"AutoProcessor": "processing_actioncodec2.ActionCodec2"}} + if include_code + else {} + ), + ) + card = save_model_card(root, str(self)) + if include_code: + from ..integration.exporting import ArtifactExporter + + files.extend( + ArtifactExporter(root, source_directory=self._source_directory).export() + ) + return [str(path), str(card), *map(str, metadata), *files] + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str | Path, + **kwargs: object, + ) -> ActionCodec2Router | ActionCodec2: + """Load a router, or bind its public processor when given ``action_space``. + + This also lets ``AutoProcessor`` bind an existing router artifact in + one call. Without ``action_space``, the result remains a router. + """ + if "action_space" in kwargs: + from .processing_actioncodec2 import ActionCodec2 + + return ActionCodec2.from_pretrained(pretrained_model_name_or_path, **kwargs) + root = resolve_pretrained_directory( + pretrained_model_name_or_path, + allow_patterns=ROUTER_ALLOW_PATTERNS, + **kwargs, + ) + path = root / ROUTER_CONFIG_NAME + payload = _mapping( + yaml.load(path.read_text(encoding="utf-8"), Loader=_UniqueKeyLoader), + "router", + ) + _reject_unknown( + payload, {"model_type", "config_version", "profiles", "presets"}, "router" + ) + if payload.get("model_type") != ROUTER_MODEL_TYPE: + raise ConfigError(f"router.model_type must be {ROUTER_MODEL_TYPE!r}") + if payload.get("config_version") != ROUTER_CONFIG_VERSION: + raise ConfigError( + f"unsupported router config_version {payload.get('config_version')!r}" + ) + raw_profiles = payload.get("profiles") + raw_presets = payload.get("presets") + if not isinstance(raw_profiles, list) or not isinstance(raw_presets, list): + raise ConfigError("router.profiles and router.presets must be lists") + profiles = [] + for index, raw in enumerate(raw_profiles): + spec = ProfileSpec.from_dict(raw, path=f"router.profiles[{index}]") + processor = ActionCodec2Processor.from_pretrained( + root / "profiles" / spec.name + ) + profiles.append((spec, processor)) + presets = [ + PresetConfig.from_dict(raw, path=f"router.presets[{index}]") + for index, raw in enumerate(raw_presets) + ] + router = cls(profiles, presets) + router._source_directory = root + return router diff --git a/runtime/tokenization/__init__.py b/runtime/tokenization/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5a9900849a39d438961bf79d9755bfa2d434e072 --- /dev/null +++ b/runtime/tokenization/__init__.py @@ -0,0 +1 @@ +"""Internal tokenization components for ActionCodec2.""" diff --git a/runtime/tokenization/bpe.py b/runtime/tokenization/bpe.py new file mode 100644 index 0000000000000000000000000000000000000000..0245d14b500d84144abf64e5f28152cdea96a0b1 --- /dev/null +++ b/runtime/tokenization/bpe.py @@ -0,0 +1,439 @@ +"""Strict ``bpe.json`` I/O for both tokenization schemes. + +The serialization layer owns only the artifact schema. Set-BPE spatial and temporal +merge semantics stay in :mod:`actioncodec2.setbpe` and plain-BPE merge +semantics in :mod:`actioncodec2.tokenization.bpe_vocabulary`, so training and deployment +cannot drift through two independent implementations of the same algorithm. + +``config.set_bpe.scheme`` selects which serializer an artifact carries: +``"set_bpe"`` (the production default) or ``"bpe"``, the FAST-style baseline +that flattens the quantized cell matrix row-major before merging. +""" + +from __future__ import annotations + +import json +from collections.abc import Mapping, Sequence +from numbers import Integral +from pathlib import Path + +import numpy as np + +from ..common._validation import active_dimensions +from .bpe_vocabulary import BPEVocabulary +from ..physical.configuration import ActionCodec2Config, ConfigError +from .setbpe_spatial import SpatialVocabulary +from .setbpe_temporal import TemporalVocabulary + +BPE_NAME = "bpe.json" + + +def _integer(value: object, path: str, *, positive: bool = False) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise ConfigError(f"{path} must be an integer") + minimum = 1 if positive else 0 + if value < minimum: + qualifier = "positive" if positive else "non-negative" + raise ConfigError(f"{path} must be {qualifier}") + return int(value) + + +def _rules(payload: object, width: int, path: str) -> tuple[tuple[int, ...], ...]: + """Parse an ordered merge list without accepting lossy coercions.""" + if not isinstance(payload, Sequence) or isinstance( + payload, (str, bytes, bytearray) + ): + raise ConfigError(f"{path} must be a list") + output = [] + for rank, raw_rule in enumerate(payload): + rule_path = f"{path}[{rank}]" + if not isinstance(raw_rule, Sequence) or isinstance( + raw_rule, (str, bytes, bytearray) + ): + raise ConfigError(f"{rule_path} must be a list of {width} integers") + if len(raw_rule) != width: + raise ConfigError(f"{rule_path} must contain exactly {width} integers") + output.append( + tuple( + _integer(value, f"{rule_path}[{index}]") + for index, value in enumerate(raw_rule) + ) + ) + return tuple(output) + + +def _unique_json_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + output: dict[str, object] = {} + for key, value in pairs: + if key in output: + raise ConfigError(f"{BPE_NAME} contains duplicate field {key!r}") + output[key] = value + return output + + +def _validate_rule_references( + *, + atom_count: int, + lookahead: int, + spatial_merges: tuple[tuple[int, int], ...], + temporal_merges: tuple[tuple[int, int, int], ...], +) -> None: + """Give artifact-level errors before constructing vocabulary metadata.""" + for rank, (left, right) in enumerate(spatial_merges): + child = atom_count + rank + if left >= child or right >= child: + raise ConfigError( + f"spatial_merges[{rank}] must reference already-created tokens" + ) + spatial_size = atom_count + len(spatial_merges) + for rank, (left, right, offset) in enumerate(temporal_merges): + child = spatial_size + rank + if left >= child or right >= child: + raise ConfigError( + f"temporal_merges[{rank}] must reference already-created tokens" + ) + if not 1 <= offset < lookahead: + raise ConfigError( + f"temporal_merges[{rank}] offset must be positive and smaller than lookahead" + ) + + +def _read_payload(directory: str | Path) -> dict[str, object]: + """Parse ``bpe.json`` from an artifact directory, rejecting duplicate keys.""" + path = Path(directory) / BPE_NAME + try: + return json.loads( + path.read_text(encoding="utf-8"), object_pairs_hook=_unique_json_object + ) + except OSError: + raise + except json.JSONDecodeError as error: + raise ConfigError(f"cannot parse {path}: {error}") from error + + +def _dense_cells( + cells: np.ndarray, + bins: Sequence[int], + active_dims: Sequence[int] | None, +) -> tuple[np.ndarray, tuple[int, ...] | None]: + """Validate one trajectory and expand it to dense ``[T,D_max]`` cells. + + Compact ``[T,D_active]`` input may provide explicit global ``active_dims``. + Without it, compact input uses prefix dimensions. Dense ``[T,D_max]`` input + may mark absent slots with ``-1``. + """ + raw = np.asarray(cells) + if raw.dtype.kind not in {"i", "u"}: + raise ValueError("cells must contain integer indices") + values = np.asarray(raw, dtype=np.int64) + if values.ndim != 2 or len(values) < 1 or values.shape[1] < 1: + raise ValueError("cells must have non-empty shape [T, D_active]") + if active_dims is None: + if values.shape[1] > len(bins): + raise ValueError(f"cells width exceeds codebook dimension {len(bins)}") + if values.shape[1] < len(bins): + dimensions = tuple(range(values.shape[1])) + dense = np.full((len(values), len(bins)), -1, dtype=np.int64) + dense[:, : values.shape[1]] = values + values = dense + else: + dimensions = None + else: + dimensions = active_dimensions(active_dims, len(bins)) + if len(dimensions) != values.shape[1]: + raise ValueError( + "active_dims must have one entry per compact cell coordinate" + ) + if np.any(values < 0): + raise ValueError("compact cells must not contain -1") + dense = np.full((len(values), len(bins)), -1, dtype=np.int64) + dense[:, dimensions] = values + values = dense + counts = np.asarray(tuple(int(value) for value in bins), dtype=np.int64) + if np.any(values < -1) or np.any(values >= counts[None, :]): + raise ValueError("cells contain an invalid active or absent index") + return values, dimensions + + +class SetBPESerializer: + """Lossless mapping between ``[T,D]`` physical cells and token IDs. + + Merge lists are interpreted in creation order. Token IDs are therefore + stable across physical save/load cycles, and decoding can infer any positive + trajectory length directly from stream coverage. + """ + + scheme = "set_bpe" + + def __init__( + self, + *, + bins: Sequence[int], + lookahead: int, + spatial_merges: Sequence[Sequence[int]], + temporal_merges: Sequence[Sequence[int]], + ) -> None: + self.lookahead = _integer(lookahead, "lookahead", positive=True) + self.bins = tuple( + _integer(value, f"bins[{index}]", positive=True) + for index, value in enumerate(bins) + ) + if not self.bins: + raise ConfigError("bins must contain at least one cell slot") + self.spatial_merges = _rules(spatial_merges, 2, "spatial_merges") + self.temporal_merges = _rules(temporal_merges, 3, "temporal_merges") + _validate_rule_references( + atom_count=sum(self.bins), + lookahead=self.lookahead, + spatial_merges=self.spatial_merges, + temporal_merges=self.temporal_merges, + ) + try: + self.spatial = SpatialVocabulary(self.bins, self.spatial_merges) + self.vocabulary = TemporalVocabulary( + self.spatial, self.lookahead, self.temporal_merges + ) + except ValueError as error: + raise ConfigError(str(error)) from error + self.patterns = self.vocabulary.patterns + self.root = self.vocabulary.root + self.vocab_size = self.vocabulary.vocab_size + self.token_cell_counts = np.asarray( + [ + sum( + int(self.spatial.support[spatial_token]).bit_count() + for _, spatial_token in self.patterns[token] + ) + for token in range(self.vocab_size) + ], + dtype=np.int64, + ) + + def encode( + self, + cells: np.ndarray, + *, + active_dims: Sequence[int] | None = None, + ) -> list[int]: + """Serialize one non-empty trajectory without emitting absent slots. + + Compact ``[T,D_active]`` input may provide explicit global + ``active_dims``. Without it, compact input uses prefix dimensions. + Dense ``[T,D_max]`` input may mark absent slots with ``-1``. + """ + values, _ = _dense_cells(cells, self.bins, active_dims) + return self.vocabulary.encode(values) + + def decode( + self, + tokens: Sequence[int], + *, + active_dims: Sequence[int] | None = None, + ) -> np.ndarray: + """Recover dense cells, writing ``-1`` in semantically absent slots.""" + if len(tokens) < 1: + raise ValueError("token stream must not be empty") + checked = [] + for index, token in enumerate(tokens): + if isinstance(token, bool) or not isinstance(token, Integral): + raise ValueError(f"tokens[{index}] must be an integer") + value = int(token) + if not 0 <= value < self.vocab_size: + raise ValueError(f"tokens[{index}] is outside the vocabulary") + checked.append(value) + return self.vocabulary.decode(checked, active_dims=active_dims) + + def to_dict(self) -> dict[str, object]: + """Return the intentionally minimal, path-independent BPE artifact.""" + return { + "spatial_merges": [list(rule) for rule in self.spatial_merges], + "temporal_merges": [list(rule) for rule in self.temporal_merges], + } + + @classmethod + def from_dict( + cls, + payload: object, + *, + bins: Sequence[int], + lookahead: int, + ) -> SetBPESerializer: + if not isinstance(payload, Mapping) or not all( + isinstance(key, str) for key in payload + ): + raise ConfigError(f"{BPE_NAME} must contain one JSON object") + allowed = {"spatial_merges", "temporal_merges"} + unknown = sorted(set(payload) - allowed) + if unknown: + raise ConfigError( + f"{BPE_NAME} contains unknown field(s): {', '.join(unknown)}" + ) + missing = sorted(allowed - set(payload)) + if missing: + raise ConfigError(f"{BPE_NAME} is missing field(s): {', '.join(missing)}") + return cls( + bins=bins, + lookahead=lookahead, + spatial_merges=payload["spatial_merges"], + temporal_merges=payload["temporal_merges"], + ) + + @classmethod + def from_pretrained( + cls, + directory: str | Path, + *, + bins: Sequence[int], + lookahead: int, + ) -> SetBPESerializer: + return cls.from_dict(_read_payload(directory), bins=bins, lookahead=lookahead) + + def save_pretrained(self, directory: str | Path) -> Path: + """Write ``bpe.json`` and return its path.""" + path = Path(directory) / BPE_NAME + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(self.to_dict(), indent=2) + "\n", encoding="utf-8") + return path + + +class BPESerializer: + """Lossless mapping between ``[T,D]`` physical cells and plain-BPE ids. + + The FAST-style baseline: cells are serialized row-major over the active + support and merged as an ordinary symbol sequence. It shares the primitive + atoms, the artifact layout, and the total vocabulary budget with + :class:`SetBPESerializer`, so the two can be fitted on one corpus and + compared directly. + + ``lookahead`` bounds a token at ``lookahead * len(bins)`` primitive cells, + the plain-BPE analogue of the Set-BPE rule that a temporal token spans at + most ``lookahead`` control steps. + """ + + scheme = "bpe" + + def __init__( + self, + *, + bins: Sequence[int], + lookahead: int, + merges: Sequence[Sequence[int]], + ) -> None: + self.lookahead = _integer(lookahead, "lookahead", positive=True) + self.bins = tuple( + _integer(value, f"bins[{index}]", positive=True) + for index, value in enumerate(bins) + ) + if not self.bins: + raise ConfigError("bins must contain at least one cell slot") + self.merges = _rules(merges, 2, "merges") + atom_count = sum(self.bins) + for rank, (left, right) in enumerate(self.merges): + child = atom_count + rank + if left >= child or right >= child: + raise ConfigError( + f"merges[{rank}] must reference already-created tokens" + ) + try: + self.vocabulary = BPEVocabulary(self.bins, self.merges) + except ValueError as error: + raise ConfigError(str(error)) from error + self.max_token_cells = self.lookahead * len(self.bins) + self.token_cell_counts = self.vocabulary.token_cells + oversized = int(np.count_nonzero(self.token_cell_counts > self.max_token_cells)) + if oversized: + raise ConfigError( + f"{oversized} merge rule(s) cover more than the {self.max_token_cells} " + "primitive cells allowed by lookahead" + ) + self.vocab_size = self.vocabulary.vocab_size + + def encode( + self, + cells: np.ndarray, + *, + active_dims: Sequence[int] | None = None, + ) -> list[int]: + """Serialize one non-empty trajectory without emitting absent slots. + + Input conventions match :meth:`SetBPESerializer.encode`: compact + ``[T,D_active]`` may carry explicit global ``active_dims`` and + otherwise uses prefix dimensions, while dense ``[T,D_max]`` marks + absent slots with ``-1``. + """ + values, _ = _dense_cells(cells, self.bins, active_dims) + return self.vocabulary.encode(values) + + def decode( + self, + tokens: Sequence[int], + *, + active_dims: Sequence[int] | None = None, + ) -> np.ndarray: + """Recover dense cells, writing ``-1`` in semantically absent slots.""" + if len(tokens) < 1: + raise ValueError("token stream must not be empty") + checked = [] + for index, token in enumerate(tokens): + if isinstance(token, bool) or not isinstance(token, Integral): + raise ValueError(f"tokens[{index}] must be an integer") + value = int(token) + if not 0 <= value < self.vocab_size: + raise ValueError(f"tokens[{index}] is outside the vocabulary") + checked.append(value) + return self.vocabulary.decode(checked, active_dims=active_dims) + + def to_dict(self) -> dict[str, object]: + """Return the intentionally minimal, path-independent BPE artifact.""" + return {"merges": [list(rule) for rule in self.merges]} + + @classmethod + def from_dict( + cls, + payload: object, + *, + bins: Sequence[int], + lookahead: int, + ) -> BPESerializer: + if not isinstance(payload, Mapping) or not all( + isinstance(key, str) for key in payload + ): + raise ConfigError(f"{BPE_NAME} must contain one JSON object") + unknown = sorted(set(payload) - {"merges"}) + if unknown: + raise ConfigError( + f"{BPE_NAME} contains unknown field(s): {', '.join(unknown)}" + ) + if "merges" not in payload: + raise ConfigError(f"{BPE_NAME} is missing field(s): merges") + return cls(bins=bins, lookahead=lookahead, merges=payload["merges"]) + + @classmethod + def from_pretrained( + cls, + directory: str | Path, + *, + bins: Sequence[int], + lookahead: int, + ) -> BPESerializer: + return cls.from_dict(_read_payload(directory), bins=bins, lookahead=lookahead) + + def save_pretrained(self, directory: str | Path) -> Path: + """Write ``bpe.json`` and return its path.""" + path = Path(directory) / BPE_NAME + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(self.to_dict(), indent=2) + "\n", encoding="utf-8") + return path + + +SERIALIZERS: dict[str, type[SetBPESerializer] | type[BPESerializer]] = { + SetBPESerializer.scheme: SetBPESerializer, + BPESerializer.scheme: BPESerializer, +} + + +def serializer_class( + config: ActionCodec2Config, +) -> type[SetBPESerializer] | type[BPESerializer]: + """Return the serializer the configured tokenization scheme requires.""" + return SERIALIZERS[config.set_bpe.scheme] diff --git a/runtime/tokenization/bpe_training.py b/runtime/tokenization/bpe_training.py new file mode 100644 index 0000000000000000000000000000000000000000..7836dd26fda40afcf3511a58505e464db8af6ad9 --- /dev/null +++ b/runtime/tokenization/bpe_training.py @@ -0,0 +1,262 @@ +"""Plain-BPE fitting implementations. + +Runtime vocabulary and serialization live in :mod:`bpe_vocabulary`; this module +is loaded only by training entry points, mirroring how ``setbpe.spatial`` and +``setbpe.spatial_training`` are split. +""" + +from __future__ import annotations + +import heapq +from collections import defaultdict +from collections.abc import Iterable, Iterator, Sequence + +import numpy as np + +from .bpe_vocabulary import BPEVocabulary +from ..backend.native_backend import require_native, resolve_backend + +__all__ = ["train_bpe", "cell_chunks"] + + +def cell_chunks( + episodes: Iterable[np.ndarray], + *, + max_chunk_rows: int = 1_000_000, +) -> Iterator[tuple[np.ndarray, np.ndarray]]: + """Yield ``(cells, episode_offsets)`` blocks for the native trainer. + + A fitted corpus already stores contiguous chunks with compact episode + offsets; reusing them avoids one corpus-sized copy and lets each replayed + chunk be released as soon as the trainer has ingested it. + """ + chunks = getattr(episodes, "chunks", None) + offsets = getattr(episodes, "offsets", None) + if getattr(episodes, "_actioncodec2_validated", False) and chunks is not None: + for index, (codes, bounds) in enumerate(zip(chunks, offsets, strict=True)): + if codes is None: + raise RuntimeError("cell corpus chunk has already been consumed") + yield ( + np.ascontiguousarray(codes), + np.ascontiguousarray(bounds, dtype=np.int64), + ) + episodes.release_chunk(index) + return + pending: list[np.ndarray] = [] + lengths: list[int] = [] + rows = 0 + + def flush() -> tuple[np.ndarray, np.ndarray]: + block = np.concatenate(pending, axis=0) + bounds = np.concatenate(([0], np.cumsum(lengths, dtype=np.int64))) + pending.clear() + lengths.clear() + return np.ascontiguousarray(block), bounds + + for episode in episodes: + values = np.asarray(episode) + if values.dtype.kind not in {"i", "u"}: + raise ValueError("episodes must contain integer cell indices") + if values.ndim != 2 or len(values) < 1: + raise ValueError("every episode must have non-empty shape [T, D]") + pending.append(values) + lengths.append(len(values)) + rows += len(values) + if rows >= max_chunk_rows: + yield flush() + rows = 0 + if pending: + yield flush() + + +def _flatten( + chunks: Iterable[tuple[np.ndarray, np.ndarray]], bins: Sequence[int] +) -> tuple[list[list[int]], int]: + """Serialize every episode row-major over its active support.""" + counts = np.asarray(tuple(int(value) for value in bins), dtype=np.int64) + offsets = np.concatenate(([0], np.cumsum(counts))) + streams: list[list[int]] = [] + actions = 0 + for codes, bounds in chunks: + values = np.asarray(codes, dtype=np.int64) + if values.ndim != 2 or values.shape[1] != len(counts): + raise ValueError("cells must have one column per bin count") + if np.any(values < -1) or np.any(values >= counts[None, :]): + raise ValueError("cells must contain -1 or an in-range cell index") + for start, stop in zip(bounds[:-1], bounds[1:], strict=True): + episode = values[int(start) : int(stop)] + if not len(episode): + raise ValueError("every episode must be non-empty") + support = episode >= 0 + if np.any(support != support[:1]): + raise ValueError( + "active dimensions must remain constant within an episode" + ) + present = np.flatnonzero(support[0]) + if not len(present): + raise ValueError( + "every control step needs at least one active coordinate" + ) + shifted = episode[:, present] + offsets[present][None, :] + streams.append(shifted.reshape(-1).tolist()) + actions += len(episode) + if not streams: + raise ValueError("BPE training requires at least one complete trajectory") + return streams, actions + + +def _train_bpe_python( + chunks: Iterable[tuple[np.ndarray, np.ndarray]], + bins: Sequence[int], + *, + vocab_size: int, + max_token_cells: int, + progress: bool = False, +) -> tuple[BPEVocabulary, dict[str, object]]: + """Fit exact greedy BPE with sparse counters over observed ordered pairs. + + The reference keeps whole episodes as Python lists and rebuilds each one + the merge touches. That is quadratic in the worst case and exists to pin + the semantics the native kernel must reproduce, not to fit a corpus. + """ + atoms = int(sum(int(value) for value in bins)) + if vocab_size < atoms: + raise ValueError( + f"vocab_size {vocab_size} is below the {atoms} primitive atoms" + ) + if max_token_cells < 1: + raise ValueError("max_token_cells must be positive") + streams, actions = _flatten(chunks, bins) + initial_length = sum(len(stream) for stream in streams) + length = initial_length + + counts: dict[tuple[int, int], int] = defaultdict(int) + carriers: dict[tuple[int, int], set[int]] = defaultdict(set) + for index, stream in enumerate(streams): + for pair in zip(stream[:-1], stream[1:], strict=True): + counts[pair] += 1 + carriers[pair].add(index) + heap = [(-value, pair) for pair, value in counts.items()] + heapq.heapify(heap) + + token_cells = [1] * atoms + merges: list[tuple[int, int]] = [] + oversized: set[tuple[int, int]] = set() + observed_cells = 1 + while len(merges) + atoms < vocab_size and heap: + negative, pair = heapq.heappop(heap) + frequency = counts.get(pair, 0) + if frequency <= 0 or frequency != -negative or pair in oversized: + continue + span = token_cells[pair[0]] + token_cells[pair[1]] + if span > max_token_cells: + oversized.add(pair) + continue + child = atoms + len(merges) + token_cells.append(span) + observed_cells = max(observed_cells, span) + touched: set[tuple[int, int]] = set() + for index in sorted(carriers.pop(pair, ())): + stream = streams[index] + output: list[int] = [] + position = 0 + while position < len(stream): + if ( + position + 1 < len(stream) + and stream[position] == pair[0] + and stream[position + 1] == pair[1] + ): + if output: + counts[(output[-1], pair[0])] -= 1 + touched.add((output[-1], pair[0])) + counts[(output[-1], child)] += 1 + carriers[(output[-1], child)].add(index) + touched.add((output[-1], child)) + if position + 2 < len(stream): + follower = stream[position + 2] + counts[(pair[1], follower)] -= 1 + touched.add((pair[1], follower)) + counts[(child, follower)] += 1 + carriers[(child, follower)].add(index) + touched.add((child, follower)) + counts[pair] -= 1 + length -= 1 + output.append(child) + position += 2 + else: + output.append(stream[position]) + position += 1 + streams[index] = output + merges.append(pair) + for member in touched: + value = counts.get(member, 0) + if value > 0: + heapq.heappush(heap, (-value, member)) + if progress and len(merges) % 256 == 0: + print( + f" bpe merge {len(merges):>5} tokens/action {length / actions:.4f}", + flush=True, + ) + + vocabulary = BPEVocabulary(tuple(int(value) for value in bins), tuple(merges)) + report: dict[str, object] = { + "episodes": len(streams), + "actions": actions, + "atoms": atoms, + "merges": len(merges), + "vocab_size": vocabulary.vocab_size, + "cells": initial_length, + "tokens_per_action_initial": initial_length / actions, + "tokens_per_action": length / actions, + "compression": initial_length / length, + "max_token_cells": int(max_token_cells), + "max_token_cells_observed": observed_cells, + "oversized_pairs_skipped": len(oversized), + "trainer": "python_sparse_exact", + } + return vocabulary, report + + +def train_bpe( + episodes: Iterable[np.ndarray], + bins: Sequence[int], + *, + vocab_size: int, + max_token_cells: int, + backend: str = "auto", + progress: bool = False, + threads: int = 0, +) -> tuple[BPEVocabulary, dict[str, object]]: + """Fit plain BPE on the row-major flattened primitive-cell stream. + + ``vocab_size`` is the *total* budget including the ``sum(bins)`` primitive + atoms, so it is directly comparable with the Set-BPE budget. A pair whose + child would cover more than ``max_token_cells`` primitive cells is never + merged; this is the plain-BPE analogue of the Set-BPE lookahead, which caps + a temporal token at ``lookahead`` control steps. + + Both backends return the identical greedy merge sequence; the native one is + preferred when built. The native trainer currently runs serially; ``threads`` + is accepted for the shared fitting API and reported as ``requested_threads``, + while the report's ``threads`` records the actual worker count. + """ + if resolve_backend(backend) == "python": + return _train_bpe_python( + cell_chunks(episodes), + bins, + vocab_size=vocab_size, + max_token_cells=max_token_cells, + progress=progress, + ) + merges, report = require_native().train_bpe( + cell_chunks(episodes), + [int(count) for count in bins], + int(vocab_size), + int(max_token_cells), + int(threads), + ) + vocabulary = BPEVocabulary( + tuple(int(count) for count in bins), + tuple((int(left), int(right)) for left, right in merges), + ) + return vocabulary, dict(report) diff --git a/runtime/tokenization/bpe_vocabulary.py b/runtime/tokenization/bpe_vocabulary.py new file mode 100644 index 0000000000000000000000000000000000000000..ebb98f828eb7ec3ff50a15213c3dad1bfa967a92 --- /dev/null +++ b/runtime/tokenization/bpe_vocabulary.py @@ -0,0 +1,179 @@ +"""Runtime plain-BPE vocabulary over the row-major flattened cell stream. + +Set-BPE encodes one control step as an unordered set of primitive cells and +merges spatially before merging across time. Plain BPE is the FAST-style +baseline: the quantized ``[T, D]`` cell matrix is serialized row-major into one +symbol sequence and textbook byte-pair encoding runs on that sequence. + +Atoms are the same per-dimension primitive cells Set-BPE uses, so a token id +carries its dimension. Absent coordinates emit no atom, exactly as in Set-BPE, +which keeps a single vocabulary usable across trajectory arities: two atoms are +adjacent in the stream only when their dimensions are adjacent in that +trajectory's active support. +""" + +from __future__ import annotations + +import heapq +from collections.abc import Sequence +from dataclasses import dataclass + +import numpy as np + +from ..common._validation import active_dimensions + +__all__ = ["BPEVocabulary"] + + +@dataclass +class BPEVocabulary: + """Primitive atoms plus ordered, adjacency-based merge rules.""" + + bins: tuple[int, ...] + merges: tuple[tuple[int, int], ...] + + def __post_init__(self) -> None: + if not self.bins or any(count < 1 for count in self.bins): + raise ValueError("every dimension needs at least one bin") + self.offsets = np.concatenate( + ([0], np.cumsum(np.asarray(self.bins, dtype=np.int64))) + ) + self.atom_count = int(self.offsets[-1]) + self.vocab_size = self.atom_count + len(self.merges) + atom_dims = np.repeat( + np.arange(len(self.bins), dtype=np.int64), + np.asarray(self.bins, dtype=np.int64), + ) + # Every token expands to a fixed run of atoms. Storing the expansions + # once keeps decoding and grammar construction table lookups instead of + # recursive walks. + expansions: list[tuple[int, ...]] = [ + (token,) for token in range(self.atom_count) + ] + for rank, (left, right) in enumerate(self.merges): + token = self.atom_count + rank + if not 0 <= left < token or not 0 <= right < token: + raise ValueError("a merge rule must reference already-created tokens") + expansions.append(expansions[left] + expansions[right]) + self.expansions = tuple(expansions) + self.token_dims = tuple( + tuple(int(atom_dims[atom]) for atom in expansion) + for expansion in expansions + ) + self.token_cells = np.asarray( + [len(expansion) for expansion in expansions], dtype=np.int64 + ) + self._rules = { + (int(left), int(right)): (self.atom_count + rank, rank) + for rank, (left, right) in enumerate(self.merges) + } + + @property + def dimension(self) -> int: + return len(self.bins) + + def atom(self, dimension: int, value: int) -> int: + return int(self.offsets[dimension]) + int(value) + + def _flatten(self, values: np.ndarray, dimensions: tuple[int, ...]) -> list[int]: + """Serialize dense ``[T, D]`` cells row-major over the active support.""" + counts = np.asarray(self.bins, dtype=np.int64) + if np.any(values < -1) or np.any(values >= counts[None, :]): + raise ValueError("cells contain an invalid active or absent index") + support = values >= 0 + if np.any(support != support[:1]): + raise ValueError( + "active dimensions must remain constant within a trajectory" + ) + present = tuple(int(dim) for dim in np.flatnonzero(support[0])) + if present != dimensions: + raise ValueError("cells do not activate exactly the requested dimensions") + if not present: + raise ValueError("every control step needs at least one active coordinate") + active = values[:, list(present)] + base = np.asarray([int(self.offsets[dim]) for dim in present], dtype=np.int64) + return (active + base[None, :]).reshape(-1).tolist() + + def encode(self, cells: np.ndarray) -> list[int]: + """Apply the merges greedily in creation-rank order to one trajectory. + + A rank-ordered priority queue over adjacent pairs applies the lowest + available rule first and, within one rule, left to right -- the exact + order both trainers count in. Ranked greedy application is also what + makes the encoder reproducible: the same cell stream always yields the + same tokens regardless of how the search that learned the rules ran. + """ + values = np.asarray(cells, dtype=np.int64) + if values.ndim != 2 or len(values) < 1 or values.shape[1] < 1: + raise ValueError("cells must have non-empty shape [T, D]") + dimensions = tuple(int(dim) for dim in np.flatnonzero(values[0] >= 0)) + stream = self._flatten(values, dimensions) + count = len(stream) + previous = list(range(-1, count - 1)) + following = list(range(1, count + 1)) + following[-1] = -1 + alive = [True] * count + rules = self._rules + heap: list[tuple[int, int]] = [] + for index in range(count - 1): + rule = rules.get((stream[index], stream[index + 1])) + if rule is not None: + heap.append((rule[1], index)) + heapq.heapify(heap) + while heap: + rank, index = heapq.heappop(heap) + if not alive[index]: + continue + follower = following[index] + if follower < 0 or not alive[follower]: + continue + rule = rules.get((stream[index], stream[follower])) + if rule is None or rule[1] != rank: + continue + stream[index] = rule[0] + alive[follower] = False + after = following[follower] + following[index] = after + if after >= 0: + previous[after] = index + successor = rules.get((stream[index], stream[after])) + if successor is not None: + heapq.heappush(heap, (successor[1], index)) + before = previous[index] + if before >= 0: + predecessor = rules.get((stream[before], stream[index])) + if predecessor is not None: + heapq.heappush(heap, (predecessor[1], before)) + return [stream[index] for index in range(count) if alive[index]] + + def decode( + self, + tokens: Sequence[int], + *, + active_dims: Sequence[int] | None = None, + ) -> np.ndarray: + """Expand tokens to dense cells, writing ``-1`` in absent coordinates.""" + dimensions = active_dimensions(active_dims, self.dimension) + atoms: list[int] = [] + for index, raw_token in enumerate(tokens): + token = int(raw_token) + if not 0 <= token < self.vocab_size: + raise ValueError(f"tokens[{index}] is outside the vocabulary") + atoms.extend(self.expansions[token]) + width = len(dimensions) + if not atoms or len(atoms) % width: + raise ValueError( + "token stream does not cover a whole number of routed control steps" + ) + values = np.asarray(atoms, dtype=np.int64).reshape(-1, width) + starts = np.asarray( + [int(self.offsets[dim]) for dim in dimensions], dtype=np.int64 + ) + stops = np.asarray( + [int(self.offsets[dim + 1]) for dim in dimensions], dtype=np.int64 + ) + if np.any(values < starts[None, :]) or np.any(values >= stops[None, :]): + raise ValueError("token stream does not follow the routed dimension order") + cells = np.full((len(values), self.dimension), -1, dtype=np.int64) + cells[:, list(dimensions)] = values - starts[None, :] + return cells diff --git a/runtime/tokenization/generation.py b/runtime/tokenization/generation.py new file mode 100644 index 0000000000000000000000000000000000000000..806de3c50f3f7d01bc4559b1cae6f4425a5de804 --- /dev/null +++ b/runtime/tokenization/generation.py @@ -0,0 +1,138 @@ +"""Explicit action-token constraints for Hugging Face generate().""" + +from __future__ import annotations + +import torch +from transformers import LogitsProcessor + +from ..routing.cache import BoundedCache + + +class ActionCodec2LogitsProcessor(LogitsProcessor): + """Constrain generated action IDs, then allow EOS only at complete coverage. + + Args: + codec: Fitted ActionCodec2 instance; its current grammar is captured. + horizon: Number of source action frames requested. + fps: Explicit source/output sampling rate in Hz. + prompt_length: Number of leading input IDs to ignore, including padding. + For encoder-decoder models this is the decoder prompt length. + eos_token_id: Model EOS ID, outside the action vocabulary interval. + token_offset: Explicit contiguous offset of codec IDs in model vocabulary. + No vocabulary mapping is inferred or installed into the model. + pad_token_id: Optional model padding ID used after a completed EOS. + + Each call takes input_ids (B,L) and scores (B,V), including expanded beam rows. + Prefixes move to CPU once per call; device masks are shared by grammar state + in a bounded cache. The original scores tensor is never modified. + """ + + def __init__( + self, + codec, + horizon, + *, + fps, + prompt_length, + eos_token_id, + token_offset=0, + pad_token_id=None, + ): + for name, value in { + "prompt_length": prompt_length, + "eos_token_id": eos_token_id, + "token_offset": token_offset, + }.items(): + if type(value) is not int or value < 0: + raise ValueError(f"{name} must be a nonnegative integer") + if pad_token_id is not None and ( + type(pad_token_id) is not int or pad_token_id < 0 + ): + raise ValueError("pad_token_id must be a nonnegative integer") + self.grammar = codec.grammar(horizon, fps=fps) + self.prompt_length = prompt_length + self.eos_token_id = eos_token_id + self.pad_token_id = pad_token_id + self.token_offset = token_offset + stop = token_offset + self.grammar.vocab_size + if token_offset <= eos_token_id < stop: + raise ValueError("eos_token_id must be outside the action token interval") + if pad_token_id is not None and token_offset <= pad_token_id < stop: + raise ValueError("pad_token_id must be outside the action token interval") + self._mask_cache = BoundedCache(64) + + def _state(self, row): + if self.eos_token_id in row: + end = row.index(self.eos_token_id) + if any(t not in (self.eos_token_id, self.pad_token_id) for t in row[end:]): + raise ValueError( + "only EOS/padding may follow completed action generation" + ) + tokens = [t - self.token_offset for t in row[:end]] + if not self.grammar.is_complete(tokens): + raise ValueError("EOS occurred before action coverage was complete") + else: + tokens = [t - self.token_offset for t in row] + state = self.grammar._consume(tokens) + if state is None: + raise ValueError( + "generated prefix violates the action grammar; check prompt_length and token_offset" + ) + return state + + def _mask(self, state, scores): + key = (state, str(scores.device), scores.shape[-1]) + mask = self._mask_cache.get(key) + if mask is None: + mask = torch.zeros(scores.shape[-1], device=scores.device, dtype=torch.bool) + index, local = state + if index == len(self.grammar.segments): + mask[self.eos_token_id] = True + else: + segment = self.grammar.segments[index] + allowed = self.grammar._grammars[index]._next_token_mask_from_state( + local + ) + start = self.token_offset + segment.profile.spec.token_offset + mask[start : start + len(allowed)] = torch.as_tensor( + allowed, device=scores.device + ) + self._mask_cache[key] = mask + return mask + + def __call__( + self, input_ids: torch.LongTensor, scores: torch.FloatTensor + ) -> torch.FloatTensor: + if ( + input_ids.ndim != 2 + or scores.ndim != 2 + or input_ids.shape[0] != scores.shape[0] + ): + raise ValueError( + "expected input_ids(B,L) and scores(B,V) with the same batch size" + ) + if input_ids.shape[1] < self.prompt_length: + raise ValueError("prompt_length exceeds the supplied input length") + required = max( + self.token_offset + self.grammar.vocab_size, + self.eos_token_id + 1, + 0 if self.pad_token_id is None else self.pad_token_id + 1, + ) + if scores.shape[1] < required: + raise ValueError( + "model vocabulary does not cover configured action/EOS/padding IDs" + ) + rows = input_ids[:, self.prompt_length :].detach().cpu().tolist() + output = scores.clone() + groups = {} + for index, row in enumerate(rows): + groups.setdefault(self._state(row), []).append(index) + for state, indices in groups.items(): + output[indices] = scores[indices].masked_fill( + ~self._mask(state, scores), -torch.inf + ) + if not torch.isfinite(output).any(dim=1).all(): + raise ValueError( + "all legal action logits were suppressed; check other generation processors" + ) + return output diff --git a/runtime/tokenization/grammar.py b/runtime/tokenization/grammar.py new file mode 100644 index 0000000000000000000000000000000000000000..d838da2749523e756e5d134c740e89e2034b5966 --- /dev/null +++ b/runtime/tokenization/grammar.py @@ -0,0 +1,232 @@ +"""Constrained-generation grammars for the two ActionCodec2 BPE schemes.""" + +from __future__ import annotations + +from collections.abc import Sequence +from numbers import Integral + +import numpy as np + +from ..common._validation import active_dimensions +from .bpe import BPESerializer, SetBPESerializer +from ..routing.cache import BoundedCache + + +class _PrefixGrammar: + """Shared public queries over scheme-specific prefix state transitions.""" + + def next_token_mask(self, tokens: Sequence[int]) -> np.ndarray: + """Return a boolean vocabulary mask for the next generated token.""" + state = self._state(tokens) + if state is None: + raise ValueError("token prefix is not a legal ActionCodec2 stream") + return self._next_token_mask_from_state(state) + + def is_complete(self, tokens: Sequence[int]) -> bool: + """Whether ``tokens`` is exactly one complete trajectory of this horizon.""" + state = self._state(tokens) + return state is not None and self._state_is_complete(state) + + +class ActionCodec2Grammar(_PrefixGrammar): + """Prefix grammar for one exact, fixed-horizon rooted token stream. + + The state is one support bitmask per timestep. A token is legal only when + its rooted pattern starts at the earliest incomplete timestep, does not + overlap covered dimensions, and stays inside ``horizon``. + + Args: + serializer: Fitted Set-BPE rules defining each token's physical cells. + horizon: Number of codec timesteps the stream must cover. + active_dims: Increasing physical dimension indices present in the stream; + omitted means all dimensions. + """ + + def __init__( + self, + serializer: SetBPESerializer, + *, + horizon: int, + active_dims: Sequence[int] | None = None, + ) -> None: + if isinstance(horizon, bool) or not isinstance(horizon, int) or horizon < 1: + raise ValueError("horizon must be a positive integer") + self.horizon = horizon + self.vocab_size = serializer.vocab_size + dims = active_dimensions(active_dims, serializer.spatial.dimension) + self._full = sum(1 << dimension for dimension in dims) + self._root_anchor = np.asarray( + [ + serializer.spatial.anchor[serializer.root[token]] + for token in range(self.vocab_size) + ], + dtype=np.int64, + ) + self._patterns = tuple( + tuple( + (offset, int(serializer.spatial.support[spatial_token])) + for offset, spatial_token in serializer.patterns[token] + ) + for token in range(self.vocab_size) + ) + self._mask_cache = BoundedCache(64) + self._state_cache = BoundedCache(256) + self._state_cache[()] = self._initial_state() + + def _advance(self, state: tuple[int, ...], token: int) -> tuple[int, ...] | None: + if not 0 <= token < self.vocab_size: + return None + coverage = list(state) + front = next( + (index for index, mask in enumerate(coverage) if mask != self._full), + self.horizon, + ) + if front == self.horizon: + return None + uncovered = self._full & ~coverage[front] + expected_anchor = int(uncovered & -uncovered).bit_length() - 1 + if int(self._root_anchor[token]) != expected_anchor: + return None + for offset, support in self._patterns[token]: + step = front + offset + if step >= self.horizon or coverage[step] & support: + return None + coverage[step] |= support + return tuple(coverage) + + def _initial_state(self) -> tuple[int, ...]: + """Return the empty coverage state for incremental internal validation.""" + return (0,) * self.horizon + + def _state_is_complete(self, state: tuple[int, ...]) -> bool: + """Whether an incrementally consumed state covers the whole trajectory.""" + return all(mask == self._full for mask in state) + + def _state(self, tokens: Sequence[int]) -> tuple[int, ...] | None: + checked: list[int] = [] + for raw_token in tokens: + if isinstance(raw_token, bool) or not isinstance(raw_token, Integral): + return None + checked.append(int(raw_token)) + key = tuple(checked) + if not key: + return self._initial_state() + missing = object() + cached = self._state_cache.get(key, missing) + if cached is not missing: + return cached + parent = key[:-1] + state = self._state_cache.get(parent, missing) + if state is missing: + state = self._initial_state() + for token in parent: + state = self._advance(state, token) + if state is None: + break + self._state_cache[parent] = state + if state is not None: + state = self._advance(state, key[-1]) + self._state_cache[key] = state + return state + + def _next_token_mask_from_state(self, state: tuple[int, ...]) -> np.ndarray: + """Build one next-token mask without replaying an already-known prefix.""" + if self._state_is_complete(state): + return np.zeros(self.vocab_size, dtype=bool) + cached = self._mask_cache.get(state) + if cached is not None: + return cached.copy() + front = next(index for index, mask in enumerate(state) if mask != self._full) + uncovered = self._full & ~state[front] + expected_anchor = int(uncovered & -uncovered).bit_length() - 1 + mask = self._root_anchor == expected_anchor + for token in np.flatnonzero(mask): + for offset, support in self._patterns[int(token)]: + step = front + offset + if step >= self.horizon or state[step] & support: + mask[token] = False + break + self._mask_cache[state] = mask + return mask.copy() + + +class ActionCodec2BPEGrammar(_PrefixGrammar): + """Prefix grammar for one exact, fixed-horizon plain-BPE token stream. + + Plain BPE emits the active cells of a trajectory in row-major order, so the + state is just how many cells the prefix already covers. A token is legal + only when its expansion continues the routed dimension cycle at that + position and stays inside ``horizon``. + + Args: + serializer: Fitted plain-BPE rules defining each token's cell sequence. + horizon: Number of codec timesteps the stream must cover. + active_dims: Increasing physical dimension indices present in the stream; + omitted means all dimensions. + """ + + def __init__( + self, + serializer: BPESerializer, + *, + horizon: int, + active_dims: Sequence[int] | None = None, + ) -> None: + if isinstance(horizon, bool) or not isinstance(horizon, int) or horizon < 1: + raise ValueError("horizon must be a positive integer") + self.horizon = horizon + self.vocab_size = serializer.vocab_size + dims = active_dimensions(active_dims, serializer.vocabulary.dimension) + self._width = len(dims) + self._cells = horizon * self._width + self._token_cells = np.asarray(serializer.token_cell_counts, dtype=np.int64) + token_dims = serializer.vocabulary.token_dims + self._phase = np.zeros((self._width, self.vocab_size), dtype=bool) + for phase in range(self._width): + for token in range(self.vocab_size): + expansion = token_dims[token] + self._phase[phase, token] = all( + dimension == dims[(phase + offset) % self._width] + for offset, dimension in enumerate(expansion) + ) + self._mask_cache = BoundedCache(64) + + def _initial_state(self) -> tuple[int, ...]: + return (0,) + + def _state_is_complete(self, state: tuple[int, ...]) -> bool: + return state[0] == self._cells + + def _advance(self, state: tuple[int, ...], token: int) -> tuple[int, ...] | None: + if not 0 <= token < self.vocab_size: + return None + covered = state[0] + length = int(self._token_cells[token]) + if covered + length > self._cells: + return None + if not self._phase[covered % self._width, token]: + return None + return (covered + length,) + + def _state(self, tokens: Sequence[int]) -> tuple[int, ...] | None: + state: tuple[int, ...] | None = self._initial_state() + for raw_token in tokens: + if isinstance(raw_token, bool) or not isinstance(raw_token, Integral): + return None + assert state is not None + state = self._advance(state, int(raw_token)) + if state is None: + return None + return state + + def _next_token_mask_from_state(self, state: tuple[int, ...]) -> np.ndarray: + covered = state[0] + if covered == self._cells: + return np.zeros(self.vocab_size, dtype=bool) + cached = self._mask_cache.get(covered) + if cached is None: + cached = self._phase[covered % self._width] & ( + self._token_cells <= self._cells - covered + ) + self._mask_cache[covered] = cached + return cached.copy() diff --git a/runtime/tokenization/setbpe_spatial.py b/runtime/tokenization/setbpe_spatial.py new file mode 100644 index 0000000000000000000000000000000000000000..010030c16fd052635e0bbffaeb5e594761564d07 --- /dev/null +++ b/runtime/tokenization/setbpe_spatial.py @@ -0,0 +1,198 @@ +"""Runtime spatial Set-BPE vocabulary and lossless row serialization.""" + +from __future__ import annotations + +import heapq +from collections.abc import Iterable, Sequence +from dataclasses import dataclass + +import numpy as np + +from ..common._validation import active_dimensions + +__all__ = ["SpatialVocabulary", "aggregate_records"] + + +def aggregate_records( + codes: np.ndarray, bins: Sequence[int] | None = None +) -> tuple[np.ndarray, np.ndarray]: + """Lossless micro-aggregation of identical quantized actions. + + Replacing ``w`` identical records by one record of weight ``w`` leaves every + pair frequency and therefore every greedy decision unchanged, because + identical actions follow identical merges forever. + """ + values = np.asarray(codes, dtype=np.int64) + if values.ndim != 2: + raise ValueError("codes must have shape [U, D]") + if bins is not None: + dimensions = tuple(int(count) for count in bins) + if len(dimensions) != values.shape[1]: + raise ValueError("bins must have one entry per code coordinate") + counts = np.asarray(dimensions, dtype=np.int64) + valid = values >= 0 + if np.any(values < -1) or np.any(values >= counts[None, :]): + raise ValueError("codes must contain -1 or an in-range cell index") + cardinality = int(np.prod(np.asarray(dimensions, dtype=object))) + if np.all(valid) and cardinality <= np.iinfo(np.int64).max: + keys = np.ravel_multi_index(values.T, dimensions) + unique_keys, counts = np.unique(keys, return_counts=True) + unique = np.stack(np.unravel_index(unique_keys, dimensions), axis=1) + return unique.astype(np.int64, copy=False), counts.astype(np.int64) + unique, counts = np.unique(values, axis=0, return_counts=True) + return unique, counts.astype(np.int64) + + +@dataclass +class SpatialVocabulary: + """Atoms, merge rules and the decoder metadata they induce.""" + + bins: tuple[int, ...] + merges: tuple[tuple[int, int], ...] + + def __post_init__(self) -> None: + if not self.bins or any(count < 1 for count in self.bins): + raise ValueError("every dimension needs at least one bin") + self.offsets = np.concatenate( + ([0], np.cumsum(np.asarray(self.bins, dtype=np.int64))) + ) + self.atom_count = int(self.offsets[-1]) + self.vocab_size = self.atom_count + len(self.merges) + dimension = len(self.bins) + support = np.zeros(self.vocab_size, dtype=np.int64) + expansion = np.full((self.vocab_size, dimension), -1, dtype=np.int64) + for index, count in enumerate(self.bins): + for value in range(count): + token = int(self.offsets[index]) + value + support[token] = 1 << index + expansion[token, index] = value + for rank, (left, right) in enumerate(self.merges): + token = self.atom_count + rank + if left >= token or right >= token: + raise ValueError("a merge rule must reference already-created tokens") + if support[left] & support[right]: + raise ValueError("a merge rule joins two overlapping supports") + support[token] = support[left] | support[right] + merged = np.maximum(expansion[left], expansion[right]) + expansion[token] = merged + self.support = support + self.expansion = expansion + self.anchor = np.asarray( + [int(mask & -mask).bit_length() - 1 for mask in support], dtype=np.int64 + ) + self.full_support = (1 << dimension) - 1 + self._rules: dict[tuple[int, int], tuple[int, int]] = {} + for rank, (left, right) in enumerate(self.merges): + key = (min(left, right), max(left, right)) + self._rules[key] = (self.atom_count + rank, rank) + + @property + def dimension(self) -> int: + return len(self.bins) + + def atom(self, dimension: int, value: int) -> int: + return int(self.offsets[dimension]) + int(value) + + def encode( + self, + code: Sequence[int], + *, + active_dims: Sequence[int] | None = None, + ) -> list[int]: + """Deterministic online Set-BPE encoding (Algorithm 8). + + Merge rules are applied in creation-rank order, which is what makes the + encoder reproducible: the same quantized action always yields the same + token set regardless of how the search that learned the rules ran. + ``-1`` in a dense row, or a missing coordinate in ``active_dims``, means + semantic absence: it emits no atom and can never participate in a merge. + """ + compact = [int(value) for value in code] + if active_dims is None: + if len(compact) != self.dimension: + raise ValueError( + f"a dense quantized action must have {self.dimension} slots" + ) + pairs = enumerate(compact) + else: + dimensions = active_dimensions(active_dims, self.dimension) + if len(dimensions) != len(compact): + raise ValueError( + "active_dims must have one entry per compact coordinate" + ) + pairs = zip(dimensions, compact, strict=True) + tokens = set() + for index, value in pairs: + if value == -1: + if active_dims is not None: + raise ValueError("compact coordinates must not contain -1") + continue + if not 0 <= value < self.bins[index]: + raise ValueError(f"bin {value} is out of range for dimension {index}") + tokens.add(self.atom(index, value)) + if not tokens: + raise ValueError( + "a quantized action must contain at least one active coordinate" + ) + heap: list[tuple[int, int, int, int]] = [] + ordered = sorted(tokens) + for position, left in enumerate(ordered): + for right in ordered[position + 1 :]: + rule = self._rules.get((left, right)) + if rule is not None: + heapq.heappush(heap, (rule[1], rule[0], left, right)) + while heap: + _, child, left, right = heapq.heappop(heap) + if left not in tokens or right not in tokens: + continue + tokens.discard(left) + tokens.discard(right) + for other in tokens: + key = (min(child, other), max(child, other)) + rule = self._rules.get(key) + if rule is not None: + heapq.heappush(heap, (rule[1], rule[0], key[0], key[1])) + tokens.add(child) + return self.serialize(tokens) + + def serialize(self, tokens: Iterable[int]) -> list[int]: + """Canonical order for autoregressive supervision. + + Set semantics carry no order, but a policy must be trained on one fixed + target. Tokens within an action have pairwise disjoint supports, so + their anchor dimensions are distinct and sorting by anchor is a total + order. Anchor order is preferred over raw token id because it is the + one that keeps constrained decoding feasible: the next token's anchor is + always the lowest uncovered dimension, and a primitive atom for that + dimension always exists, so no prefix can paint itself into a corner. + """ + return sorted(tokens, key=lambda token: (int(self.anchor[token]), int(token))) + + def decode( + self, + tokens: Sequence[int], + *, + active_dims: Sequence[int] | None = None, + ) -> np.ndarray: + """Expand tokens to a dense row, using ``-1`` for absent coordinates.""" + code = np.full(self.dimension, -1, dtype=np.int64) + covered = 0 + for token in tokens: + value = int(token) + if not 0 <= value < self.vocab_size: + raise ValueError(f"token {value} is outside the vocabulary") + if covered & int(self.support[value]): + raise ValueError("token supports overlap; this is not a valid action") + covered |= int(self.support[value]) + local = self.expansion[value] + code = np.where(local >= 0, local, code) + if active_dims is None: + expected = self.full_support + else: + dimensions = active_dimensions(active_dims, self.dimension) + expected = sum(1 << dimension for dimension in dimensions) + if covered != expected: + raise ValueError( + "token set does not exactly cover the requested active dimensions" + ) + return code diff --git a/runtime/tokenization/setbpe_spatial_training.py b/runtime/tokenization/setbpe_spatial_training.py new file mode 100644 index 0000000000000000000000000000000000000000..90c1e9b43c6d64cf9e62d5854db44fe20818cf99 --- /dev/null +++ b/runtime/tokenization/setbpe_spatial_training.py @@ -0,0 +1,334 @@ +"""Spatial Set-BPE fitting implementations. + +Runtime vocabulary and serialization live in ``spatial``; this module is +loaded only by training entry points. +""" + +from __future__ import annotations + +import heapq +from collections.abc import Sequence + +import numpy as np + +from ..backend.native_backend import require_native, resolve_backend +from .setbpe_spatial import SpatialVocabulary + +__all__ = ["train_spatial"] + + +class _SparseSpatialTrainer: + """Columnar Set-BPE state with counters for observed pairs only.""" + + def __init__( + self, + codes: np.ndarray, + weights: np.ndarray, + bins: Sequence[int], + ) -> None: + self.bins = tuple(int(value) for value in bins) + self.dimension = len(self.bins) + self.records = np.asarray(codes, dtype=np.int64) + self.weights = np.asarray(weights, dtype=np.int64) + if self.records.ndim != 2 or self.records.shape[1] != self.dimension: + raise ValueError("codes must have shape [U, D] matching the bin counts") + if self.weights.shape != (self.records.shape[0],): + raise ValueError("weights must have one entry per record") + counts = np.asarray(self.bins, dtype=np.int64) + if np.any(self.records < -1) or np.any(self.records >= counts[None, :]): + raise ValueError("codes must contain -1 or an in-range cell index") + if np.any(np.all(self.records < 0, axis=1)): + raise ValueError("every record must contain at least one active coordinate") + self.count = int(self.records.shape[0]) + self.offsets = np.concatenate( + ([0], np.cumsum(np.asarray(self.bins, dtype=np.int64))) + ) + self.atom_count = int(self.offsets[-1]) + + self.occurrence: list[np.ndarray] = [] + self.support: list[int] = [] + self.anchor: list[int] = [] + self.frequency: list[int] = [] + for index in range(self.dimension): + column = self.records[:, index] + for value in range(self.bins[index]): + members = np.flatnonzero(column == value).astype(np.int64) + self._append_token(members, 1 << index, index) + self._scratch = np.zeros(self.count, dtype=bool) + self.merges: list[tuple[int, int]] = [] + self.statistics = { + "sparse_updates": 0, + "sparse_heap_rebuilds": 0, + "sparse_pair_types_peak": 0, + } + # Row-wise mirror used to initialize and incrementally update pair counts. + self.member = np.where( + self.records >= 0, + self.offsets[:-1][None, :] + self.records, + -1, + ) + self._slot_pairs = tuple( + (first, second) + for first in range(self.dimension) + for second in range(first + 1, self.dimension) + ) + # The observed-pair backend is built on demand. Keys use two fixed + # 32-bit token fields, so appending vocabulary entries never invalidates + # the counter map as a width-based packed index would. + self._sparse_counts: dict[int, int] | None = None + self._sparse_heap: list[tuple[int, int, int]] = [] + + def _append_token(self, members: np.ndarray, support: int, anchor: int) -> int: + token = len(self.occurrence) + self.occurrence.append(members) + self.support.append(support) + self.anchor.append(anchor) + self.frequency.append(int(self.weights[members].sum()) if members.size else 0) + return token + + def _intersect(self, left: np.ndarray, right: np.ndarray) -> np.ndarray: + if left.size > right.size: + left, right = right, left + if left.size == 0: + return left + self._scratch[left] = True + result = right[self._scratch[right]] + self._scratch[left] = False + return result + + @staticmethod + def _sparse_key(left: int, right: int) -> int: + low, high = (left, right) if left < right else (right, left) + return (int(low) << 32) | int(high) + + @staticmethod + def _sparse_pair(key: int) -> tuple[int, int]: + return key >> 32, key & 0xFFFFFFFF + + def build_sparse_index(self) -> None: + """Materialize exact counters only for currently observed pairs once.""" + if self._sparse_counts is not None: + return + width = len(self.occurrence) + counts = np.zeros(width * width, dtype=np.int64) + for first, second in self._slot_pairs: + left = self.member[:, first] + right = self.member[:, second] + valid = (left >= 0) & (right >= 0) + if not valid.any(): + continue + low = np.minimum(left[valid], right[valid]) + high = np.maximum(left[valid], right[valid]) + np.add.at(counts, low * width + high, self.weights[valid]) + observed = np.flatnonzero(counts) + sparse: dict[int, int] = {} + heap: list[tuple[int, int, int]] = [] + for packed in observed: + left, right = divmod(int(packed), width) + key = self._sparse_key(left, right) + value = int(counts[packed]) + sparse[key] = value + heap.append((-value, left, right)) + heapq.heapify(heap) + self._sparse_counts = sparse + self._sparse_heap = heap + self.statistics["sparse_pair_types_peak"] = len(sparse) + + def sparse_best_pair(self) -> tuple[int, int, int] | None: + """Return the exact winner from the observed-pair max heap.""" + if self._sparse_counts is None: + self.build_sparse_index() + assert self._sparse_counts is not None + while self._sparse_heap: + negative, left, right = heapq.heappop(self._sparse_heap) + value = self._sparse_counts.get(self._sparse_key(left, right), 0) + if value == -negative: + return left, right, value + return None + + def _update_sparse_index( + self, + left: int, + right: int, + child: int, + rows: np.ndarray, + members: np.ndarray, + mass: int, + ) -> None: + if self._sparse_counts is None: + return + counts = self._sparse_counts + heap = self._sparse_heap + updates = 0 + + def change(key: int, delta: int) -> None: + nonlocal updates + new = counts.get(key, 0) + delta + if new < 0: + pair = self._sparse_pair(key) + raise RuntimeError( + f"sparse spatial pair counter became negative for {pair}" + ) + if new: + counts[key] = new + pair = self._sparse_pair(key) + heapq.heappush(heap, (-new, pair[0], pair[1])) + else: + counts.pop(key, None) + updates += 1 + + change(self._sparse_key(left, right), -mass) + other_mass = np.zeros(len(self.occurrence), dtype=np.int64) + weights = self.weights[members] + for slot in range(rows.shape[1]): + values = rows[:, slot] + valid = (values >= 0) & (values != left) & (values != right) + if valid.any(): + np.add.at(other_mass, values[valid], weights[valid]) + for raw_other in np.flatnonzero(other_mass): + other = int(raw_other) + value = int(other_mass[other]) + change(self._sparse_key(left, other), -value) + change(self._sparse_key(right, other), -value) + change(self._sparse_key(child, other), value) + + self.statistics["sparse_updates"] += updates + self.statistics["sparse_pair_types_peak"] = max( + self.statistics["sparse_pair_types_peak"], len(counts) + ) + if len(heap) > 2 * len(counts) + 65_536: + self._sparse_heap = [ + (-value, *self._sparse_pair(key)) for key, value in counts.items() + ] + heapq.heapify(self._sparse_heap) + self.statistics["sparse_heap_rebuilds"] += 1 + + def merge(self, left: int, right: int) -> int: + """Apply one exact incremental merge update.""" + members = self._intersect(self.occurrence[left], self.occurrence[right]) + mass = int(self.weights[members].sum()) if members.size else 0 + rows = self.member[members] if members.size else self.member[:0].copy() + self.occurrence[left] = np.setdiff1d( + self.occurrence[left], members, assume_unique=True + ) + self.occurrence[right] = np.setdiff1d( + self.occurrence[right], members, assume_unique=True + ) + self.frequency[left] -= mass + self.frequency[right] -= mass + token = self._append_token( + members, + self.support[left] | self.support[right], + min(self.anchor[left], self.anchor[right]), + ) + self._update_sparse_index(left, right, token, rows, members, mass) + # Keep the row mirror in step. Slot order inside a row carries no + # meaning, so the child simply takes the left parent's slot. + if members.size: + rows[rows == right] = -1 + rows[rows == left] = token + self.member[members] = rows + self.merges.append((left, right)) + return token + + +def _train_spatial_python( + codes: np.ndarray, + weights: np.ndarray, + bins: Sequence[int], + *, + vocab_size: int, + progress: bool = False, +) -> tuple[SpatialVocabulary, dict[str, object]]: + """Train Set-BPE with exact counters for observed row-local pairs. + + ``vocab_size`` is the total budget including the ``sum(bins)`` primitive + atoms. Training stops early when no positive-frequency compatible pair + remains, which happens once every record has collapsed to a single token. + + The one-time counter build costs one dense pass; later work is limited to + records changed by the committed merge. + """ + trainer = _SparseSpatialTrainer(codes, weights, bins) + if vocab_size < trainer.atom_count: + raise ValueError( + f"vocab_size {vocab_size} is below the {trainer.atom_count} primitive atoms" + ) + total_weight = int(trainer.weights.sum()) + active_width = np.sum(trainer.records >= 0, axis=1, dtype=np.int64) + initial = int(np.dot(trainer.weights, active_width)) + length = initial + trainer.build_sparse_index() + for step in range(vocab_size - trainer.atom_count): + winner = trainer.sparse_best_pair() + if winner is None: + break + left, right, frequency = winner + trainer.merge(left, right) + length -= frequency + if progress and step % 256 == 0: + print( + f" merge {step:>5} tokens/action {length / total_weight:.4f}", + flush=True, + ) + vocabulary = SpatialVocabulary( + tuple(int(value) for value in bins), tuple(trainer.merges) + ) + report = { + "records": int(trainer.count), + "weight": total_weight, + "atoms": trainer.atom_count, + "merges": len(trainer.merges), + "vocab_size": vocabulary.vocab_size, + "tokens_per_action_initial": initial / total_weight, + "tokens_per_action": length / total_weight, + "compression": initial / length, + "sparse_updates": trainer.statistics["sparse_updates"], + "sparse_heap_rebuilds": trainer.statistics["sparse_heap_rebuilds"], + "sparse_pair_types_peak": trainer.statistics["sparse_pair_types_peak"], + } + return vocabulary, report + + +def train_spatial( + codes: np.ndarray, + weights: np.ndarray, + bins: Sequence[int], + *, + vocab_size: int, + backend: str = "auto", + progress: bool = False, + **kwargs: object, +) -> tuple[SpatialVocabulary, dict[str, object]]: + """Fit spatial Set-BPE, preferring the native kernel when it is built. + + ``codes`` is one row per distinct quantized action and ``weights`` its + multiplicity (see :func:`aggregate_records`); collapsing duplicates first is + lossless because identical actions follow identical merges forever. + + ``vocab_size`` is the *total* budget including the ``sum(bins)`` primitive + atoms. Training stops early once no positive-frequency compatible pair + remains. Both backends return the identical greedy merge sequence -- the + test suite proves it -- so ``backend`` only trades speed for portability. + """ + if resolve_backend(backend) == "python": + # ``threads`` controls only the native OpenMP kernel. The profile fit API carries + # it for either backend, but the Python reference has no such argument. + kwargs.pop("threads", None) + return _train_spatial_python( + codes, weights, bins, vocab_size=vocab_size, progress=progress, **kwargs + ) + values = np.ascontiguousarray(codes, dtype=np.int32) + masses = np.ascontiguousarray(weights, dtype=np.int64) + merges, report = require_native().train_spatial( + values, + masses, + [int(count) for count in bins], + int(vocab_size), + int(kwargs.get("threads", 0)), + ) + vocabulary = SpatialVocabulary( + tuple(int(count) for count in bins), + tuple((int(left), int(right)) for left, right in merges), + ) + return vocabulary, dict(report) diff --git a/runtime/tokenization/setbpe_temporal.py b/runtime/tokenization/setbpe_temporal.py new file mode 100644 index 0000000000000000000000000000000000000000..604d45282b5706f0842ab74783f98f3b67312acd --- /dev/null +++ b/runtime/tokenization/setbpe_temporal.py @@ -0,0 +1,258 @@ +"""Runtime rooted-temporal Set-BPE vocabulary and lossless serialization.""" + +from __future__ import annotations + +import heapq +from collections.abc import Sequence +from dataclasses import dataclass + +import numpy as np + +from ..common._validation import active_dimensions +from .setbpe_spatial import SpatialVocabulary + +__all__ = ["TemporalVocabulary"] + +Rule = tuple[int, int, int] +Pattern = tuple[tuple[int, int], ...] + + +def _active_support(active_dims: Sequence[int] | None, dimension: int) -> int: + dims = active_dimensions(active_dims, dimension) + return sum(1 << value for value in dims) + + +def _validate_episode(values: np.ndarray, spatial: SpatialVocabulary) -> int: + if values.ndim != 2 or values.shape[1] != spatial.dimension or len(values) < 1: + raise ValueError( + f"every episode must have non-empty shape [T, {spatial.dimension}], got {values.shape}" + ) + counts = np.asarray(spatial.bins, dtype=np.int64) + if np.any(values < -1) or np.any(values >= counts[None, :]): + raise ValueError("episodes must contain -1 or an in-range cell index") + support = values >= 0 + if np.any(np.sum(support, axis=1) == 0): + raise ValueError("every timestep needs at least one active coordinate") + if np.any(support != support[:1]): + raise ValueError("active dimensions must remain constant within an episode") + return sum(1 << int(value) for value in np.flatnonzero(support[0])) + + +@dataclass +class TemporalVocabulary: + """A spatial vocabulary plus rooted, forward-only temporal merge rules.""" + + spatial: SpatialVocabulary + lookahead: int + merges: tuple[Rule, ...] + + def __post_init__(self) -> None: + if self.lookahead < 1: + raise ValueError("lookahead must be positive") + self.patterns: list[Pattern] = [ + ((0, token),) for token in range(self.spatial.vocab_size) + ] + self.spans: list[int] = [1] * self.spatial.vocab_size + self.root: list[int] = list(range(self.spatial.vocab_size)) + self.support_masks: list[tuple[int, ...]] = [ + (int(self.spatial.support[token]),) + for token in range(self.spatial.vocab_size) + ] + self._rules: dict[Rule, tuple[int, int]] = {} + for rank, rule in enumerate(self.merges): + left, right, offset = (int(value) for value in rule) + child = self.spatial.vocab_size + rank + if left >= child or right >= child or left < 0 or right < 0: + raise ValueError( + "a temporal merge must reference already-created tokens" + ) + pattern, masks = _merge_pattern( + self.patterns[left], + self.patterns[right], + offset, + self.lookahead, + self.spatial, + ) + self.patterns.append(pattern) + self.spans.append(len(masks)) + self.root.append(self.root[left]) + self.support_masks.append(masks) + key = (left, right, offset) + if key in self._rules: + raise ValueError("duplicate temporal merge rule") + self._rules[key] = (child, rank) + + @property + def vocab_size(self) -> int: + return self.spatial.vocab_size + len(self.merges) + + def encode(self, codes: np.ndarray) -> list[int]: + """Encode a dense episode; ``-1`` slots are semantically absent.""" + values = np.asarray(codes, dtype=np.int64) + _validate_episode(values, self.spatial) + + at: list[dict[int, int]] = [dict() for _ in range(values.shape[0])] + tokens: list[int] = [] + anchors: list[int] = [] + active: list[bool] = [] + for anchor, row in enumerate(values): + for token in self.spatial.encode(row): + instance = len(tokens) + tokens.append(token) + anchors.append(anchor) + active.append(True) + at[anchor][token] = instance + + frontier: list[tuple[int, int, int, int, int, int]] = [] + + def offer(left_instance: int, right_instance: int) -> None: + left = tokens[left_instance] + right = tokens[right_instance] + offset = anchors[right_instance] - anchors[left_instance] + rule = self._rules.get((left, right, offset)) + if rule is not None: + child, rank = rule + heapq.heappush( + frontier, + ( + rank, + anchors[left_instance], + int(self.spatial.anchor[self.root[child]]), + left_instance, + right_instance, + child, + ), + ) + + for anchor, left_items in enumerate(at): + for offset in range(1, min(self.lookahead, len(at) - anchor)): + for left_instance in left_items.values(): + for right_instance in at[anchor + offset].values(): + offer(left_instance, right_instance) + + while frontier: + _, anchor, _, left_instance, right_instance, child = heapq.heappop(frontier) + if not active[left_instance] or not active[right_instance]: + continue + if anchors[left_instance] != anchor: + continue + rule = ( + tokens[left_instance], + tokens[right_instance], + anchors[right_instance] - anchor, + ) + resolved = self._rules.get(rule) + if resolved is None or resolved[0] != child: + continue + + active[left_instance] = False + active[right_instance] = False + del at[anchor][tokens[left_instance]] + del at[anchors[right_instance]][tokens[right_instance]] + new_instance = len(tokens) + tokens.append(child) + anchors.append(anchor) + active.append(True) + at[anchor][child] = new_instance + + for offset in range(1, self.lookahead): + future = anchor + offset + if future < len(at): + for other in at[future].values(): + offer(new_instance, other) + past = anchor - offset + if past >= 0: + for other in at[past].values(): + offer(other, new_instance) + + ordered: list[tuple[int, int, int]] = [] + for anchor, instances in enumerate(at): + for token, instance in instances.items(): + if active[instance]: + ordered.append( + (anchor, int(self.spatial.anchor[self.root[token]]), token) + ) + ordered.sort() + return [token for _, _, token in ordered] + + def decode( + self, + stream: Sequence[int], + *, + active_dims: Sequence[int] | None = None, + ) -> np.ndarray: + """Decode ids using the route's requested active-dimension support.""" + if not stream: + return np.empty((0, self.spatial.dimension), dtype=np.int64) + full = _active_support(active_dims, self.spatial.dimension) + coverage: list[int] = [0] + spatial_tokens: list[list[int]] = [[]] + front = 0 + for raw_token in stream: + token = int(raw_token) + if not 0 <= token < self.vocab_size: + raise ValueError(f"temporal token {token} is outside the vocabulary") + while front < len(coverage) and coverage[front] == full: + front += 1 + if front == len(coverage): + coverage.append(0) + spatial_tokens.append([]) + + uncovered = full & ~coverage[front] + expected_anchor = int(uncovered & -uncovered).bit_length() - 1 + actual_anchor = int(self.spatial.anchor[self.root[token]]) + if actual_anchor != expected_anchor: + raise ValueError( + "temporal stream violates root order at the earliest incomplete step" + ) + + for offset, spatial_token in self.patterns[token]: + step = front + offset + while step >= len(coverage): + coverage.append(0) + spatial_tokens.append([]) + support = int(self.spatial.support[spatial_token]) + if coverage[step] & support: + raise ValueError( + "temporal stream contains overlapping spatial support" + ) + coverage[step] |= support + spatial_tokens[step].append(spatial_token) + + if any(mask != full for mask in coverage): + raise ValueError("temporal stream ended with an incomplete control step") + return np.stack( + [ + self.spatial.decode(tokens, active_dims=active_dims) + for tokens in spatial_tokens + ] + ) + + +def _merge_pattern( + left: Pattern, + right: Pattern, + offset: int, + lookahead: int, + spatial: SpatialVocabulary, +) -> tuple[Pattern, tuple[int, ...]]: + if not 1 <= offset < lookahead: + raise ValueError("a temporal merge offset must be in [1, lookahead)") + shifted = tuple((time + offset, token) for time, token in right) + pattern = tuple(sorted((*left, *shifted))) + span = 1 + max(time for time, _ in pattern) + if span > lookahead: + raise ValueError("temporal merge exceeds the lookahead horizon") + masks = [0] * span + roots = 0 + for time, token in pattern: + support = int(spatial.support[token]) + if masks[time] & support: + raise ValueError("temporal merge joins overlapping spatial supports") + masks[time] |= support + roots += int(time == 0) + if roots != 1: + raise ValueError( + "a rooted temporal token must have exactly one time-zero spatial token" + ) + return pattern, tuple(masks) diff --git a/runtime/tokenization/setbpe_temporal_training.py b/runtime/tokenization/setbpe_temporal_training.py new file mode 100644 index 0000000000000000000000000000000000000000..704b1bc4020fb0c739d174ffd04c3fa308697d48 --- /dev/null +++ b/runtime/tokenization/setbpe_temporal_training.py @@ -0,0 +1,666 @@ +"""Temporal Set-BPE fitting implementations. + +Runtime vocabulary and serialization live in ``temporal``; this module is +loaded only by training entry points. +""" + +from __future__ import annotations + +import heapq +from array import array +from collections import Counter +from collections.abc import Iterable, Iterator, Sequence + +import numpy as np + +from ..backend.native_backend import require_native, resolve_backend +from .setbpe_spatial import SpatialVocabulary +from .setbpe_temporal import ( + Pattern, + Rule, + TemporalVocabulary, + _merge_pattern, + _validate_episode, +) + +__all__ = ["train_temporal"] + + +class _SparseTemporalTrainer: + """Mutable sparse local graph with exact shifted self-rule certification.""" + + def __init__( + self, + episodes: Sequence[np.ndarray], + spatial: SpatialVocabulary, + lookahead: int, + max_vocab_size: int, + ) -> None: + if lookahead < 1: + raise ValueError("lookahead must be positive") + self.spatial = spatial + self.lookahead = lookahead + self.max_vocab_size = max_vocab_size + self.patterns: list[Pattern] = [ + ((0, token),) for token in range(spatial.vocab_size) + ] + self.spans: list[int] = [1] * spatial.vocab_size + self.masks: list[tuple[int, ...]] = [ + (int(spatial.support[token]),) for token in range(spatial.vocab_size) + ] + self._mask_table = np.zeros((max_vocab_size, lookahead), dtype=np.uint64) + self._mask_table[: spatial.vocab_size, 0] = spatial.support.astype( + np.uint64, copy=False + ) + self.merges: list[Rule] = [] + self.tokens = array("I") + self.anchors = array("I") + self.slots = array("I") + total_anchors = sum(len(episode) + lookahead for episode in episodes) + self.member = np.full((total_anchors, spatial.dimension), -1, dtype=np.int32) + self.occupied = np.zeros(total_anchors, dtype=np.uint64) + self._slot_mask = (1 << spatial.dimension) - 1 + self.occurrences: list[dict[int, int]] = [ + dict() for _ in range(spatial.vocab_size) + ] + self.episode_count = len(episodes) + self.action_count = sum(len(episode) for episode in episodes) + self.initial_count = 0 + self.counts: dict[int, int] = {} + self.self_versions: dict[int, int] = {} + self.heap: list[tuple[int, int, int, int, int, bool]] = [] + self.self_cache: dict[int, tuple[int, list[tuple[int, int]]]] = {} + self._legal_table = np.zeros( + (lookahead, max_vocab_size, max_vocab_size), dtype=np.uint8 + ) + for offset in range(1, lookahead): + self._legal_table[offset, : spatial.vocab_size, : spatial.vocab_size] = 2 + self.statistics = Counter() + self._active_count = 0 + + self._load_episodes(episodes) + self.counts = self._initial_pair_counts() + self.heap = self._initial_heap() + self.statistics["initial_edge_types"] = len(self.counts) + self.statistics["initial_edges"] = sum(self.counts.values()) + + def _load_episodes(self, episodes: Sequence[np.ndarray]) -> None: + """Populate sparse token instances from spatially encoded episodes.""" + + base = 0 + encoding_cache: dict[bytes, tuple[int, ...]] = {} + for episode in episodes: + values = np.asarray(episode, dtype=np.int64) + _validate_episode(values, self.spatial) + needed = base + len(values) + self.lookahead + for time, row in enumerate(values): + cache_key = row.tobytes() + encoded = encoding_cache.get(cache_key) + if encoded is None: + encoded = tuple(self.spatial.encode(row)) + encoding_cache[cache_key] = encoded + anchor = base + time + for token in encoded: + self._new_instance(token, anchor) + self.initial_count += 1 + base = needed + + def _initial_pair_counts(self) -> dict[int, int]: + """Count all initially adjacent temporal token pairs.""" + base_vocab = self.spatial.vocab_size + dense_size = base_vocab * base_vocab * self.lookahead + if dense_size <= 1 << 25: + dense_counts = np.zeros(dense_size, dtype=np.int64) + for offset in range(1, self.lookahead): + left_rows = self.member[:-offset] + right_rows = self.member[offset:] + for left_slot in range(self.spatial.dimension): + left = left_rows[:, left_slot] + for right_slot in range(self.spatial.dimension): + right = right_rows[:, right_slot] + valid = (left >= 0) & (right >= 0) + if not valid.any(): + continue + packed = ( + left[valid].astype(np.int64) * base_vocab + right[valid] + ) * self.lookahead + offset + np.add.at(dense_counts, packed, 1) + initial_counts: dict[int, int] = {} + for raw_key in np.flatnonzero(dense_counts): + pair, offset = divmod(int(raw_key), self.lookahead) + left, right = divmod(pair, base_vocab) + packed_key = ( + left * self.max_vocab_size + right + ) * self.lookahead + offset + initial_counts[packed_key] = int(dense_counts[raw_key]) + return initial_counts + return self._initial_pair_counts_sparse() + + def _initial_pair_counts_sparse(self) -> dict[int, int]: + """Count initial pairs without allocating a dense pair table.""" + fallback: Counter[int] = Counter() + for anchor, left_row in enumerate(self.member): + left_tokens = left_row[left_row >= 0] + for offset in range(1, self.lookahead): + target = anchor + offset + if target >= len(self.member): + break + right_tokens = self.member[target] + right_tokens = right_tokens[right_tokens >= 0] + for raw_left in left_tokens: + for raw_right in right_tokens: + packed_key = ( + int(raw_left) * self.max_vocab_size + int(raw_right) + ) * self.lookahead + offset + fallback[packed_key] += 1 + return dict(fallback) + + def _initial_heap(self) -> list[tuple[int, int, int, int, int, bool]]: + heap: list[tuple[int, int, int, int, int, bool]] = [] + for packed_key, count in self.counts.items(): + pair, offset = divmod(packed_key, self.lookahead) + left, right = divmod(pair, self.max_vocab_size) + heap.append((-count, left, right, offset, 0, False)) + heapq.heapify(heap) + return heap + + def _new_instance(self, token: int, anchor: int) -> int: + while token >= len(self.occurrences): + self.occurrences.append({}) + if anchor in self.occurrences[token]: + raise RuntimeError("a temporal type cannot occur twice at one anchor") + occupied = int(self.occupied[anchor]) + available = self._slot_mask ^ occupied + if available == 0: + raise RuntimeError("temporal anchor exceeded the spatial token-width bound") + lowest = available & -available + slot = lowest.bit_length() - 1 + self.occupied[anchor] = occupied | lowest + instance = len(self.tokens) + self.tokens.append(token) + self.anchors.append(anchor) + self.slots.append(slot) + self.member[anchor, slot] = token + self._active_count += 1 + self.occurrences[token][anchor] = instance + return instance + + def _initialize_legality(self, child: int) -> None: + """Fill all legality entries involving one newly created token.""" + count = child + 1 + spans = np.asarray(self.spans[:count], dtype=np.int64) + masks = self._mask_table[:count] + child_masks = self._mask_table[child] + child_span = self.spans[child] + for offset in range(1, self.lookahead): + overlap = np.zeros(count, dtype=bool) + for relative in range(self.lookahead - offset): + overlap |= (child_masks[offset + relative] & masks[:, relative]) != 0 + legal = (offset + spans <= self.lookahead) & ~overlap + self._legal_table[offset, child, :count] = np.where(legal, 2, 1) + + overlap.fill(False) + for relative in range(min(child_span, self.lookahead - offset)): + overlap |= (masks[:, offset + relative] & child_masks[relative]) != 0 + legal = (offset + child_span <= self.lookahead) & ~overlap + self._legal_table[offset, :count, child] = np.where(legal, 2, 1) + + def _entry( + self, packed_key: int, priority: int, certified: bool + ) -> tuple[int, int, int, int, int, bool]: + pair, offset = divmod(packed_key, self.lookahead) + left, right = divmod(pair, self.max_vocab_size) + return ( + -priority, + left, + right, + offset, + self.self_versions.get(packed_key, 0) if left == right else 0, + certified, + ) + + def _select(self, key: Rule) -> list[tuple[int, int]]: + left, right, offset = key + left_occurrences = self.occurrences[left] + right_occurrences = self.occurrences[right] + if left != right: + return [ + (left_instance, right_occurrences[anchor + offset]) + for anchor, left_instance in left_occurrences.items() + if anchor + offset in right_occurrences + ] + selected: list[tuple[int, int]] = [] + consumed: set[int] = set() + for anchor in sorted(left_occurrences): + target = anchor + offset + if ( + anchor in consumed + or target in consumed + or target not in right_occurrences + ): + continue + selected.append((left_occurrences[anchor], right_occurrences[target])) + consumed.add(anchor) + consumed.add(target) + return selected + + def best(self) -> tuple[Rule, list[tuple[int, int]]] | None: + while self.heap: + negative, left, right, offset, version, certified = heapq.heappop(self.heap) + self.statistics["heap_pops"] += 1 + key = (left, right, offset) + packed_key = (left * self.max_vocab_size + right) * self.lookahead + offset + current_count = self.counts.get(packed_key, 0) + if current_count <= 0: + continue + if left == right: + current_version = self.self_versions.get(packed_key, 0) + stale = current_version != version + else: + stale = current_count != -negative + if stale: + heapq.heappush(self.heap, self._entry(packed_key, current_count, False)) + continue + if left == right and not certified: + matches = self._select(key) + self.self_cache[packed_key] = (version, matches) + heapq.heappush(self.heap, self._entry(packed_key, len(matches), True)) + self.statistics["self_certifications"] += 1 + continue + if certified: + cached = self.self_cache.get(packed_key) + if cached is None or cached[0] != version: + continue + matches = cached[1] + else: + matches = self._select(key) + gain = -negative + if len(matches) != gain: + raise RuntimeError( + "sparse temporal counter disagrees with exact matching" + ) + return key, matches + return None + + def _accumulate_incident( + self, + instances: Sequence[int] | set[int], + multiplier: int, + deltas: Counter[int], + ) -> None: + """Vectorize all local edges incident to a changed instance set.""" + grouped: dict[int, list[int]] = {} + for instance in instances: + grouped.setdefault(self.tokens[instance], []).append(self.anchors[instance]) + width = len(self.patterns) + total_anchors = self.member.shape[0] + vocab_width = self.max_vocab_size + lookahead = self.lookahead + for token, raw_anchors in grouped.items(): + anchors = np.asarray(raw_anchors, dtype=np.int64) + for offset in range(1, self.lookahead): + future_anchors = anchors[anchors + offset < total_anchors] + offset + if future_anchors.size: + values = self.member[future_anchors].reshape(-1) + values = values[values >= 0] + if values.size: + counts = np.bincount(values, minlength=width) + neighbors = np.flatnonzero(counts) + neighbors = neighbors[ + self._legal_table[offset, token, neighbors] == 2 + ] + for raw_neighbor in neighbors: + neighbor = int(raw_neighbor) + packed_key = ( + token * vocab_width + neighbor + ) * lookahead + offset + deltas[packed_key] += multiplier * int(counts[neighbor]) + past_anchors = anchors[anchors >= offset] - offset + if past_anchors.size: + values = self.member[past_anchors].reshape(-1) + values = values[values >= 0] + if values.size: + counts = np.bincount(values, minlength=width) + neighbors = np.flatnonzero(counts) + neighbors = neighbors[ + self._legal_table[offset, neighbors, token] == 2 + ] + for raw_neighbor in neighbors: + neighbor = int(raw_neighbor) + packed_key = ( + neighbor * vocab_width + token + ) * lookahead + offset + deltas[packed_key] += multiplier * int(counts[neighbor]) + + def _correct_internal_edges( + self, + instances: Sequence[int] | set[int], + correction: int, + deltas: Counter[int], + ) -> None: + """Incident aggregation sees changed-to-changed edges from both ends.""" + by_anchor: dict[int, list[int]] = {} + for instance in instances: + by_anchor.setdefault(self.anchors[instance], []).append( + self.tokens[instance] + ) + vocab_width = self.max_vocab_size + lookahead = self.lookahead + for anchor, left_tokens in by_anchor.items(): + for offset in range(1, self.lookahead): + right_tokens = by_anchor.get(anchor + offset) + if right_tokens is None: + continue + for left in left_tokens: + for right in right_tokens: + if self._legal_table[offset, left, right] == 2: + packed_key = ( + left * vocab_width + right + ) * lookahead + offset + deltas[packed_key] += correction + + def _apply_deltas(self, deltas: Counter[int], selected: Rule) -> None: + """Update counts while deferring heap repair for decreasing keys. + + Edges added by a commit all contain its newly allocated child token. + Therefore a pre-existing key can only decrease (or keep the same + count). Its old heap priority remains an admissible upper bound and is + repaired lazily when popped. New keys and the already-popped selected + key need fresh entries immediately. + """ + counts = self.counts + self_versions = self.self_versions + heap = self.heap + updates = 0 + vocab_width = self.max_vocab_size + lookahead = self.lookahead + selected_key = (selected[0] * vocab_width + selected[1]) * lookahead + selected[ + 2 + ] + for packed_key, delta in deltas.items(): + previous = counts.get(packed_key, 0) + count = previous + delta + if count < 0: + raise RuntimeError("temporal edge counter became negative") + # Only shifted self-rules need generations: for distinct token + # types raw edge count equals matching cardinality, so an + # unchanged count leaves an old heap entry exact. + pair = packed_key // lookahead + is_self = pair % (vocab_width + 1) == 0 + if is_self: + version = self_versions.get(packed_key, 0) + 1 + self_versions[packed_key] = version + self.self_cache.pop(packed_key, None) + else: + version = 0 + if count: + counts[packed_key] = count + if previous == 0 or packed_key == selected_key: + pair, offset = divmod(packed_key, lookahead) + left, right = divmod(pair, vocab_width) + heapq.heappush(heap, (-count, left, right, offset, version, False)) + else: + counts.pop(packed_key, None) + updates += abs(delta) + self.statistics["counter_updates"] += updates + + def commit(self, key: Rule, matches: Sequence[tuple[int, int]]) -> None: + left, right, offset = key + pattern, masks = _merge_pattern( + self.patterns[left], + self.patterns[right], + offset, + self.lookahead, + self.spatial, + ) + child = len(self.patterns) + self.patterns.append(pattern) + self.spans.append(len(masks)) + self.masks.append(masks) + self.occurrences.append({}) + self.merges.append(key) + self._mask_table[child, : len(masks)] = masks + self._initialize_legality(child) + + removed = {instance for pair in matches for instance in pair} + if len(removed) != 2 * len(matches): + raise RuntimeError("selected temporal matches overlap") + deltas: Counter[int] = Counter() + self._accumulate_incident(removed, -1, deltas) + self._correct_internal_edges(removed, 1, deltas) + + for instance in removed: + token = self.tokens[instance] + anchor = self.anchors[instance] + self.member[anchor, self.slots[instance]] = -1 + self._active_count -= 1 + self.occupied[anchor] = int(self.occupied[anchor]) & ~( + 1 << self.slots[instance] + ) + del self.occurrences[token][anchor] + + children = [ + self._new_instance(child, self.anchors[left_instance]) + for left_instance, _ in matches + ] + self._accumulate_incident(children, 1, deltas) + self._correct_internal_edges(children, -1, deltas) + self._apply_deltas(deltas, key) + + @property + def active_count(self) -> int: + return self._active_count + + +def _prepare_episodes( + episodes: Iterable[np.ndarray], spatial: SpatialVocabulary +) -> Sequence[np.ndarray]: + if getattr(episodes, "_actioncodec2_validated", False): + if not isinstance(episodes, Sequence) or len(episodes) == 0: + raise ValueError("validated temporal corpus must be a non-empty sequence") + return episodes + prepared: list[np.ndarray] = [] + for episode in episodes: + values = np.asarray(episode) + if values.dtype.kind not in {"i", "u"}: + raise ValueError("episodes must contain integer cell indices") + _validate_episode(values, spatial) + prepared.append(np.ascontiguousarray(values)) + if not prepared or sum(len(episode) for episode in prepared) == 0: + raise ValueError("temporal training requires at least one non-empty episode") + return prepared + + +def _train_temporal_python( + episodes: Iterable[np.ndarray], + spatial: SpatialVocabulary, + *, + vocab_size: int, + lookahead: int, + progress: bool = False, +) -> tuple[TemporalVocabulary, dict[str, object]]: + """Fit exact rooted temporal BPE with sparse observed-edge counters.""" + prepared = _prepare_episodes(episodes, spatial) + if vocab_size < spatial.vocab_size: + raise ValueError( + f"vocab_size {vocab_size} is below the {spatial.vocab_size} spatial tokens" + ) + trainer = _SparseTemporalTrainer(prepared, spatial, lookahead, vocab_size) + length = trainer.initial_count + for step in range(vocab_size - spatial.vocab_size): + winner = trainer.best() + if winner is None: + break + key, matches = winner + if not matches: + break + trainer.commit(key, matches) + length -= len(matches) + if progress and step % 64 == 0: + print( + f" temporal merge {step:>5} tokens/action {length / trainer.action_count:.4f}", + flush=True, + ) + if length != trainer.active_count: + raise RuntimeError("temporal token accounting disagrees with active tiling") + vocabulary = TemporalVocabulary(spatial, lookahead, tuple(trainer.merges)) + report: dict[str, object] = { + "episodes": trainer.episode_count, + "actions": trainer.action_count, + "lookahead": lookahead, + "spatial_vocab_size": spatial.vocab_size, + "merges": len(trainer.merges), + "vocab_size": vocabulary.vocab_size, + "tokens_per_action_initial": trainer.initial_count / trainer.action_count, + "tokens_per_action": length / trainer.action_count, + "compression": trainer.initial_count / length, + "trainer": "sparse_exact", + **{key: int(value) for key, value in trainer.statistics.items()}, + } + return vocabulary, report + + +def _flatten_instances_native( + episodes: Sequence[np.ndarray], + spatial: SpatialVocabulary, + lookahead: int, + threads: int, +) -> tuple[Iterator[tuple[np.ndarray, np.ndarray]], int, int]: + """Spatially encode bounded chunks without a corpus-sized final copy.""" + if getattr(episodes, "_actioncodec2_validated", False): + return _flatten_cell_corpus(episodes, spatial, lookahead, threads) + actions = sum(len(episode) for episode in episodes) + total_anchors = actions + len(episodes) * lookahead + max_batch_actions = 1_000_000 + merges = np.asarray(spatial.merges, dtype=np.int32).reshape(-1, 2) + + def generate() -> Iterator[tuple[np.ndarray, np.ndarray]]: + pending: list[np.ndarray] = [] + pending_anchors: list[np.ndarray] = [] + pending_actions = 0 + base = 0 + + def encode() -> tuple[np.ndarray, np.ndarray]: + codes = np.concatenate(pending, axis=0).astype(np.int32, copy=False) + row_anchors = np.concatenate(pending_anchors) + tokens, anchors = require_native().spatial_encode_instances( + np.ascontiguousarray(codes), + np.ascontiguousarray(row_anchors), + [int(count) for count in spatial.bins], + np.ascontiguousarray(merges), + int(threads), + ) + return np.asarray(tokens, dtype=np.int32), np.asarray( + anchors, dtype=np.int64 + ) + + for episode in episodes: + length = len(episode) + pending.append(episode) + pending_anchors.append(np.arange(base, base + length, dtype=np.int64)) + pending_actions += length + base += length + lookahead + if pending_actions >= max_batch_actions: + yield encode() + pending.clear() + pending_anchors.clear() + pending_actions = 0 + if pending: + yield encode() + + return generate(), total_anchors, actions + + +def _flatten_cell_corpus( + episodes: Sequence[np.ndarray], + spatial: SpatialVocabulary, + lookahead: int, + threads: int, +) -> tuple[Iterator[tuple[np.ndarray, np.ndarray]], int, int]: + """Encode contiguous quantizer chunks while vectorizing episode gaps.""" + actions = sum(len(codes) for codes in episodes.chunks if codes is not None) + total_anchors = actions + len(episodes) * lookahead + merges = np.asarray(spatial.merges, dtype=np.int32).reshape(-1, 2) + chunks = episodes.chunks + offsets_by_chunk = episodes.offsets + + def generate() -> Iterator[tuple[np.ndarray, np.ndarray]]: + base = 0 + for chunk_index, (codes, offsets) in enumerate( + zip(chunks, offsets_by_chunk, strict=True) + ): + if codes is None: + raise RuntimeError("cell corpus chunk has already been consumed") + lengths = np.diff(offsets) + row_anchors = ( + base + + np.arange(len(codes), dtype=np.int64) + + np.repeat( + np.arange(len(lengths), dtype=np.int64) * lookahead, lengths + ) + ) + tokens, anchors = require_native().spatial_encode_instances( + np.ascontiguousarray(codes), + row_anchors, + [int(count) for count in spatial.bins], + np.ascontiguousarray(merges), + int(threads), + ) + episodes.release_chunk(chunk_index) + yield ( + np.asarray(tokens, dtype=np.int32), + np.asarray(anchors, dtype=np.int64), + ) + base += len(codes) + len(lengths) * lookahead + + return generate(), total_anchors, actions + + +def train_temporal( + episodes: Iterable[np.ndarray], + spatial: SpatialVocabulary, + *, + vocab_size: int, + lookahead: int, + backend: str = "auto", + progress: bool = False, + threads: int = 0, +) -> tuple[TemporalVocabulary, dict[str, object]]: + """Fit rooted temporal Set-BPE on top of a frozen spatial vocabulary. + + ``lookahead`` bounds how many control steps one temporal symbol may span. + Both backends return the identical greedy merge sequence; the native one is + preferred when built. + """ + prepared = _prepare_episodes(episodes, spatial) + if resolve_backend(backend) == "python": + return _train_temporal_python( + prepared, + spatial, + vocab_size=vocab_size, + lookahead=lookahead, + progress=progress, + ) + episode_count = len(prepared) + instance_chunks, total_anchors, actions = _flatten_instances_native( + prepared, spatial, lookahead, threads + ) + merges, report = require_native().train_temporal( + instance_chunks, + int(total_anchors), + np.ascontiguousarray(spatial.support, dtype=np.uint64), + np.ascontiguousarray(spatial.anchor, dtype=np.int32), + int(spatial.dimension), + int(lookahead), + int(vocab_size), + int(actions), + int(threads), + ) + vocabulary = TemporalVocabulary( + spatial, + int(lookahead), + tuple(tuple(int(value) for value in rule) for rule in merges), + ) + payload = dict(report) + payload["episodes"] = episode_count + return vocabulary, payload diff --git a/runtime/training/__init__.py b/runtime/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1fa83db71f8531e6f9ee6d53c9a026cce3f96508 --- /dev/null +++ b/runtime/training/__init__.py @@ -0,0 +1 @@ +"""Internal training components for ActionCodec2.""" diff --git a/runtime/training/fitting.py b/runtime/training/fitting.py new file mode 100644 index 0000000000000000000000000000000000000000..eaad5044ffc5164580da3768a5175cffb3d3e36b --- /dev/null +++ b/runtime/training/fitting.py @@ -0,0 +1,188 @@ +"""Replayable, bounded source staging for the public ActionCodec2.fit interface.""" + +from __future__ import annotations + +from collections.abc import Mapping +from contextlib import ExitStack +from pathlib import Path +from tempfile import TemporaryDirectory + +import numpy as np +import torch + +from ..physical.boundary import bootstrap_second_order_trajectory +from ..routing.routing import ActionCodec2Router +from ..frontend.tokenizer import _numpy +from .training import ProfileTrainer +from ..physical.trajectory import ActiveTrajectory + + +class TrajectorySpool: + """Stage canonical episodes in two sequential files, without retaining source data. + + Args: + directory: Task-owned scratch directory, removed by the caller. + """ + + def __init__(self, directory): + self.directory = Path(directory) + self.counts = {} + self._stack = ExitStack() + self._writers = {} + + def __enter__(self): + return self + + def __exit__(self, *exc): + return self._stack.__exit__(*exc) + + def append(self, profile, trajectory): + if profile not in self._writers: + self._writers[profile] = self._stack.enter_context( + (self.directory / f"{profile}.npy").open("wb") + ) + stream = self._writers[profile] + for value in ( + trajectory.values, + trajectory.active_dims, + [] + if trajectory.previous_first_order is None + else trajectory.previous_first_order, + ): + np.save(stream, np.asarray(value), allow_pickle=False) + self.counts[profile] = self.counts.get(profile, 0) + 1 + + def trajectories(self, profile): + self._writers[profile].flush() + with (self.directory / f"{profile}.npy").open("rb") as stream: + for _ in range(self.counts[profile]): + values = np.load(stream, allow_pickle=False) + dims = tuple(int(d) for d in np.load(stream, allow_pickle=False)) + previous = np.load(stream, allow_pickle=False) + yield ActiveTrajectory( + values, + dims, + previous_first_order=previous if previous.size else None, + ) + + +class ActionCodec2Fitter: + """Route user layouts and fit replacement profiles without mutating the caller. + + Args: + codec: Public processor supplying action-space definitions and profile configs. + bootstrap_context: Explicitly consume two canonical context frames for second order. + """ + + def __init__(self, codec, *, bootstrap_context=False): + if not isinstance(bootstrap_context, bool): + raise TypeError("bootstrap_context must be a boolean") + self.codec = codec + self.router = codec.action_tokenizer + self.bootstrap_context = bootstrap_context + + def _route(self, record): + unknown = set(record) - { + "actions", + "action_space", + "fps", + "current_state", + "previous_first_order", + } + if unknown: + raise ValueError(f"unknown training record fields: {sorted(unknown)}") + if "actions" not in record: + raise ValueError("training record requires an 'actions' array") + preset = self.router.preset( + record.get("action_space", self.codec.action_space.name) + ) + values = _numpy(torch.as_tensor(record["actions"])) + if values.ndim == 2: + values = values[None] + if values.ndim != 3 or values.shape[0] == 0 or values.shape[1] == 0: + raise ValueError( + "training actions must have nonempty shape (T,D) or (B,T,D)" + ) + current = _numpy(record.get("current_state")) + if preset.input == "absolute" and current is None: + raise ValueError( + f"absolute action space {preset.name!r} requires current_state(B,D)" + ) + if preset.input == "delta" and current is not None: + raise ValueError(f"delta action space {preset.name!r}: omit current_state") + previous = self.codec._numpy_boundary(record.get("previous_first_order")) + kwargs = dict( + preset=preset.name, + fps=self.codec._source_fps(record.get("fps"), preset), + current_state=current, + ) + if not self.bootstrap_context: + return self.router.route_trajectories( + values, previous_first_order=previous, **kwargs + ) + if previous is not None: + raise ValueError( + "bootstrap_context and previous_first_order are mutually exclusive" + ) + canonical = preset.canonicalize( + values, source_hz=kwargs["fps"], current_state=current + ) + routed = {} + for component in preset.components: + processor = self.router.profiles[component.space].processor + trajectories = canonical[component.name] + if processor.requires_previous_first_order: + trajectories = [ + bootstrap_second_order_trajectory(t, processor.config) + for t in trajectories + ] + routed.setdefault(component.space, []).extend(trajectories) + return routed + + def fit(self, episodes, *, work_dir=None, metadata=None, **options): + """Fit observed profiles from arrays or records, replaying staged episodes per profile.""" + metadata = {k: v for k, v in (metadata or {}).items() if v is not None} + if isinstance(episodes, (np.ndarray, torch.Tensor, Mapping)): + episodes = (episodes,) + reports = {} + processors = { + name: runtime.processor for name, runtime in self.router.profiles.items() + } + records = 0 + with TemporaryDirectory(prefix="actioncodec2-fit-", dir=work_dir) as directory: + with TrajectorySpool(directory) as spool: + for episode in episodes: + record = ( + dict(episode) + if isinstance(episode, Mapping) + else {"actions": episode} + ) + duplicates = set(record).intersection(metadata) + if duplicates: + raise ValueError( + f"training arguments provided twice: {sorted(duplicates)}" + ) + record.update(metadata) + for profile, trajectories in self._route(record).items(): + for trajectory in trajectories: + spool.append(profile, trajectory) + records += 1 + if not spool.counts: + raise ValueError("fit requires at least one nonempty episode") + for name in spool.counts: + processors[name], reports[name] = ProfileTrainer( + processors[name].config + ).fit(spool.trajectories(name), **options) + router = ActionCodec2Router.from_processors( + processors, + token_budgets={ + name: runtime.spec.token_budget + for name, runtime in self.router.profiles.items() + }, + presets=list(self.router.action_spaces.values()), + ) + return router, { + "records": records, + "bootstrap_context": self.bootstrap_context, + "profiles": reports, + } diff --git a/runtime/training/training.py b/runtime/training/training.py new file mode 100644 index 0000000000000000000000000000000000000000..e45a17e492746c8ec8e6eac91ea2c30d6deaaf84 --- /dev/null +++ b/runtime/training/training.py @@ -0,0 +1,756 @@ +"""Whole-trajectory primitive construction and BPE fitting. + +A training sample is always one complete episode/trajectory. Physical error +feedback is reset only between trajectories, and the tokenizer receives the +same trajectory boundaries. Chunking belongs exclusively to downstream policy +sampling and is intentionally absent from this module. + +``config.set_bpe.scheme`` selects the tokenizer fitted on those cells: rooted +spatial-temporal Set-BPE, or plain BPE over the row-major flattened stream. +""" + +from __future__ import annotations + +import math +import time +from bisect import bisect_right +from collections import defaultdict +from collections.abc import Iterable, Sequence + +import numpy as np + +from ..common._validation import broadcast_vector +from ..tokenization.bpe import BPESerializer, SetBPESerializer +from ..tokenization.bpe_training import train_bpe +from ..physical.configuration import ActionCodec2Config, ConfigError +from ..physical.physical_codec import PhysicalCodec +from ..physical.physical_processor import ActionCodec2Processor +from ..physical.primitives import PrimitiveStore +from ..tokenization.setbpe_spatial import aggregate_records +from ..tokenization.setbpe_spatial_training import train_spatial +from ..tokenization.setbpe_temporal_training import train_temporal +from ..physical.trajectory import ( + TrajectoryLike, + prepare_trajectory, + trajectory_boundary, +) + + +def _zero_anchored_grid( + lower: np.float32, upper: np.float32, spacing: np.float32 +) -> np.ndarray: + """Build a float32 endpoint-preserving grid with an exact zero.""" + below = int(np.ceil(-lower / spacing)) + above = int(np.ceil(upper / spacing)) + negative = np.linspace(lower, np.float32(0.0), below + 1, dtype=np.float32)[:-1] + positive = np.linspace(np.float32(0.0), upper, above + 1, dtype=np.float32) + return np.concatenate((negative, positive)) + + +def _additive_axis( + lower: np.float32, + upper: np.float32, + distortion_limit: np.float32, + *, + zero_primitive: bool, +) -> np.ndarray: + """Construct an endpoint-preserving additive cover at the distortion limit.""" + width = np.float32(upper - lower) + if width == 0.0: + return np.asarray([lower], dtype=np.float32) + spacing = np.float32(2.0) * distortion_limit + if zero_primitive and lower < 0.0 < upper: + return _zero_anchored_grid(lower, upper, spacing) + count = 1 + int(np.ceil(width / spacing)) + return np.linspace(lower, upper, count, dtype=np.float32) + + +def build_primitives(config: ActionCodec2Config) -> PrimitiveStore: + """Derive primitive arrays and force an exact zero on every axis. + + Additive domains, SO(3) radii, scales, and distortion limits are read from + ``physical`` and never duplicated under ``train``. ``zero`` is a deployment + invariant, not a tuning option: every newly fitted primitive axis contains + an exact floating-point zero, including continuous gripper-servo axes. + """ + if config.train is None: + raise ConfigError("config.train is required to build primitives") + output: list[np.ndarray | None] = [None] * config.physical.input_dim + for component in config.physical.components.values(): + width = len(component.dims) + if component.space == "additive": + domain = component.constraints[f"order_{component.primitive_order}"] + lower = broadcast_vector( + domain.minimum, width, dtype=np.dtype(np.float32), name="minimum" + ) + upper = broadcast_vector( + domain.maximum, width, dtype=np.dtype(np.float32), name="maximum" + ) + epsilon = broadcast_vector( + component.distortion_limit, + width, + dtype=np.dtype(np.float32), + name="distortion_limit", + ) + if component.primitive_order == 2: + # Relaxed decoder: mu=(epsilon, 2*epsilon, 2*epsilon). + lower = lower - 2 * epsilon + upper = upper + 2 * epsilon + axes = [ + _additive_axis( + lower[index], + upper[index], + epsilon[index], + zero_primitive=True, + ) + for index in range(width) + ] + else: + domain = component.constraints[f"order_{component.primitive_order}"] + radius_dtype = ( + np.float64 + if component.representation == "rotation_vector" + else np.float32 + ) + lower = broadcast_vector( + domain.minimum, width, dtype=np.dtype(radius_dtype), name="minimum" + ) + upper = broadcast_vector( + domain.maximum, width, dtype=np.dtype(radius_dtype), name="maximum" + ) + corner = np.maximum(np.abs(lower), np.abs(upper)) + source_radius = ( + float(np.linalg.norm(corner)) + if component.representation == "rotation_vector" + else float(np.sum(corner, dtype=np.float32)) + ) + source_radius = min(math.pi, source_radius) + epsilon = float(component.distortion_limit) + allowance = (3 if component.primitive_order == 2 else 1) * epsilon + cover_radius = float(min(math.pi, source_radius + allowance)) + maximum_spacing = 2.0 * epsilon / math.sqrt(3.0) + count = int(math.ceil(2.0 * cover_radius / maximum_spacing)) + 1 + if count % 2 == 0: + count += 1 + axis = np.linspace(-cover_radius, cover_radius, count, dtype=np.float64) + if count % 2 == 1: + axis[count // 2] = 0.0 + axes = [axis.copy() for _ in range(width)] + for dimension, axis in zip(component.dims, axes, strict=True): + output[dimension] = axis + if any(axis is None for axis in output): + raise RuntimeError( + "internal error: primitive construction left an unassigned dimension" + ) + store = PrimitiveStore([axis for axis in output if axis is not None]) + missing_zero = [ + dimension + for dimension, axis in enumerate(store.values) + if not np.any(np.asarray(axis) == 0) + ] + if missing_zero: + raise ConfigError( + "every fitted primitive axis must contain exact zero; missing dimension(s) " + f"{missing_zero}. Ensure each configured domain contains zero." + ) + return store + + +def _validate_component_support( + active_dims: tuple[int, ...], config: ActionCodec2Config, *, index: int +) -> None: + active = set(active_dims) + for name, component in config.physical.components.items(): + present = active.intersection(component.dims) + if component.space == "so3" and present and present != set(component.dims): + raise ValueError( + f"trajectory {index} activates only part of SO(3) component {name!r}" + ) + + +def _length_buckets( + pending: Sequence[tuple[np.ndarray, tuple[int, ...], np.ndarray | None]], + *, + maximum_ratio: float = 1.25, +) -> list[list[int]]: + """Group nearby lengths so ragged episodes still use vectorized quantization.""" + ordered = sorted(range(len(pending)), key=lambda index: len(pending[index][0])) + buckets: list[list[int]] = [] + current: list[int] = [] + shortest = 0 + for index in ordered: + length = len(pending[index][0]) + if current and length > shortest * maximum_ratio: + buckets.append(current) + current = [] + if not current: + shortest = length + current.append(index) + if current: + buckets.append(current) + return buckets + + +class _CellCorpus(Sequence[np.ndarray]): + """Contiguous cell chunks plus compact episode offsets.""" + + _actioncodec2_validated = True + + def __init__(self, chunks: list[np.ndarray], offsets: list[np.ndarray]) -> None: + self.chunks: list[np.ndarray | None] = list(chunks) + self.offsets = tuple(offsets) + counts = np.asarray([len(value) - 1 for value in offsets], dtype=np.int64) + self._episode_ends = np.cumsum(counts) + + def __len__(self) -> int: + return int(self._episode_ends[-1]) if len(self._episode_ends) else 0 + + def __getitem__(self, index: int) -> np.ndarray: + if not isinstance(index, int): + raise TypeError("cell corpus indices must be integers") + if index < 0: + index += len(self) + if not 0 <= index < len(self): + raise IndexError(index) + chunk = bisect_right(self._episode_ends, index) + previous = 0 if chunk == 0 else int(self._episode_ends[chunk - 1]) + local = index - previous + start, stop = self.offsets[chunk][local : local + 2] + values = self.chunks[chunk] + if values is None: + raise RuntimeError("cell corpus chunk has already been consumed") + return values[int(start) : int(stop)] + + def __iter__(self): + for chunk, offsets in zip(self.chunks, self.offsets, strict=True): + if chunk is None: + raise RuntimeError("cell corpus chunk has already been consumed") + for start, stop in zip(offsets[:-1], offsets[1:], strict=True): + yield chunk[int(start) : int(stop)] + + def release_chunk(self, index: int) -> None: + """Release a replayed chunk while retaining compact episode metadata.""" + self.chunks[index] = None + + +def _quantize_pending_batch( + pending: Sequence[tuple[np.ndarray, tuple[int, ...], np.ndarray | None]], + physical: PhysicalCodec, + cell_dtype: type[np.signedinteger], +) -> list[np.ndarray]: + """Quantize one bounded ragged batch and restore its original order.""" + result: list[np.ndarray | None] = [None] * len(pending) + dimension = physical.config.physical.input_dim + for indices in _length_buckets(pending): + length = max(len(pending[index][0]) for index in indices) + batch = np.zeros((len(indices), length, dimension), dtype=np.float32) + active_by_item: list[tuple[int, ...]] = [] + for local, index in enumerate(indices): + trajectory, active_dims, _ = pending[index] + item = batch[local, : len(trajectory)] + item[:, active_dims] = trajectory + active_by_item.append(active_dims) + boundaries = [pending[index][2] for index in indices] + previous = ( + None + if all(value is None for value in boundaries) + else np.stack( + [ + np.zeros(dimension, dtype=np.float32) if value is None else value + for value in boundaries + ] + ) + ) + encoded = physical.encode_batch(batch, previous_first_order=previous).astype( + cell_dtype, copy=False + ) + for local, index in enumerate(indices): + item_length = len(pending[index][0]) + dense = np.full((item_length, dimension), -1, dtype=cell_dtype) + dense[:, active_by_item[local]] = encoded[local][ + :item_length, active_by_item[local] + ] + result[index] = dense + if any(value is None for value in result): + raise RuntimeError("internal error: trajectory batch was not fully quantized") + return [value for value in result if value is not None] + + +def _quantize_trajectories( + trajectories: Iterable[TrajectoryLike], + physical: PhysicalCodec, + *, + batch_size: int, + max_cell_bytes: int | None, +) -> tuple[_CellCorpus, int, int, dict[int, int]]: + """Quantize trajectories in length-homogeneous vectorized batches. + + Grouping is local to each bounded input batch, so generators remain + streaming and peak source memory is ``O(batch_size)``. Output order and + episode boundaries are preserved exactly for Temporal Set-BPE. + """ + if batch_size < 1: + raise ValueError("batch_size must be positive") + cell_chunks: list[np.ndarray] = [] + offset_chunks: list[np.ndarray] = [] + cell_parts: list[np.ndarray] = [] + cell_lengths: list[int] = [] + chunk_actions = 0 + action_count = 0 + active_value_count = 0 + arity_counts: dict[int, int] = defaultdict(int) + source_index = 0 + retained_cell_bytes = 0 + cell_dtype = np.int16 if max(physical.bins) <= np.iinfo(np.int16).max else np.int32 + iterator = iter(trajectories) + + def flush_cells() -> None: + nonlocal chunk_actions, retained_cell_bytes + if not cell_parts: + return + chunk = np.concatenate(cell_parts, axis=0) + projected = retained_cell_bytes + chunk.nbytes + if max_cell_bytes is not None and projected > max_cell_bytes: + raise MemoryError( + "quantized ActionCodec2 corpus would retain at least " + f"{projected / 2**30:.2f} GiB, exceeding the configured " + f"{max_cell_bytes / 2**30:.2f} GiB limit; use stratified fitting, " + "lower coverage, or explicitly raise --max-cell-memory-gib" + ) + cell_chunks.append(chunk) + retained_cell_bytes = projected + offset_chunks.append( + np.concatenate(([0], np.cumsum(cell_lengths, dtype=np.int64))) + ) + cell_parts.clear() + cell_lengths.clear() + chunk_actions = 0 + + while True: + pending: list[tuple[np.ndarray, tuple[int, ...], np.ndarray | None]] = [] + for _ in range(batch_size): + try: + raw = next(iterator) + except StopIteration: + break + trajectory, active_dims = prepare_trajectory( + raw, + index=source_index, + dimension=physical.config.physical.input_dim, + dtype=np.float32, + ) + _validate_component_support( + active_dims, physical.config, index=source_index + ) + source_index += 1 + action_count += len(trajectory) + active_value_count += len(trajectory) * len(active_dims) + arity_counts[len(active_dims)] += 1 + pending.append( + ( + trajectory, + active_dims, + trajectory_boundary(raw, physical.config, active_dims), + ) + ) + if not pending: + break + + for value in _quantize_pending_batch(pending, physical, cell_dtype): + cell_parts.append(value) + cell_lengths.append(len(value)) + chunk_actions += len(value) + if chunk_actions >= 1_000_000: + flush_cells() + flush_cells() + cells = _CellCorpus(cell_chunks, offset_chunks) + if not cells: + raise ValueError("training requires at least one complete trajectory") + return cells, action_count, active_value_count, dict(sorted(arity_counts.items())) + + +def _aggregate_trajectory_records( + cells: Sequence[np.ndarray], bins: Sequence[int], action_count: int +) -> tuple[np.ndarray, np.ndarray]: + """Aggregate records with memory bounded by one row chunk plus unique rows.""" + del action_count # retained in the signature for artifact-report compatibility + dimension = len(tuple(bins)) + widths = tuple(int(count).bit_length() for count in bins) + if sum(widths) <= 64: + return _aggregate_packed_records(cells, bins, widths) + counts_by_row: dict[bytes, int] = defaultdict(int) + pending: list[np.ndarray] = [] + pending_rows = 0 + chunk_rows = 1_000_000 + + def flush() -> None: + nonlocal pending_rows + if not pending: + return + block = ( + np.asarray(pending[0], dtype=np.int32) + if len(pending) == 1 + else np.concatenate(pending, axis=0).astype(np.int32, copy=False) + ) + unique, counts = aggregate_records(block, bins) + for row, count in zip(unique.astype(np.int32, copy=False), counts, strict=True): + counts_by_row[row.tobytes()] += int(count) + pending.clear() + pending_rows = 0 + + for trajectory in cells: + pending.append(np.asarray(trajectory, dtype=np.int32)) + pending_rows += len(trajectory) + if pending_rows >= chunk_rows: + flush() + flush() + keys = sorted(counts_by_row) + records = np.frombuffer(b"".join(keys), dtype=np.int32).reshape(-1, dimension) + weights = np.fromiter( + (counts_by_row[key] for key in keys), dtype=np.int64, count=len(keys) + ) + return records, weights + + +def _aggregate_packed_records( + cells: Sequence[np.ndarray], bins: Sequence[int], widths: Sequence[int] +) -> tuple[np.ndarray, np.ndarray]: + """Aggregate short ragged cell rows using one exact uint64 key per row.""" + counts_by_key: dict[int, int] = defaultdict(int) + pending: list[np.ndarray] = [] + pending_rows = 0 + chunk_rows = 1_000_000 + + def flush() -> None: + nonlocal pending_rows + if not pending: + return + block = pending[0] if len(pending) == 1 else np.concatenate(pending, axis=0) + keys = np.zeros(len(block), dtype=np.uint64) + shift = 0 + for dimension, width in enumerate(widths): + encoded = np.asarray(block[:, dimension], dtype=np.int64) + 1 + keys |= encoded.astype(np.uint64) << np.uint64(shift) + shift += width + unique, counts = np.unique(keys, return_counts=True) + for key, count in zip(unique, counts, strict=True): + counts_by_key[int(key)] += int(count) + pending.clear() + pending_rows = 0 + + for trajectory in cells: + pending.append(trajectory) + pending_rows += len(trajectory) + if pending_rows >= chunk_rows: + flush() + flush() + + keys = np.fromiter(sorted(counts_by_key), dtype=np.uint64, count=len(counts_by_key)) + records = np.empty((len(keys), len(tuple(bins))), dtype=np.int32) + shift = 0 + for dimension, width in enumerate(widths): + mask = np.uint64((1 << width) - 1) + records[:, dimension] = ((keys >> np.uint64(shift)) & mask).astype(np.int32) - 1 + shift += width + weights = np.fromiter( + (counts_by_key[int(key)] for key in keys), dtype=np.int64, count=len(keys) + ) + return records, weights + + +def _fit_report( + config: ActionCodec2Config, + primitives: PrimitiveStore, + codec: ActionCodec2Processor, + *, + trajectory_count: int, + actions: int, + active_values: int, + arity_counts: dict[int, int], + cell_dtype: str, + unique_records: int, + spatial_report: dict[str, object], + temporal_report: dict[str, object], + timestamps: tuple[float, float, float, float, float, float], +) -> dict[str, object]: + started, primitive_done, quantized, aggregated, spatial_done, finished = timestamps + scheme = config.set_bpe.scheme + stages: dict[str, object] = ( + {"bpe": temporal_report} + if scheme == "bpe" + else {"spatial": spatial_report, "temporal": temporal_report} + ) + return { + "trajectories": trajectory_count, + "actions": actions, + "active_values": active_values, + "trajectory_arities": {str(key): value for key, value in arity_counts.items()}, + "bpe_boundary": "complete_component_trajectory", + "bpe_training_domain": "quantized_primitive_cells", + "physical_boundary": ( + "explicit_previous_first_order" + if codec.requires_previous_first_order + else "zero_position_identity_orientation" + ), + "component_primitive_orders": { + name: component.primitive_order + for name, component in config.physical.components.items() + }, + "allow_cross_component_merge": config.set_bpe.allow_cross_component_merge, + "lookahead": config.set_bpe.lookahead, + "physical_bins": list(primitives.bins), + "cell_dtype": cell_dtype, + "unique_spatial_records": unique_records, + "primitive_atoms": sum(primitives.bins), + "tokenization_scheme": scheme, + **stages, + "vocab_size": codec.vocab_size, + "tokens_per_action": float(temporal_report["tokens_per_action"]), + "seconds": { + "primitives": primitive_done - started, + "quantize": quantized - primitive_done, + "aggregate": aggregated - quantized, + "spatial_train": spatial_done - aggregated, + "spatial": spatial_done - quantized, + "temporal": finished - spatial_done, + "total": finished - started, + }, + } + + +def _fit_set_bpe( + cells: Sequence[np.ndarray], + config: ActionCodec2Config, + primitives: PrimitiveStore, + actions: int, + *, + backend: str, + threads: int, + progress: bool, +) -> tuple[SetBPESerializer, dict[str, object], dict[str, object], int, float, float]: + records, weights = _aggregate_trajectory_records(cells, primitives.bins, actions) + aggregated = time.perf_counter() + bpe = config.train.bpe if config.train is not None else None + assert bpe is not None + spatial, spatial_report = train_spatial( + records, + weights, + primitives.bins, + vocab_size=bpe.spatial_vocab_size, + backend=backend, + threads=threads, + progress=progress, + ) + spatial_done = time.perf_counter() + vocabulary, temporal_report = train_temporal( + cells, + spatial, + vocab_size=bpe.vocab_size, + lookahead=config.set_bpe.lookahead, + backend=backend, + threads=threads, + progress=progress, + ) + serializer = SetBPESerializer( + bins=primitives.bins, + lookahead=config.set_bpe.lookahead, + spatial_merges=spatial.merges, + temporal_merges=vocabulary.merges, + ) + return ( + serializer, + spatial_report, + temporal_report, + len(records), + aggregated, + spatial_done, + ) + + +def _fit_bpe( + cells: Sequence[np.ndarray], + config: ActionCodec2Config, + primitives: PrimitiveStore, + *, + backend: str, + threads: int, + progress: bool, +) -> tuple[BPESerializer, dict[str, object], dict[str, object], int, float, float]: + """Fit the FAST-style baseline on the same quantized cells. + + Plain BPE has a single merge stage and no record aggregation, so the + spatial slots of the shared fit signature stay empty and both stage + timestamps collapse onto the start of the one stage that runs. + """ + started = time.perf_counter() + bpe = config.train.bpe if config.train is not None else None + assert bpe is not None + vocabulary, report = train_bpe( + cells, + primitives.bins, + vocab_size=bpe.vocab_size, + max_token_cells=config.set_bpe.lookahead * config.physical.input_dim, + backend=backend, + threads=threads, + progress=progress, + ) + serializer = BPESerializer( + bins=primitives.bins, + lookahead=config.set_bpe.lookahead, + merges=vocabulary.merges, + ) + return serializer, {}, report, 0, started, started + + +def _validate_fit_options( + config: ActionCodec2Config, max_cell_bytes: int | None +) -> None: + if config.train is None: + raise ConfigError("config.train is required for fitting") + if max_cell_bytes is not None and ( + isinstance(max_cell_bytes, bool) + or not isinstance(max_cell_bytes, int) + or max_cell_bytes < 1 + ): + raise ValueError("max_cell_bytes must be a positive integer or None") + if config.set_bpe.allow_cross_component_merge: + raise ConfigError( + "allow_cross_component_merge=true requires the future multi-component " + "corpus representation; v1 fitting supports the default isolated boundary only" + ) + if config.set_bpe.scheme == "bpe" and config.train.bpe.spatial_vocab_size != ( + config.train.bpe.vocab_size + ): + raise ConfigError( + "set_bpe.scheme=bpe has a single merge stage; set " + "train.bpe.spatial_vocab_size equal to train.bpe.vocab_size to make the " + "unused split explicit" + ) + + +class ProfileTrainer: + """Internal whole-trajectory fitter for one physical profile. + + User workflows enter through ``ActionCodec2.fit``; repository pretraining + jobs also use this class after their own source routing and quality checks. + + Args: + config: Physical and BPE configuration including a ``train`` section. + """ + + def __init__(self, config: ActionCodec2Config) -> None: + self.config = config + + def fit( + self, + trajectories: Iterable[TrajectoryLike], + *, + backend: str = "auto", + threads: int = 0, + batch_size: int = 256, + progress: bool = False, + max_cell_bytes: int | None = None, + ) -> tuple[ActionCodec2Processor, dict[str, object]]: + """Fit primitives and the configured tokenization scheme on whole trajectories. + + Neither scheme consumes the floating-point trajectories directly. The + physical codec first converts them to primitive-cell streams; frequency + fitting operates only on those discrete cells, so ``set_bpe.scheme`` + changes the tokenizer and nothing upstream of it. + + Parameters + ---------- + trajectories: + Iterable of complete variable-width episodes. Plain ``[T,D]`` arrays + use prefix dimensions; :class:`ActiveTrajectory` supplies explicit + global dimension ids for layouts such as a fixed gripper slot. + Second-order components require ``ActiveTrajectory.previous_first_order`` + with one preceding reconstructed value per active coordinate. + backend, threads: + Exact BPE backend and its worker count. ``"auto"`` uses the native + sparse trainers when available. + batch_size: + Maximum trajectories buffered for vectorized physical quantization. + progress: + Forwarded to both Set-BPE trainers. + max_cell_bytes: + Optional hard bound on retained quantized-corpus bytes. Temporal BPE + requires replaying the full cell corpus, so callers operating on large + datasets should set this instead of relying on available RAM. + """ + config = self.config + _validate_fit_options(config, max_cell_bytes) + started = time.perf_counter() + primitives = build_primitives(config) + primitive_done = time.perf_counter() + + atoms = sum(primitives.bins) + bpe = config.train.bpe + if bpe.spatial_vocab_size < atoms: + raise ConfigError( + f"train.bpe.spatial_vocab_size must be at least the {atoms} primitive atoms" + ) + + physical = PhysicalCodec(config, primitives) + cells, actions, active_values, arity_counts = _quantize_trajectories( + trajectories, + physical, + batch_size=batch_size, + max_cell_bytes=max_cell_bytes, + ) + trajectory_count = len(cells) + first_cell_dtype = str(cells[0].dtype) + quantized = time.perf_counter() + + if config.set_bpe.scheme == "bpe": + fitted = _fit_bpe( + cells, + config, + primitives, + backend=backend, + threads=threads, + progress=progress, + ) + else: + fitted = _fit_set_bpe( + cells, + config, + primitives, + actions, + backend=backend, + threads=threads, + progress=progress, + ) + ( + serializer, + spatial_report, + temporal_report, + unique_records, + aggregated, + spatial_done, + ) = fitted + codec = ActionCodec2Processor(config, primitives, serializer) + finished = time.perf_counter() + + report = _fit_report( + config, + primitives, + codec, + trajectory_count=trajectory_count, + actions=actions, + active_values=active_values, + arity_counts=arity_counts, + cell_dtype=first_cell_dtype, + unique_records=unique_records, + spatial_report=spatial_report, + temporal_report=temporal_report, + timestamps=( + started, + primitive_done, + quantized, + aggregated, + spatial_done, + finished, + ), + ) + return codec, report diff --git a/setup.py b/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..d707f12d9c66dad1cd1d9daffe22a9839a01f002 --- /dev/null +++ b/setup.py @@ -0,0 +1,15 @@ +from setuptools import setup +from pybind11.setup_helpers import Pybind11Extension, build_ext + +native = "actioncodec2_kernels" +setup( + ext_modules=[Pybind11Extension( + "actioncodec2_kernels._actioncodec2_native", + sources=[f"{native}/{name}.cpp" for name in ("module", "physical", "so3", "setbpe", "resample")], + include_dirs=[native], + depends=[f"{native}/parallel.hpp"], + extra_compile_args=["-O3", "-fvisibility=hidden", "-ffp-contract=off"], + cxx_std=17, + )], + cmdclass={"build_ext": build_ext}, +)