"""Hypothesis batches: run B parallel model/solver hypotheses through the batched
solver each MPC tick and let reality pick the winner.
This is the "batch-as-identity" API: the batch dimension carries *hypotheses*
(disturbance guesses, solver-parameter spreads, ...) rather than independent
problems. Lifecycle per tick:
hypotheses.apply(solver, x_measured) # before solve: program the batch
result = solver.solve(...)
best = hypotheses.select(solver, result, x_prev, u_prev, x_measured, dt)
At batch_size == 1 both hooks are no-ops (select returns 0), so a controller
built with hypotheses=None or B=1 behaves as a plain single-solve MPC.
"""
from abc import ABC, abstractmethod
import numpy as np
from .common import _require_pin
[docs]
def world_wrench_to_gato_slot(model, data, ee_frame_id, q, f_world):
"""World-axes wrench at the EE frame origin -> GRiD's last-body f_ext slot.
GRiD's slot is the last body's spatial force in the LAST JOINT's local frame
about the joint origin, Featherstone-ordered [angular(3); linear(3)] — verified
against pin.aba to 1e-8 (fext_frame_probe 2026-07-07). An older version got
both the frame chain and the [linear; angular] ordering wrong, which is why a
seeded-truth hypothesis LOST to zero-wrench rollouts (~45% win rate)."""
from gato.common import world_wrench_to_joint_local
_, Fj = world_wrench_to_joint_local(model, data, q, f_world, ee_frame_id)
return np.concatenate([Fj.angular, Fj.linear])
[docs]
class HypothesisBatch(ABC):
"""B parallel hypotheses evaluated by the batched solver each MPC tick."""
batch_size: int
[docs]
@abstractmethod
def apply(self, solver, x):
"""Sample/refresh B hypotheses and program them into the solver
(set_f_ext_B, set_rho_penalty_batch, set_mu_batch, ...)."""
[docs]
def select(self, solver, result, x_prev, u_prev, x_meas, dt):
"""Pick the hypothesis that best explains reality; return its index.
Default scorer: roll the PREVIOUS state/control one step under each
hypothesis' dynamics on the GPU and compare to the measured next state:
x_next_B = solver.sim_forward(x_prev, u_prev, dt) # per-hypothesis f_ext
best = argmin_b || x_next_B[b] - x_meas ||
Hypothesis types that don't alter the dynamics (e.g. per-batch rho/mu
spreads) should override select() and score from ``result.stats``
(e.g. argmin final_merit) instead.
"""
x_next_batch = solver.sim_forward(np.asarray(x_prev, dtype=np.float32),
np.asarray(u_prev, dtype=np.float32), dt)
errors = np.linalg.norm(x_next_batch - np.asarray(x_meas)[None, :], axis=1)
best_id = int(np.argmin(errors))
self.update(best_id, errors)
return best_id
[docs]
@abstractmethod
def update(self, best_id, errors):
"""Feed the selection outcome back into the hypothesis sampler."""
[docs]
def get_stats(self):
return {}
[docs]
class ForceHypothesisBatch(HypothesisBatch):
"""Wrench-disturbance hypotheses driven by a ForceEstimator/CEMForceEstimator.
``estimator`` is duck-typed: generate_batch() -> (B, 6) WORLD-frame wrenches,
update(best, errors, batch_used), reset(), get_stats(). The adapter owns the
world->GATO frame transform (URDF ``ee_frame`` -> parent-joint local wrench)
and uploads the transformed batch via solver.set_f_ext_B.
"""
def __init__(self, estimator, model, ee_frame="EE"):
_require_pin()
self.estimator = estimator
self.batch_size = estimator.batch_size
self.model = model # solver model (no pendulum augmentation)
self._data = model.createData() # one Data, reused every tick
self._ee_frame_id = model.getFrameId(ee_frame)
if self._ee_frame_id >= model.nframes:
raise ValueError(f"URDF has no frame named {ee_frame!r}")
self._jid_ee = model.frames[self._ee_frame_id].parentJoint
self._last_world_batch = None
[docs]
def apply(self, solver, x):
world = self.estimator.generate_batch() # (B, 6) world frame
self._last_world_batch = world
q = np.asarray(x)[: self.model.nq]
gato_frame = np.stack([self._world_to_gato(q, w) for w in world])
solver.set_f_ext_B(gato_frame.astype(np.float32))
[docs]
def update(self, best_id, errors):
# estimator state lives in the WORLD frame -> pass the world-frame batch
self.estimator.update(best_id, errors, batch_used=self._last_world_batch)
[docs]
def reset(self):
self.estimator.reset()
self._last_world_batch = None
[docs]
def get_stats(self):
return self.estimator.get_stats()
def _world_to_gato(self, q, f_world):
return world_wrench_to_gato_slot(self.model, self._data, self._ee_frame_id, q, f_world)
[docs]
class IdentifiedWrenchBatch(HypothesisBatch):
"""B=1 "batch": the wrench IDENTIFIED from the last one-step motion residual.
Structurally identical to ForceHypothesisBatch — same world->GATO frame
transform, same ``set_f_ext_B`` injection — so an A/B against it isolates
exactly one variable: how the wrench is obtained. ForceHypothesisBatch
SAMPLES B candidates and lets a GPU rollout pick a winner;
this SOLVES for the wrench with one least-squares fit and needs no batch.
Lifecycle note: the fit needs (x_prev, u_prev, x_measured), which the
controller supplies to ``select()`` — i.e. AFTER the solve. The estimate
therefore lands one tick later, in the next ``apply()``. That one-tick lag
is inherent to any observer (the FE has it too: its batch is scored against
the previous interval).
Args:
identifier: an OneStepWrenchIdentifier over the SOLVER model.
model: the solver's pinocchio model (robot only).
ee_frame: URDF frame the wrench acts at.
n_actuated: width of the applied-torque head of the control vector
(``solver.n_actuated``); on contact-force builds ``u_prev`` also
carries the solver's fc tail, which is NOT an applied torque and
must not enter the residual.
"""
def __init__(self, identifier, model, ee_frame="EE", n_actuated=None):
_require_pin()
self.identifier = identifier
self.batch_size = 1
self.model = model
self._data = model.createData()
self._ee_frame_id = model.getFrameId(ee_frame)
if self._ee_frame_id >= model.nframes:
raise ValueError(f"URDF has no frame named {ee_frame!r}")
self.n_actuated = int(n_actuated) if n_actuated is not None else model.nv
[docs]
def apply(self, solver, x):
w = self.identifier.generate_batch() # (1, 6) world frame
q = np.asarray(x)[: self.model.nq]
solver.set_f_ext_B(
np.stack([self._world_to_gato(q, w[0])]).astype(np.float32))
[docs]
def select(self, solver, result, x_prev, u_prev, x_meas, dt):
"""No-op: there is only one hypothesis.
The fit itself is driven at SENSOR rate by the driver
(MPC_GATO._observe_substep), not from this tick-rate hook — the
control-period motion sample is too coarse to identify a wrench from
(9.4% vs 0.5% median error; see OneStepWrenchIdentifier.identify).
A caller running its own sim loop should call identifier.identify()
at its measurement rate instead.
"""
return 0
[docs]
def update(self, best_id, errors):
"""No-op: select() already folded the observation into the fit."""
[docs]
def reset(self):
self.identifier.reset()
[docs]
def get_stats(self):
return self.identifier.get_stats()
def _world_to_gato(self, q, f_world):
return world_wrench_to_gato_slot(self.model, self._data, self._ee_frame_id, q, f_world)