| """Configuration-driven physical quantization.""" |
|
|
| 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 ._validation import broadcast_vector |
| from .configuration import ActionCodec2Config, ComponentConfig, ConfigError |
| from .native_backend import native_kernel, require_native, resolve_backend |
| from .primitives import PrimitiveStore |
| from .semantic import difference_additive as _difference_additive |
| from .semantic import integrate_additive as _integrate_additive |
|
|
| _FLOAT32_SLACK = np.float32(8 * np.finfo(np.float32).eps) |
| _IDENTITY = np.asarray((0.0, 0.0, 0.0, 1.0), dtype=np.float64) |
|
|
| _SELECTION_CODES = {"nearest": 0, "smooth": 1, "momentum_lazy": 2} |
|
|
|
|
| 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: |
| values = np.asarray(actions, dtype=np.float32) * _scale(component) |
| if component.input_transform.clip: |
| bounds = _constraint(component, component.input_order) |
| if bounds is None: |
| 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) |
|
|
|
|
| class _ComponentRuntime(Protocol): |
| 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: |
| 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: |
| flat_request = np.asarray(request, dtype=np.float32).reshape(-1, self.width) |
| flat_direction = np.asarray(direction, dtype=np.float32).reshape(-1, self.width) |
| kernel = native_kernel("choose_cells") |
| if kernel is not None: |
| return kernel( |
| self._native_table, |
| self._native_bins, |
| flat_request, |
| flat_direction, |
| self.limit, |
| self.config.primitive_selection == "smooth", |
| ).reshape(request.shape) |
| 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: |
| 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 resolve_backend("auto") == "native": |
| lower, upper = ( |
| self.position_bounds |
| if self.position_bounds is not None |
| else ( |
| np.full(self.width, -np.inf, dtype=np.float32), |
| np.full(self.width, np.inf, dtype=np.float32), |
| ) |
| ) |
| 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(lower, dtype=np.float32), |
| np.ascontiguousarray(upper, 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) |
| kernel = native_kernel("additive_decode") |
| if kernel is not None: |
| lower, upper = ( |
| self.position_bounds |
| if self.position_bounds is not None |
| else ( |
| np.full(self.width, -np.inf, dtype=np.float32), |
| np.full(self.width, np.inf, dtype=np.float32), |
| ) |
| ) |
| return kernel( |
| self._native_table, |
| indices, |
| _scale(self.config), |
| lower, |
| upper, |
| self.config.input_order, |
| self.order, |
| ) |
| 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: |
| 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) |
|
|
| 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): |
| 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, |
| ) |
|
|
| 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 |
|
|
| 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: |
| kernel = native_kernel("additive_reconstruct") |
| if kernel is not None: |
| return kernel( |
| self._native_table, |
| cells, |
| previous_first_order, |
| _scale(self.config), |
| self.config.input_order, |
| self.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: |
| 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) |
| self.lazy = component.primitive_selection == "momentum_lazy" |
| self._zero = np.asarray([int(np.argmin(np.abs(values))) for values in self.primitives]) |
| if self.lazy and any( |
| values[index] != 0.0 for values, index in zip(self.primitives, self._zero, strict=True) |
| ): |
| raise ValueError("momentum_lazy selection requires an exact zero primitive per axis") |
|
|
| 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, |
| previous: np.ndarray | None = None, |
| ) -> np.ndarray: |
| 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 | None = None |
| best_preferred: tuple | None = None |
| repeat = previous is not None and tuple(int(v) for v in previous) |
| 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 |
| if self.lazy and choice == repeat: |
| return np.asarray(choice, dtype=np.int64) |
| offset = np.abs(np.asarray(choice) - self._zero) |
| candidate = ( |
| (int(np.count_nonzero(offset)), int(offset.sum()), error, choice) |
| if self.lazy |
| else (error, choice) |
| ) |
| if best_any is None or candidate < best_any: |
| best_any = candidate |
|
|
| 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, |
| _SELECTION_CODES[self.config.primitive_selection], |
| 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) |
| previous = np.broadcast_to(self._zero, (batch, 3)) |
| 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 != "nearest": |
| choice = np.stack( |
| [ |
| self._smooth( |
| residual[index], |
| reconstruction[index], |
| true[index], |
| direction[index], |
| choice[index], |
| previous[index] if self.lazy else None, |
| ) |
| for index in range(batch) |
| ] |
| ) |
| previous = choice |
| 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.""" |
| kernel = native_kernel("so3_reconstruct") |
| if kernel is not None: |
| previous = previous_first_order |
| if previous is None: |
| previous = np.zeros((len(cells), 3), dtype=np.float32) |
| return kernel( |
| self._native_table, |
| cells, |
| previous, |
| _scale(self.config), |
| self.config.input_order, |
| self.config.primitive_order, |
| self.left, |
| self.config.representation == "rotation_vector", |
| getattr(self, "second_order_radius", np.inf), |
| ) |
| 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) |
|
|
| 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: |
| 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: |
| 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: |
| 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: |
| 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.""" |
|
|
| 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.""" |
| 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.""" |
| 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 |
|
|
