from copy import deepcopy
import dataclasses
import io
import os
import pathlib
import time
import numpy as np
from PIL import Image
from scipy.spatial.transform import Rotation
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
JPEG_QUALITY = 90
SIDES = ("left", "right")
POSE_KEY = "observation.pose.left_hand_root_to_right_hand_root.absolute"
_K = ("observation.image.left", "observation.image.right", POSE_KEY,
"observation.joint_states", "prompt")
CUP_PA1 = "Pick up the cup with your right hand and set it on the plate"
CUP_PA2 = "Pick up the cup with your left hand and return the cup to its original position"
CUP_HANDS = {CUP_PA1: ("right", 15), CUP_PA2: ("left", 14)}
CUP_ADVANCE = {CUP_PA1: 8, CUP_PA2: 4}
CUP_NOISE = {CUP_PA1: .15, CUP_PA2: .10}
PEN_PROMPTS = (
"Pick up a pencil on the table and put it into the pen holder",
"Pick up a ballpoint pen on the table and put it into the pen holder",
)
def _asset_id(ckpt):
found = sorted(p.parent.name for p in (ckpt / "assets").glob("*/norm_stats.json"))
if len(found) != 1:
raise FileNotFoundError(f"expected one assets/<id>/norm_stats.json in {ckpt}, found {found}")
return found[0]
def _m(r):
a, b = r[..., :3], r[..., 3:6]
c0 = a / np.linalg.norm(a, axis=-1, keepdims=True)
b = b - np.sum(c0 * b, axis=-1, keepdims=True) * c0
c1 = b / np.linalg.norm(b, axis=-1, keepdims=True)
return np.stack([c0, c1, np.cross(c0, c1)], axis=-1)
def _d(x):
rot = Rotation.from_matrix(_m(x[:, 3:9]))
pr = Rotation.concatenate([Rotation.identity(), rot[:-1]])
pp = np.concatenate([np.zeros((1, 3)), x[:-1, :3]])
return np.concatenate([pr.inv().apply(x[:, :3] - pp), (pr.inv() * rot).as_quat(canonical=True)], axis=1)
def _image(img):
from openpi.shared import image_tools
a = np.asarray(image_tools.resize_with_pad(np.asarray(img, dtype=np.uint8), 224, 224))
b = io.BytesIO(); Image.fromarray(a).save(b, format="JPEG", quality=JPEG_QUALITY)
return np.asarray(Image.open(io.BytesIO(b.getvalue())).convert("RGB"))
def _build_transforms():
from openpi import transforms as T
from openpi.training.config import DataConfigFactory, ModelTransformFactory
@dataclasses.dataclass(frozen=True)
class _P(T.DataTransformFn):
def __call__(self, d):
return {k: d[k] for k in _K if k in d}
@dataclasses.dataclass(frozen=True)
class _I(T.DataTransformFn):
n: int = 32
def __call__(self, d):
l, r = _image(d[_K[0]]), _image(d[_K[1]])
p = np.array(d[_K[2]], dtype=np.float64)
if p[6] < 0:
p[3:] *= -1
s = np.concatenate([p, np.asarray(d[_K[3]], dtype=np.float64)]).astype(np.float32)
o = {"state": T.pad_to_dim(s, self.n),
"image": {"base_0_rgb": np.zeros_like(l), "left_wrist_0_rgb": l, "right_wrist_0_rgb": r},
"image_mask": {"base_0_rgb": np.False_, "left_wrist_0_rgb": np.True_, "right_wrist_0_rgb": np.True_}}
if "prompt" in d:
o["prompt"] = d["prompt"]
return o
@dataclasses.dataclass(frozen=True)
class _O(T.DataTransformFn):
def __call__(self, d):
a = np.asarray(d["actions"], dtype=np.float64)
return {"actions": np.concatenate([_d(a[:, 0:9]), _d(a[:, 9:18]), a[:, 18:20]], axis=1).astype(np.float32)}
@dataclasses.dataclass(frozen=True)
class _C(DataConfigFactory):
def create(self, assets_dirs, model_config):
return dataclasses.replace(
self.create_base_config(assets_dirs, model_config),
repack_transforms=T.Group(inputs=[_P()]),
data_transforms=T.Group(inputs=[_I(model_config.action_dim)], outputs=[_O()]),
model_transforms=ModelTransformFactory()(model_config))
return _C
class CupContinuousAdvance:
def __init__(self, advance_rows):
if type(advance_rows) is not int or not 1 <= advance_rows <= 16:
raise ValueError('Continuous place advance must be in 1..16 rows')
self.advance_rows = advance_rows
self.state = dict(prompt=None, clock=0, saw_open=False, armed=False, deadline=None,
waveform=None, latched=False, cursor=0, done=False)
self._generation = 0
self._owner = object()
def propose(self, actions, obs, adopt_rows):
a = np.asarray(actions)
if a.shape != (32, 16) or a.dtype != np.float32 or not np.isfinite(a).all():
raise ValueError('Continuous place advance expects finite float32[32,16]')
if type(adopt_rows) is not int or not 1 <= adopt_rows <= 32:
raise ValueError('Continuous place advance requires 1..32 adopted rows')
prompt = obs.get('prompt', '').strip().rstrip('.')
if prompt not in CUP_HANDS:
raise ValueError('Continuous place advance requires task1 PA1/PA2')
fingers = np.asarray(obs['observation.joint_states'], dtype=float)
if fingers.shape != (2,) or not np.isfinite(fingers).all():
raise ValueError('Expected two finite measured finger angles')
side, column = CUP_HANDS[prompt]
hand = column - 14
state = deepcopy(self.state)
if state['prompt'] != prompt:
state = dict(prompt=prompt, clock=0, saw_open=False, armed=False, deadline=None,
waveform=None, latched=False, cursor=0, done=False)
output = a.copy()
if not state['done'] and not state['armed']:
if fingers[hand] > .53:
state['saw_open'] = True
elif state['saw_open'] and 3 * 2 * np.pi / 4096 <= fingers[hand] <= .53:
state['armed'] = True
raw = a[:, column]
threshold = min(.53, float(raw[0]) + .05)
opening = None
if np.any(raw >= .56):
for i in range(31):
if raw[i] >= threshold and raw[i + 1] >= threshold:
opening = i + 1
break
if state['armed'] and fingers[hand] >= .56 and raw[0] >= .56:
state.update(latched=False, done=True, deadline=None, waveform=None)
start = None
if state['armed'] and not state['latched'] and not state['done']:
if opening is None:
state.update(deadline=None, waveform=None)
else:
if state['deadline'] is None:
state['deadline'] = state['clock'] + opening
state['waveform'] = raw[opening - 1:].tolist()
desired = state['deadline'] - state['clock'] - self.advance_rows
candidate = max(desired, 1)
if candidate <= adopt_rows and candidate < opening:
start = candidate
natural_start = (opening if state['armed'] and not state['done']
and not state['latched'] and start is None
and opening is not None and opening <= adopt_rows else None)
report = dict(enabled=True, kind='cup_continuous_place_advance_v1',
hand=side, advance_rows=self.advance_rows, clock=state['clock'],
forecast_opening_row=opening, opening_threshold_rad=threshold,
anchored_opening_row=state['deadline'], forced_start_row=start,
natural_start_row=natural_start,
already_latched=state['latched'], modified=False,
waveform='continuous values from the first reserved forecast',
pick_advance=False, height_or_robot_state_used=False)
release_row = start or natural_start
hold_end = ((release_row - 1) // 4 + 1) * 4 if release_row else 0
states = []
for index in range(adopt_rows):
if start == index + 1:
state.update(latched=True, cursor=0, open_floor=float(fingers[hand]))
if state['latched']:
wave = state['waveform']
if not wave:
raise RuntimeError('Missing reserved continuous opening waveform')
value = max(float(fingers[hand]), state.get('open_floor', 0.),
wave[min(state['cursor'], len(wave)-1)])
output[index, column] = value
state.update(cursor=state['cursor']+1, open_floor=value)
report['modified'] = True
if release_row and release_row <= index + 1 <= hold_end:
output[index, :14] = np.tile([0., 0., 0., 0., 0., 0., 1.], 2)
report['modified'] = True
state['clock'] += 1
states.append(deepcopy(state))
return dict(actions=output, states=states, report=report,
adopt_rows=adopt_rows, _owner=self._owner, _generation=self._generation)
def commit(self, proposal, executed_rows):
if proposal.get('_owner') is not self._owner or proposal.get('_generation') != self._generation:
raise ValueError('Stale continuous place advance proposal')
if type(executed_rows) is not int or not 0 <= executed_rows <= proposal['adopt_rows']:
raise ValueError('Invalid verified executed prefix')
if executed_rows:
self.state = deepcopy(proposal['states'][executed_rows - 1])
self._generation += 1
PEN_POLICY = dict(kind='u1_pen_delayed_bidirectional_32_retention_narrow_retry_v2',
lookahead_rows=16, retained_forecast_rows=32, command_row_1based=8,
delay_completed_chunks=1, close_below_rad=.07, release_at_or_above_rad=.27,
preclose_open_range_rad=[.07, .27], close_cancel_opening_delta_rad=.02,
close_target_rad=0., open_target='first_remaining_raw_release_value',
latch_max_completed_chunks=2, latch_release='natural_prefix_catchup_or_countermand_or_forecast_withdrawal_or_bound',
initial_closed_is_verified_grasp=False, narrow_retry='current_not_holding_plus_future_close_and_no_pending_open_or_open_latch_or_prefix_open', state_commit='verified_full_eight_rows_only',
raw_four_stage_widths_preserved=True, ground_truth_inputs=False)
def policy_for_rows(rows=8, *, place_advance_rows=0, close_lookahead=True, close_defer_only=False):
if type(rows) is not int or rows not in (8, 16):
raise ValueError('U1 pen lookahead execution rows must be8 or16')
if type(close_defer_only) is not bool or (close_defer_only and (not close_lookahead or rows != 8)):
raise ValueError('Confirmed defer-only requires close lookahead enabled and adopt8')
if type(close_lookahead) is not bool:
raise ValueError('close_lookahead must be boolean')
if type(place_advance_rows) is not int or not 0 <= place_advance_rows <= 16:
raise ValueError('Pen place advance must be an integer in0..16')
policy = (deepcopy(PEN_POLICY) if rows == 8 else
dict(deepcopy(PEN_POLICY), kind='u1_pen_delayed_bidirectional_32_retention_adopt16_v3',
lookahead_rows=32, command_row_1based=16, state_commit='verified_full_sixteen_rows_only'))
if place_advance_rows:
policy.update(kind=policy['kind']+'_place_advance_v1', place_advance_rows=place_advance_rows,
release_rule='first raw >=0.27 forecast absolute executed-row deadline minus advance; earliest current row1',
release_target='first current raw opening value; existing bounded open latch',
release_forecast_rows=32, release_withdrawal_cancels=True,
pick_advance_changed=False, arm_actions_changed=False)
if not close_lookahead:
policy.update(kind=policy['kind']+'_open_only_v1', close_lookahead_enabled=False,
close_rule='raw model timing only; no forced-close reservation or latch',
normal_model_closing_preserved=True)
if close_defer_only:
policy.update(kind=policy['kind']+'_confirmed_close_defer_v1', close_defer_only=True,
close_defer_near_rows=12, close_defer_confirmations=2, close_defer_reserve_rows=16,
close_defer_deadline='first eligible raw close forecast; absolute verified executed-row clock',
close_defer_trigger='deadline already expired, raw close within first12 and delayed for two consecutive completed observations',
close_defer_command_row=rows, normal_model_closing_preserved=True)
return policy
class PenGraspSchedule:
def __init__(self, rows=8, *, place_advance_rows=0, close_lookahead=True, close_defer_only=False):
self.rows = rows
self.policy = policy_for_rows(rows, place_advance_rows=place_advance_rows, close_lookahead=close_lookahead, close_defer_only=close_defer_only)
self.close_lookahead = close_lookahead
self.close_defer_only = close_defer_only
self.close_defer_state = {s: None for s in SIDES}
self.place_advance_rows = place_advance_rows
self.release_deadlines = {s: None for s in SIDES}
self.executed_row_clock = 0
self.pending = {s: None for s in SIDES}
self.phases = {s: None for s in SIDES}
self.verified_close = {s: False for s in SIDES}
self.seen_open = {s: False for s in SIDES}
self.latches = {s: None for s in SIDES}
self.last_committed_cycle = -1
self._generation = 0
self._owner = id(self)
def propose(self, cycle, predictions, targets, measured_angles, readings, limits, *, allow_clipping):
prediction = np.asarray(predictions)
selected = np.asarray(targets, dtype=float).copy()
measured = np.asarray(measured_angles, dtype=float)
if (type(cycle) is not int or cycle <= self.last_committed_cycle
or prediction.shape != (32, 16) or prediction.dtype.kind != 'f'
or not np.isfinite(prediction).all() or selected.shape != (self.rows, 2)
or not np.isfinite(selected).all() or measured.shape != (2,)
or not np.isfinite(measured).all()):
raise ValueError('Invalid U1 pen transition inputs/cycle')
next_close_defer = deepcopy(self.close_defer_state)
next_release_deadlines = deepcopy(self.release_deadlines)
next_pending = {s: None for s in SIDES}
next_phases = deepcopy(self.phases)
next_verified = deepcopy(self.verified_close)
next_seen = deepcopy(self.seen_open)
next_latches = {s: None for s in SIDES}
report = dict(enabled=True, policy=deepcopy(self.policy), cycle=cycle,
state_committed=False, pending_before=deepcopy(self.pending),
place_advance_rows=self.place_advance_rows, executed_row_clock=self.executed_row_clock,
close_defer_only=self.close_defer_only, close_defer_before=deepcopy(self.close_defer_state),
release_deadlines_before=deepcopy(self.release_deadlines),
phases_before=deepcopy(self.phases), latches_before=deepcopy(self.latches),
verified_close_before=deepcopy(self.verified_close), hands={},
_owner=self._owner, _generation=self._generation)
for col, side in enumerate(SIDES):
values = prediction[:, 14 + col]
adopted = selected[:, col].copy()
angle = float(measured[col])
holding = readings.get(side, {}).get('contact_force_hold') is True
verified = bool(self.verified_close[side] or holding)
seen_open = bool(self.seen_open[side] or angle >= .07)
if angle >= .27 and not holding:
verified = False
state = ('closed_holding' if holding else 'wide_open' if angle >= .27
else 'preclose_open' if angle >= .07 and not verified
else 'release_started' if angle >= .07
else 'closed' if verified else 'initial_closed_unverified')
phase = ('closed' if holding or (angle < .07 and verified)
else 'release_started' if verified else 'wide_open' if angle >= .27
else 'preclose_open' if angle >= .07 else 'initial_closed_unverified')
closes = np.flatnonzero(values < .07)
opens = np.flatnonzero(values >= .27)
close16 = closes[closes < self.policy['lookahead_rows']]
open16 = opens[opens < self.policy['lookahead_rows']]
prefix_close = np.flatnonzero(adopted < .07)
prefix_open = np.flatnonzero(adopted >= .27)
open_then_close = bool(len(prefix_open) and np.any(adopted[prefix_open[0]+1:] < .07))
pending_before = self.pending[side]
latch_before = self.latches[side]
retry_preclose = bool(verified and not holding and .07 <= angle < .27
and len(closes) and not len(prefix_close) and not len(prefix_open)
and not (pending_before and pending_before['direction'] == 'open')
and not (latch_before and latch_before['direction'] == 'open')
and (not len(opens) or closes[0] < opens[0]))
if retry_preclose:
phase = 'preclose_retry'
hand = dict(measured_angle_rad=angle, contact_force_hold=holding,
source_state=state, source_phase=phase, verified_close_history=verified, retry_after_verified_close=retry_preclose,
close_prediction_rows_1based=(closes + 1).tolist(),
open_prediction_rows_1based=(opens + 1).tolist(),
original_selected_row_rad=float(adopted[self.rows-1]), selected_row_rad=float(adopted[self.rows-1]),
model_row_rad=float(values[self.rows-1]), command_row_1based=self.rows, event='none', override=False,
scheduled_in_adopted_plan=False, calibrated_clipping=False,
override_rows_1based=[], latch_event='none', close_lookahead_enabled=self.close_lookahead)
hand[f'original_selected_row{self.rows}_rad'] = float(adopted[self.rows-1])
hand[f'selected_row{self.rows}_rad'] = float(adopted[self.rows-1])
hand[f'model_row{self.rows}_rad'] = float(values[self.rows-1])
def apply(rows, requested):
lo, hi = (float(limits[side][k]) for k in ('min_angle_rad', 'max_angle_rad'))
if not np.isfinite([lo, hi]).all() or lo > hi:
raise ValueError('Invalid measured gripper limits')
if not lo <= requested <= hi and not allow_clipping:
raise ValueError(side + ': pen transition target outside calibrated limits')
target = float(np.clip(requested, lo, hi))
selected[rows, col] = target
hand.update(override=True, requested_target_rad=float(requested),
calibrated_clipping=bool(hand['calibrated_clipping'] or target != requested))
hand['override_rows_1based'].extend((np.arange(self.rows)[rows] + 1).tolist())
return target
pending = self.pending[side]
latch = self.latches[side]
if not self.close_lookahead:
if pending is not None and pending['direction'] == 'close':
pending = None
hand['discarded_close_reservation'] = True
if latch is not None and latch['direction'] == 'close':
latch = None
hand['discarded_close_latch'] = True
if latch is not None:
direction = latch['direction']
remaining = closes if direction == 'close' else opens
catchup = prefix_close if direction == 'close' else prefix_open
countermand = bool(len(prefix_open)) if direction == 'close' else open_then_close
if countermand:
hand['latch_event'] = 'released_explicit_countermand'
elif not len(remaining):
hand['latch_event'] = 'released_forecast_withdrawn'
elif latch['completed_chunks'] >= self.policy['latch_max_completed_chunks']:
hand['latch_event'] = 'released_bounded_duration'
else:
stop = int(catchup[0]) if len(catchup) else self.rows
if stop:
apply(slice(0, stop), latch['target_rad'])
hand['latch_event'] = 'natural_prefix_catchup' if len(catchup) else 'maintain_short_latch'
if not len(catchup):
next_latches[side] = dict(latch, completed_chunks=latch['completed_chunks'] + 1)
hand['event'] = 'existing_latch'
elif pending is not None:
direction = pending['direction']
hand.update(direction=direction, detected_cycle=pending['detected_cycle'], due_cycle=pending['due_cycle'])
if cycle != pending['due_cycle']:
hand['event'] = 'cancelled_missed_due_cycle'
elif direction == 'close' and (holding or (angle < .07 and verified)):
hand['event'] = 'already_closed_measured'
elif direction == 'open' and state == 'wide_open':
hand['event'] = 'already_open_measured'
elif direction == 'close' and phase not in ('preclose_open', 'preclose_retry'):
hand['event'] = 'cancelled_source_state_changed'
elif direction == 'open' and not verified:
hand['event'] = 'cancelled_source_state_changed'
elif not len(closes if direction == 'close' else opens):
hand['event'] = 'cancelled_prediction_retracted_32'
elif direction == 'close' and np.any(adopted > pending['anchor_angle_rad'] + .02):
hand['event'] = 'cancelled_opening_command_before_deadline'
elif direction == 'open' and (open_then_close or (phase == 'release_started' and len(prefix_close))):
hand['event'] = 'cancelled_closing_command_before_deadline'
elif direction == 'close' and len(prefix_close):
hand['event'] = 'normal_policy_close_in_chunk'
elif direction == 'open' and len(prefix_open):
hand['event'] = 'normal_policy_open_in_chunk'
else:
requested = 0. if direction == 'close' else float(values[opens[0]])
actual = apply(slice(self.rows-1, self.rows), requested)
hand['event'] = 'apply_due_transition'
next_latches[side] = dict(direction=direction, target_rad=actual, completed_chunks=0)
else:
direction = None
if self.close_lookahead and not self.close_defer_only and phase in ('preclose_open', 'preclose_retry') and len(close16):
if np.any(adopted > angle + .02):
hand['event'] = 'not_armed_opening_command'
elif len(prefix_close):
hand['event'] = 'normal_policy_close_in_chunk'
else:
direction = 'close'
elif not self.place_advance_rows and verified and (holding or angle < .27) and len(open16):
if open_then_close:
hand['event'] = 'not_armed_open_then_close'
elif len(prefix_open):
hand['event'] = 'normal_policy_open_in_chunk'
else:
direction = 'open'
if direction:
next_pending[side] = dict(direction=direction, detected_cycle=cycle,
due_cycle=cycle+1, anchor_angle_rad=angle)
hand.update(event='arm_next_chunk', direction=direction,
detected_cycle=cycle, due_cycle=cycle+1)
if self.close_defer_only:
first_close = int(closes[0])+1 if len(closes) else None
check = dict(enabled=True, first_close_row=first_close, near_rows=32,
confirmations_required=1, previous=deepcopy(self.close_defer_state[side]),
modified=False, event='ineligible')
eligible = bool(phase in ('preclose_open','preclose_retry') and not holding
and len(closes) and not len(prefix_close) and not len(prefix_open)
and not (pending_before and pending_before['direction']=='open')
and not (latch_before and latch_before['direction']=='open')
and not np.any(adopted > angle+.02)
and (not len(opens) or closes[0] < opens[0]))
if not eligible or latch is not None:
next_close_defer[side] = None
check['event'] = 'natural_close_in_prefix' if len(prefix_close) else 'reset_ineligible_or_latch'
else:
old = self.close_defer_state[side]
deferred = deepcopy(old) if old is not None else None
if deferred is None and first_close <= 32:
deferred = dict(deadline=self.executed_row_clock+first_close, confirmations=0)
check['event'] = 'reserve_first_close_deadline'
if deferred is not None:
expected = deferred['deadline']-self.executed_row_clock
expired = self.executed_row_clock >= deferred['deadline']
near = first_close <= 32
delayed = first_close >= expected+2
confirm = bool(expired and near and delayed)
deferred['confirmations'] = deferred['confirmations']+1 if confirm else 0
check.update(deadline=deferred['deadline'], expected_close_row=expected,
deadline_expired=expired, near_close=near, delayed=delayed,
confirmations=deferred['confirmations'], event='observe_close_deferral')
if deferred['confirmations'] >= 1:
actual = apply(slice(self.rows-1,self.rows),0.)
hand.update(event='apply_confirmed_deferred_close',direction='close')
check.update(event='apply_confirmed_deferred_close',modified=True,
command_row_1based=self.rows,requested_target_rad=0.,selected_target_rad=actual)
next_pending[side] = None
next_latches[side] = dict(direction='close',target_rad=actual,completed_chunks=0)
deferred = None
next_close_defer[side] = deferred
hand['close_defer'] = check
if self.place_advance_rows:
advance = dict(enabled=True, advance_rows=self.place_advance_rows,
forecast_open_row=None if not len(opens) else int(opens[0])+1,
original_deadline=self.release_deadlines[side], target_row=None,
modified=False, event='not_in_release_state')
if not len(opens) or open_then_close:
next_release_deadlines[side] = None
advance['event'] = 'forecast_withdrawn' if not len(opens) else 'prefix_open_then_close'
elif latch is not None and latch['direction'] == 'open':
next_release_deadlines[side] = None
advance['event'] = 'existing_open_latch'
elif verified and (holding or angle < .27):
first_open = int(opens[0])+1
deadline = self.release_deadlines[side]
if deadline is None:
deadline = self.executed_row_clock + first_open
next_release_deadlines[side] = deadline
expected = min(deadline-self.executed_row_clock, first_open)
target_row = max(1, expected-self.place_advance_rows)
advance.update(original_deadline=deadline, expected_open_row=expected,
target_row=target_row, event='reserved_open_deadline')
if target_row <= self.rows and target_row < first_open:
requested = float(values[opens[0]])
actual = apply(slice(target_row-1,self.rows), requested)
advance.update(modified=True,event='apply_advanced_open',requested_target_rad=requested,
selected_target_rad=actual, previous_transition_event=hand['event'])
hand.update(event='apply_advanced_open',direction='open')
next_pending[side] = None
next_latches[side] = dict(direction='open',target_rad=actual,completed_chunks=0)
next_release_deadlines[side] = None
elif len(prefix_open):
advance['event'] = 'natural_open_in_prefix'
next_release_deadlines[side] = None
else:
next_release_deadlines[side] = None
hand['place_advance'] = advance
final = float(selected[self.rows-1, col])
next_verified[side] = bool(verified or (final < .07 and seen_open))
if final >= .27:
next_verified[side] = False
next_seen[side] = seen_open
next_phases[side] = ('closed' if next_verified[side] and final < .07
else 'release_started' if next_verified[side]
else 'wide_open' if final >= .27 else 'preclose_open' if final >= .07
else 'initial_closed_unverified')
hand.update(selected_row_rad=final, phase_after_selected_chunk=next_phases[side],
verified_close_after_selected_chunk=next_verified[side],
override_rows_1based=sorted(set(hand['override_rows_1based'])))
hand[f'selected_row{self.rows}_rad'] = final
report['hands'][side] = hand
report.update(proposed_close_defer=next_close_defer,
proposed_release_deadlines=next_release_deadlines,
proposed_pending=next_pending, proposed_phases=next_phases,
proposed_verified_close=next_verified, proposed_seen_open=next_seen,
proposed_latches=next_latches)
return selected, report
def annotate_plan(self, report, adopted_rows):
report['adopted_rows'] = adopted_rows
for hand in report['hands'].values():
hand['scheduled_in_adopted_plan'] = bool(hand['override'] and adopted_rows == self.rows)
def commit(self, report, *, executed_rows):
if (report.get('_owner') != self._owner or report.get('_generation') != self._generation
or report.get('state_committed') or report['cycle'] <= self.last_committed_cycle):
raise ValueError('Pen transition proposal already committed/stale/foreign')
if type(executed_rows) is not int or not 0 <= executed_rows <= self.rows:
raise ValueError('Invalid verified executed rows')
full = executed_rows == self.rows
if full:
self.close_defer_state = deepcopy(report['proposed_close_defer'])
self.release_deadlines = deepcopy(report['proposed_release_deadlines'])
self.pending = deepcopy(report['proposed_pending'])
self.phases = deepcopy(report['proposed_phases'])
self.verified_close = deepcopy(report['proposed_verified_close'])
self.seen_open = deepcopy(report['proposed_seen_open'])
self.latches = deepcopy(report['proposed_latches'])
else:
self.close_defer_state = {s: None for s in SIDES}
self.release_deadlines = {s: None for s in SIDES}
self.pending = {s: None for s in SIDES}
self.latches = {s: None for s in SIDES}
self.executed_row_clock += executed_rows
self.last_committed_cycle = report['cycle']
self._generation += 1
report.update(state_committed=True, executed_rows=executed_rows,
commit_result='full_chunk' if full else 'shortened_drop_pending',
close_defer_after=deepcopy(self.close_defer_state),
release_deadlines_after=deepcopy(self.release_deadlines),
executed_row_clock_after=self.executed_row_clock,
pending_after=deepcopy(self.pending), phases_after=deepcopy(self.phases),
latches_after=deepcopy(self.latches), verified_close_after=deepcopy(self.verified_close))
ARM_POLICY = dict(
kind="u1_pen_grasp_axis_pose_retention_v9",
correction_scale=1.25, approach_axis_hand=[1., 0., 0.],
grasp_width_min_rad=3 * 2 * np.pi / 4096,
closed_command_threshold_rad=.07, raw_forecast_horizon_rows=32,
progress_epsilon_m=1e-6, reference_max_translation_m=.005,
reference_max_rotation_rad=.05,
coordinate_reference="opposite hand, assumed stationary; contract inter-hand pose only",
direction="negative active-hand X at width-gated grasp, frozen in opposite-hand frame",
translation="configured scale times the old/latest target deficit along estimated up only",
history_source="raw predictions, never corrected predictions",
orientation="retain missing rotation progress around a frozen forecast axis; same correction scale as height",
rotation_axis_min_angle_rad=np.pi/180, rotation_progress_epsilon_rad=1e-5,
rotation_blend_in="smoothstep over executed prefix; full deficit at prefix endpoint",
grippers="unchanged; any independent client lookahead stays separate",
fixed_postgrasp_row_cutoff=False, world_inputs=False, robot_fk_inputs=False,
calibrated_transform_inputs=False, current_or_contact_inputs=False,
lift_success_classifier=False,
)
def regime(x):
return "closed" if x < .07 else "open" if x >= .27 else "narrow"
def pair(x):
a = np.asarray([x[s] for s in SIDES] if isinstance(x, dict) else x, dtype=float)
if a.shape != (2,) or not np.isfinite(a).all():
raise ValueError("Two finite finger values required")
return a
def matrix(pose):
p = np.asarray(pose, dtype=float)
if p.shape != (7,) or not np.isfinite(p).all():
raise ValueError("Finite xyz/xyzw pose required")
q = p[3:]; norm = np.linalg.norm(q)
if abs(norm - 1.) > .02:
raise ValueError("Unit xyzw quaternion required")
x, y, z, w = q / norm
t = np.eye(4)
t[:3, :3] = [[1-2*(y*y+z*z), 2*(x*y-z*w), 2*(x*z+y*w)],
[2*(x*y+z*w), 1-2*(x*x+z*z), 2*(y*z-x*w)],
[2*(x*z-y*w), 2*(y*z+x*w), 1-2*(x*x+y*y)]]
t[:3, 3] = p[:3]
return t
def inverse(t):
result = np.eye(4)
result[:3, :3] = t[:3, :3].T
result[:3, 3] = -result[:3, :3] @ t[:3, 3]
return result
def angle(r):
return float(np.arccos(np.clip((np.trace(r)-1.)/2., -1., 1.)))
def quaternion(r):
trace=float(np.trace(r))
if trace > 0:
s=2*np.sqrt(max(0.,trace+1.))
q=np.array([(r[2,1]-r[1,2])/s,(r[0,2]-r[2,0])/s,(r[1,0]-r[0,1])/s,s/4])
else:
i=int(np.argmax(np.diag(r)));j=(i+1)%3;k=(i+2)%3
s=2*np.sqrt(max(0.,1.+r[i,i]-r[j,j]-r[k,k]))
q=np.zeros(4);q[i]=s/4;q[j]=(r[i,j]+r[j,i])/s;q[k]=(r[i,k]+r[k,i])/s
q[3]=(r[k,j]-r[j,k])/s
q/=np.linalg.norm(q)
return -q if q[3]<0 else q
def axis_rotation(axis, radians):
x,y,z=axis
skew=np.array([[0.,-z,y],[z,0.,-x],[-y,x,0.]])
return np.eye(3)+np.sin(radians)*skew+(1-np.cos(radians))*(skew@skew)
def twist_angle(r, axis, reference=0.):
q=quaternion(r);v=float(q[:3]@axis);w=float(q[3])
if np.hypot(v,w)<1e-8:
raise ValueError("Rotation twist is singular")
value=2*np.arctan2(v,w)
return float(reference+(value-reference+np.pi)%(2*np.pi)-np.pi)
def twist_path(rotations, base, axis, reference):
result=[]
for r in rotations:
reference=twist_angle(r@base.T,axis,reference);result.append(reference)
return np.asarray(result)
def forecast(raw, starts):
poses = np.empty((len(raw), 2, 4, 4))
for h in range(2):
p = starts[h].copy()
for i, row in enumerate(raw):
p = p @ matrix(row[h*7:h*7+7])
poses[i, h] = p
return poses
class GraspAxisRetention:
def __init__(self, execution_rows=None, *, correction_scale=None):
if execution_rows is not None and (type(execution_rows) is not int
or execution_rows not in (8, 16)):
raise ValueError("Execution rows must be 8, 16, or None")
scale = ARM_POLICY["correction_scale"] if correction_scale is None else correction_scale
if isinstance(scale, (bool, np.bool_)) or not isinstance(scale, (int, float, np.integer, np.floating)) or not np.isfinite(scale) or not 0 <= scale <= 2:
raise ValueError("correction_scale must be a finite number in [0,2]")
self.scale = float(scale)
self.n = execution_rows
self.policy = dict(ARM_POLICY, execution_rows=execution_rows, correction_scale=self.scale,
max_old_age_chunks=None if execution_rows is None else 32//execution_rows-1)
self.reset()
def reset(self):
self.history = []
self.clock = self.generation = 0
self.held_last = [False, False]
self.held_rows = [0, 0]
self.grasp_directions = [None, None]
self.retention_done = [False, False]
self.rotation_bases = [None, None]
self.rotation_axes = [None, None]
self.rotation_retention_done = [False, False]
self.rotation_progress = [0., 0.]
self.observed_execution_rows = None
self.pending = None
self.last_prompt = None
self.last_report = None
def propose(self, raw, obs, targets=None, previous_applied=None):
raw = np.asarray(raw)
if raw.shape != (32, 16) or raw.dtype.kind != "f" or not np.isfinite(raw).all():
raise ValueError("Finite floating (32,16) raw predictions required")
relative = matrix(obs[POSE_KEY])
current = [inverse(relative), relative]
measured = pair(obs["observation.joint_states"])
previous = measured if previous_applied is None else pair(previous_applied)
n = self.n or 32
selected = np.asarray(raw[:n, 14:16] if targets is None else targets, dtype=float)
if selected.shape != (n, 2) or not np.isfinite(selected).all():
raise ValueError("Selected gripper targets must match the returned prefix")
poses = forecast(raw, current)
relative_motion = forecast(raw, [np.eye(4), np.eye(4)])
ref_rows = self.n or 16
stable = []
for h in range(2):
other = 1-h
translation = np.max(np.linalg.norm(relative_motion[:ref_rows, other, :3, 3], axis=1))
rotation = max(angle(t[:3, :3]) for t in relative_motion[:ref_rows, other])
stable.append(bool(translation <= ARM_POLICY["reference_max_translation_m"]
and rotation <= ARM_POLICY["reference_max_rotation_rad"]))
held, directions, done, resets, counts, reasons = [], [], [], [], [], []
for h in range(2):
gate = selected[:, h] if self.n else selected[:1, h]
sequence = [regime(previous[h])] + [regime(v) for v in gate]
transition = any(a != b for a, b in zip(sequence, sequence[1:]))
closing = all(x == "closed" for x in sequence)
width = ARM_POLICY["grasp_width_min_rad"]-1e-12 <= measured[h] < .07
active = bool(closing and width and stable[h])
held.append(active); resets.append(transition)
counts.append(self.held_rows[h] if active and self.held_last[h] else 0)
direction = self.grasp_directions[h] if active and self.held_last[h] else None
if active and direction is None:
direction = -current[h][:3, :3][:, 0].copy()
directions.append(direction)
done.append(self.retention_done[h] if active and self.held_last[h] else False)
why = []
if transition: why.append("selected_grip_transition")
if not closing or not width: why.append("width_gate_inactive")
if not stable[h]: why.append("opposite_hand_prediction_moves")
if done[h]: why.append("postgrasp_lift_phase_complete")
reasons.append(why)
output = raw.copy(); reports = {}; rotation_states = []
for h, side in enumerate(SIDES):
direction = directions[h]; start = current[h][:3, 3]
fresh = poses[:n, h, :3, 3]
additions = np.zeros(n); best = 0.; chosen = None
pending_forward = False; old_peak = 0.; fresh_peak = 0.
if held[h] and not done[h]:
progress = (fresh-start) @ direction
closed_end = next((i for i in range(32) if regime(raw[i, 14+h]) != "closed"), 32)
if closed_end:
fresh_peak = max(0., float(np.max((poses[:closed_end, h, :3, 3]-start) @ direction)))
for entry in self.history:
offset = self.clock-entry["birth"]
if resets[h] or not entry["valid"][h] or not 1 <= offset < 32:
continue
length = min(n, 32-offset)
if self.n and length != n: continue
mismatch = next((i for i in range(length)
if regime(entry["grips"][offset+i, h]) != regime(selected[i, h])
or regime(selected[i, h]) != "closed"), length)
if self.n and mismatch != length:
reasons[h].append("old_grip_regime_mismatch"); continue
length = mismatch
if not length: continue
old_progress = (entry["poses"][offset:offset+length, h, :3, 3]-start) @ direction
peak = float(np.max(old_progress)); old_peak = max(old_peak, peak)
if peak <= ARM_POLICY["progress_epsilon_m"]: continue
pending_forward = True
delta = np.zeros(n)
delta[:length] = np.maximum(np.maximum(old_progress, 0.)-progress[:length], 0.) * self.scale
amount = float(np.max(delta))
if amount > best:
best, additions, chosen = amount, delta, (entry, offset)
end = min(n, closed_end)
if (not pending_forward and fresh_peak <= ARM_POLICY["progress_epsilon_m"]
and end and progress[end-1] < -ARM_POLICY["progress_epsilon_m"]):
done[h] = True
reasons[h].append("no_unexecuted_upward_forecast; return_to_latest_model")
modified = False
rotation=self._rotation_proposal(h,current,poses,selected,held[h],resets[h],raw[:,14+h])
rotation_states.append(rotation["state"])
if chosen is not None or rotation["modified"]:
prior_p = start; prior_r = current[h][:3, :3]
for i in range(n):
retained = fresh[i] if chosen is None else fresh[i] + additions[i]*direction
output[i, h*7:h*7+3] = prior_r.T @ (retained-prior_p)
target_r=rotation["rotations"][i]
if rotation["modified"]:
q=quaternion(prior_r.T@target_r)
if q@raw[i,h*7+3:h*7+7]<0:q=-q
output[i,h*7+3:h*7+7]=q
prior_p, prior_r = retained, target_r
modified = chosen is not None and best>0
reports[side] = dict(
blended=False, translation_retained=modified, grasp_by_width=held[h],
measured_width_rad=float(measured[h]), previous_applied_command_rad=float(previous[h]),
width_threshold_rad=ARM_POLICY["grasp_width_min_rad"], postgrasp_executed_rows=counts[h],
reset_reasons=reasons[h], retention_phase_complete=done[h],
reference_hand=SIDES[1-h], reference_assumed_stationary=stable[h],
direction_frame="opposite hand at grasp, assumed stationary",
correction_direction=None if direction is None else direction.tolist(),
correction_direction_hand=None if direction is None else (current[h][:3, :3].T @ direction).tolist(),
pending_old_forward_forecast=pending_forward,
old_forward_peak_m=old_peak, fresh_forward_peak_m=fresh_peak,
old_sources=[] if chosen is None else [dict(generation=chosen[0]["generation"],
age_chunks=self.generation-chosen[0]["generation"], consumed_rows=chosen[1])],
added_forward_m=additions.tolist() if chosen is not None else [],
max_target_shift_m=best, **rotation["report"])
assert np.array_equal(output[:, 14:], raw[:, 14:])
for h in range(2):
if not reports[SIDES[h]]["rotation_retained"]:
assert np.array_equal(output[:, h*7+3:h*7+7], raw[:, h*7+3:h*7+7])
assert np.array_equal(output[n:], raw[n:])
report = dict(enabled=True, policy=self.policy, generation=self.generation,
executed_row_clock=self.clock, modified=not np.array_equal(output, raw),
hands=reports, retained_prefix_rows=n, committed=False,
future_rows="unchanged raw predictions" if self.n else "finite old-tail overlap only")
return dict(actions=output, report=report, generation=self.generation,
poses=poses, grips=raw[:, 14:].copy(), selected=selected.copy(),
reset=resets, held=held, held_rows=counts, directions=directions, phase_done=done,
rotation_states=rotation_states,
relative_path=self._relative_path(relative, output))
def _rotation_proposal(self,h,current,poses,selected,held,reset,raw_grips):
n=self.n or 32;fresh=poses[:n,h,:3,:3]
carry=held and self.held_last[h]
base=self.rotation_bases[h] if carry else current[h][:3,:3].copy() if held else None
axis=self.rotation_axes[h] if carry else None
complete=self.rotation_retention_done[h] if carry else False
previous=self.rotation_progress[h] if carry else 0.
state=dict(base=base,axis=axis,done=complete,progress=previous)
report=dict(rotation_retained=False,rotation_phase_complete=complete,
rotation_axis_frame="opposite hand, assumed stationary",rotation_axis=None,
rotation_progress_rad=previous,added_rotation_rad=[],rotation_old_sources=[],
max_target_rotation_change_rad=0.,rotation_reset_reasons=[])
result=dict(rotations=fresh.copy(),modified=False,state=state,report=report)
if not held:return result
if complete:
report["rotation_reset_reasons"].append("postgrasp_rotation_phase_complete")
return result
closed_end=next((i for i,v in enumerate(raw_grips) if regime(v)!="closed"),32)
if axis is None:
candidates=[quaternion(r@base.T) for r in poses[:closed_end,h,:3,:3]]
if candidates:
q=max(candidates,key=lambda q:2*np.arctan2(np.linalg.norm(q[:3]),q[3]))
magnitude=2*np.arctan2(np.linalg.norm(q[:3]),q[3])
if magnitude>=ARM_POLICY["rotation_axis_min_angle_rad"]:
axis=q[:3]/np.linalg.norm(q[:3]);state["axis"]=axis
if axis is None:
report["rotation_reset_reasons"].append("waiting_for_rotation_forecast_axis")
return result
report["rotation_axis"]=axis.tolist()
try:
present=twist_angle(current[h][:3,:3]@base.T,axis,previous)
fresh_angles=twist_path(poses[:,h,:3,:3],base,axis,present)
state["progress"]=present;report["rotation_progress_rad"]=present
best=0.;chosen=None;additions=np.zeros(n);pending=False
old_peak=0.;epsilon=ARM_POLICY["rotation_progress_epsilon_rad"]
for entry in self.history:
offset=self.clock-entry["birth"]
if reset or not entry["valid"][h] or not 1<=offset<32:continue
length=min(n,32-offset)
if self.n and length!=n:continue
mismatch=next((i for i in range(length) if regime(entry["grips"][offset+i,h])!=regime(selected[i,h])
or regime(selected[i,h])!="closed"),length)
if self.n and mismatch!=length:continue
length=mismatch
if not length:continue
old_angles=twist_path(entry["poses"][:,h,:3,:3],base,axis,present)
progress=old_angles[offset:offset+length]-present
peak=float(np.max(progress));old_peak=max(old_peak,peak)
if peak<=epsilon:continue
pending=True;delta=np.zeros(n)
delta[:length]=np.maximum(np.maximum(progress,0.)-(fresh_angles[:length]-present),0.)*self.scale
amount=float(np.max(delta))
if amount>best:best,chosen,additions=amount,(entry,offset),delta
fresh_peak=max(0.,float(np.max(fresh_angles[:closed_end]-present))) if closed_end else 0.
end=min(n,closed_end)
if not pending and fresh_peak<=epsilon and end and fresh_angles[end-1]-present<=epsilon:
state["done"]=True;report["rotation_phase_complete"]=True
report["rotation_reset_reasons"].append("no_unexecuted_rotation_forecast; return_to_latest_model")
report.update(pending_old_rotation_forecast=pending,old_rotation_peak_rad=old_peak,
fresh_rotation_peak_rad=fresh_peak)
if chosen is not None:
if best>=np.pi:raise ValueError("Rotation correction reaches ambiguous half-turn")
blend_rows=self.n or self.observed_execution_rows or 16
t=np.minimum(np.arange(1,n+1,dtype=float)/blend_rows,1.)
unramped=additions.copy()
additions=additions*(3*t*t-2*t*t*t)
for i in range(n):result["rotations"][i]=axis_rotation(axis,additions[i])@fresh[i]
result["modified"]=True
report.update(rotation_retained=True,added_rotation_rad=additions.tolist(),
unramped_added_rotation_rad=unramped.tolist(),rotation_blend_rows=blend_rows,
max_target_rotation_change_rad=float(np.max(additions)),rotation_old_sources=[dict(
generation=chosen[0]["generation"],age_chunks=self.generation-chosen[0]["generation"],consumed_rows=chosen[1])])
except ValueError as exc:
state["done"]=True;report["rotation_phase_complete"]=True
report["rotation_reset_reasons"].append(str(exc))
return result
@staticmethod
def _relative_path(relative, output):
p = forecast(output, [np.eye(4), relative])
return np.asarray([inverse(row[0]) @ row[1] for row in p])
def commit(self, proposal, executed_rows, applied=None, *, verified=True):
if proposal["generation"] != self.generation:
raise ValueError("Stale grasp-axis proposal")
limit = self.n or 32
if type(executed_rows) is not int or not 0 <= executed_rows <= limit:
raise ValueError("Invalid executed row count")
if not executed_rows or (self.n and executed_rows != self.n):
self.history.clear(); self.held_last = [False, False]; self.held_rows = [0, 0]
self.grasp_directions = [None, None]; self.retention_done = [False, False]
self.rotation_bases = [None, None]; self.rotation_axes = [None, None]
self.rotation_retention_done = [False, False]; self.rotation_progress = [0., 0.]
self.clock += executed_rows; self.generation += 1
proposal["report"].update(committed=True, executed_rows=executed_rows, history_reset="shortened_prefix")
return
end = [regime(v) for v in proposal["selected"][executed_rows-1]]
valid = [x == "closed" and proposal["report"]["hands"][side]["reference_assumed_stationary"]
for side,x in zip(SIDES,end)]
if applied is not None:
actual = pair(applied)
valid = [v and regime(a) == e for v, a, e in zip(valid, actual, end)]
reset = list(proposal["reset"])
if self.n is None:
for h in range(2):
if any(regime(x) != "closed" for x in proposal["selected"][:executed_rows, h]):
reset[h] = True; proposal["held"][h] = False
kept = []
for entry in self.history:
entry = dict(entry, valid=[a and b and not r and h for a,b,r,h in
zip(entry["valid"], valid, reset, proposal["held"])])
if any(entry["valid"]) and self.clock+executed_rows-entry["birth"] < 32:
kept.append(entry)
kept.append(dict(birth=self.clock, generation=self.generation,
poses=proposal["poses"].copy(), grips=proposal["grips"].copy(), valid=valid))
self.history = [e for e in kept if any(e["valid"])]
self.held_last = [h and v for h,v in zip(proposal["held"], valid)]
self.held_rows = [c+executed_rows if h else 0 for c,h in zip(proposal["held_rows"],self.held_last)]
self.grasp_directions = [None if not h or d is None else d.copy()
for h,d in zip(self.held_last,proposal["directions"])]
self.retention_done = [bool(h and d) for h,d in zip(self.held_last,proposal["phase_done"])]
self.rotation_bases=[s["base"].copy() if held and s["base"] is not None else None
for held,s in zip(self.held_last,proposal["rotation_states"])]
self.rotation_axes=[s["axis"].copy() if held and s["axis"] is not None else None
for held,s in zip(self.held_last,proposal["rotation_states"])]
self.rotation_retention_done=[bool(held and s["done"]) for held,s in zip(self.held_last,proposal["rotation_states"])]
self.rotation_progress=[s["progress"] if held else 0. for held,s in zip(self.held_last,proposal["rotation_states"])]
self.clock += executed_rows; self.generation += 1
proposal["report"].update(committed=True, executed_rows=executed_rows,
endpoint_verified=verified, commit_boundary="verified_client_execution" if verified else "observed_relative_pose_alignment",
postgrasp_executed_rows_after=dict(zip(SIDES,self.held_rows)),
history_generations=[e["generation"] for e in self.history])
def _slow_half_speed(actions):
out = np.repeat(actions, 2, axis=0)
for start in (0, 7):
rotation = Rotation.from_quat(actions[:, start+3:start+7])
half = Rotation.from_rotvec(rotation.as_rotvec() * .5)
translation = actions[:, start:start+3] * .5
out[0::2, start:start+3] = translation
out[1::2, start:start+3] = half.inv().apply(translation)
out[0::2, start+3:start+7] = half.as_quat(canonical=True)
out[1::2, start+3:start+7] = half.as_quat(canonical=True)
return out.astype(np.float32)
def _pa1_pick_compression(actions, measured):
measured = np.asarray(measured, dtype=float)
if measured[1] <= .56 or np.max(np.abs(actions[0, 14:16]-measured)) > .05:
return None
if np.max(np.ptp(actions[:, 14:16], axis=0)) > .03 or np.min(actions[:, 15]) < .56:
return None
distances = np.stack([np.linalg.norm(actions[:, k:k+3], axis=1) for k in (0, 7)], axis=1)
angles = np.stack([Rotation.from_quat(actions[:, k+3:k+7]).magnitude() for k in (0, 7)], axis=1)
if np.max(distances[:4].sum(axis=0)) >= .002 or np.max(angles[:4].sum(axis=0)) >= np.deg2rad(1.):
return None
groups, counts = [], []
index = 0
while index < len(actions) and len(groups) < 8:
end = index+1
length = distances[index].copy()
angle = angles[index].copy()
while end < min(index+8, len(actions)):
next_length = length+distances[end]
next_angle = angle+angles[end]
if np.max(next_length) > .002 or np.max(next_angle) > np.deg2rad(.5):
break
length, angle = next_length, next_angle
end += 1
combined = actions[end-1].astype(np.float64).copy()
for start in (0, 7):
position = np.zeros(3)
rotation = Rotation.identity()
for row in actions[index:end]:
position += rotation.apply(row[start:start+3])
rotation = rotation*Rotation.from_quat(row[start+3:start+7])
combined[start:start+3] = position
combined[start+3:start+7] = rotation.as_quat(canonical=True)
groups.append(combined)
counts.append(end-index)
index = end
consumed = sum(counts[:4])
if consumed <= 4:
return None
while len(groups) < 8:
hold = groups[-1].copy()
hold[:14] = np.tile([0., 0., 0., 0., 0., 0., 1.], 2)
groups.append(hold)
return _slow_half_speed(np.asarray(groups, dtype=np.float32)), consumed
def _cup_tip_pivot(actions, column, release_row, hold_angle):
out = actions.copy()
opening = 2*(release_row-1)
chunk = opening//8*8
stop = max(opening, chunk+4)
start = stop-4
delay = stop-opening
if delay:
out[stop:, column] = actions[opening:len(actions)-delay, column]
out[start:stop, column] = hold_angle
rotation = Rotation.from_rotvec([0., -np.deg2rad(10.)/4., 0.])
pivot = np.array([.09, 0., 0.])
translation = pivot-rotation.apply(pivot)
offset = 0 if column == 14 else 7
out[start:stop, offset:offset+3] = translation
out[start:stop, offset+3:offset+7] = rotation.as_quat(canonical=True)
return out
def _pen_preclose_approach(actions, measured):
closes = np.flatnonzero(actions[:16, 15] < .07)
if not len(closes) or measured[1] < .07:
return None
first = int(closes[0])
stop = max(first, 4)
start = stop-4
delay = stop-first
out = actions.copy()
if delay:
out[stop:, 15] = actions[first:len(actions)-delay, 15]
out[start:stop, 15] = measured[1]
rotation = Rotation.identity()
for i in range(start, stop):
out[i, 7:10] += rotation.inv().apply([.005/4., 0., 0.])
rotation = rotation*Rotation.from_quat(out[i, 10:14])
return out
def _pen_extra_lift(actions, direction):
out = actions.copy()
rotation = Rotation.identity()
for i in range(16):
out[i, 7:10] += rotation.inv().apply(np.asarray(direction) * .0025)
rotation = rotation * Rotation.from_quat(out[i, 10:14])
return out
def _pa2_initial_lift_direction(actions):
position = np.zeros(3)
rotation = Rotation.identity()
for row in actions[:8]:
if row[14] >= .56:
break
position += rotation.apply(row[:3])
rotation = rotation * Rotation.from_quat(row[3:7])
length = np.linalg.norm(position)
return position / length if length > 1e-6 else None
def _pa2_scale_lift(actions, direction, release_row):
out = actions.copy()
rotation = Rotation.identity()
limit = min(16, 2 * (release_row - 1)) if release_row is not None else 16
finished = False
for i in range(limit):
if out[i, 14] >= .56:
finished = True
break
projected = float(rotation.apply(out[i, :3]) @ direction)
if projected < -1e-5:
finished = True
break
if projected > 0:
out[i, :3] += rotation.inv().apply(.5 * projected * direction)
rotation = rotation * Rotation.from_quat(out[i, 3:7])
return out, finished or release_row is not None
class Policy:
def __init__(self, checkpoint_dir: str):
from openpi.models import pi0_config
from openpi.policies import policy_config
from openpi.training.config import AssetsConfig, DataConfig, TrainConfig
checkpoint = pathlib.Path(checkpoint_dir)
if not (checkpoint / "params").is_dir():
checkpoint = checkpoint / "step_042000"
config = TrainConfig(
name="u1_42000_submission",
model=pi0_config.Pi0Config(pi05=True, action_dim=32, action_horizon=32),
data=_build_transforms()(
assets=AssetsConfig(asset_id=_asset_id(checkpoint)),
base_config=DataConfig(prompt_from_task=True)),
)
self._model = policy_config.create_trained_policy(config, str(checkpoint))
self.reset()
rng = self._model._rng
image = np.random.default_rng(0).integers(0, 256, (480, 640, 3), dtype=np.uint8)
obs = {_K[0]: image, _K[1]: image.copy(),
POSE_KEY: np.array([.23, -.05, .27, -.58, -.09, .72, .34], np.float32),
_K[3]: np.array([.24, .78], np.float32)}
started = time.monotonic()
try:
for prompt in (PEN_PROMPTS[0], CUP_PA1, CUP_PA2):
self._sample(dict(obs, prompt=prompt), prompt)
finally:
self._model._rng = rng
self.reset()
print(f"[policy] pre-warmed in {time.monotonic()-started:.1f}s", flush=True)
def reset(self):
self._prompt = None
self._second_half = None
self._pending = None
self._cycle = 0
self._previous_grip = None
self._cup = None
self._cup_seen_open = False
self._cup_grasped = False
self._cup_tilted = False
self._pa2_lift_axis = None
self._pa2_lift_done = False
self._pen_approach_done = False
self._pen_lift_added = False
self._pen = PenGraspSchedule(8, close_defer_only=True)
self._arm = GraspAxisRetention(execution_rows=8, correction_scale=1.0)
def _sample(self, obs, prompt):
clean = {key: obs[key] for key in _K if key in obs}
mapped = prompt
if prompt.startswith("Pick 2nd ") and " part and place into cell" in prompt:
mapped = "Pick 1st " + prompt[len("Pick 2nd "):]
clean["prompt"] = mapped
sigma = CUP_NOISE.get(prompt, 1.0)
kwargs = {}
if sigma != 1.0:
import jax
_, key = jax.random.split(self._model._rng)
noise = jax.random.normal(key, (1, 32, 32))
kwargs["noise"] = noise * sigma
actions = np.asarray(self._model.infer(clean, **kwargs)["actions"], dtype=np.float32)
if actions.shape != (32, 16) or not np.isfinite(actions).all():
raise ValueError("Expected finite float32 model actions [32,16]")
return actions.copy()
def _finish_previous_pair(self):
if self._pending is None:
return
kind, proposals, applied = self._pending
if kind == "cup":
self._cup.commit(proposals, proposals['adopt_rows'])
else:
transition, arm = proposals
self._pen.commit(transition, executed_rows=8)
self._arm.commit(arm, 8, applied=applied, verified=False)
self._previous_grip = applied.copy()
self._pending = None
self._cycle += 1
def infer(self, obs: dict) -> np.ndarray:
prompt = obs.get("prompt", "").strip().rstrip(".")
if prompt != self._prompt:
self.reset()
self._prompt = prompt
if prompt in CUP_HANDS:
self._cup = CupContinuousAdvance(CUP_ADVANCE[prompt])
if self._second_half is not None:
result = self._second_half
self._second_half = None
return result.copy()
self._finish_previous_pair()
raw = self._sample(obs, prompt)
tip_release = None
lift_direction = None
if prompt in CUP_HANDS:
finger = float(obs[_K[3]][CUP_HANDS[prompt][1]-14])
if finger > .53:
self._cup_seen_open = True
elif self._cup_seen_open and finger >= 3*2*np.pi/4096:
self._cup_grasped = True
if prompt == CUP_PA1 and not self._cup_grasped:
compressed = _pa1_pick_compression(raw, obs[_K[3]])
if compressed is not None:
output, consumed = compressed
proposal = self._cup.propose(raw, dict(obs, prompt=prompt), consumed)
self._pending = ('cup', proposal, None)
return output.copy()
proposal = self._cup.propose(raw, dict(obs, prompt=prompt), 8)
controlled = proposal["actions"]
self._pending = ("cup", proposal, None)
tip_release = proposal["report"]["forced_start_row"] or proposal["report"]["natural_start_row"]
elif prompt in PEN_PROMPTS:
measured = np.asarray(obs[_K[3]], dtype=float)
if measured[1] >= .27 and np.any(raw[:8, 15] >= .27):
self._pen_approach_done = False
self._pen_lift_added = False
limits = {side: dict(min_angle_rad=-1e10, max_angle_rad=1e10) for side in SIDES}
targets, transition = self._pen.propose(
self._cycle, raw, raw[:8, 14:16], measured, {}, limits,
allow_clipping=False)
arm = self._arm.propose(raw, obs, targets,
previous_applied=self._previous_grip)
hand = arm["report"]["hands"]["right"]
if not self._pen_lift_added and hand["translation_retained"]:
lift_direction = hand["correction_direction_hand"]
controlled = arm["actions"].copy()
controlled[:8, 14:16] = targets
self._pending = ("pen", (transition, arm), targets[-1].copy())
else:
controlled = raw
slow = _slow_half_speed(controlled)
if prompt == CUP_PA2 and self._cup_grasped and not self._pa2_lift_done:
relative_rotation = Rotation.from_quat(np.asarray(obs[POSE_KEY])[3:])
if self._pa2_lift_axis is None:
direction = _pa2_initial_lift_direction(slow)
if direction is not None:
self._pa2_lift_axis = relative_rotation.inv().apply(direction)
if self._pa2_lift_axis is not None:
direction = relative_rotation.apply(self._pa2_lift_axis)
slow, self._pa2_lift_done = _pa2_scale_lift(slow, direction, tip_release)
if lift_direction is not None:
slow = _pen_extra_lift(slow, lift_direction)
self._pen_lift_added = True
if prompt in PEN_PROMPTS and not self._pen_approach_done:
approached = _pen_preclose_approach(slow, np.asarray(obs[_K[3]], dtype=float))
if approached is not None:
slow = approached
self._pen_approach_done = True
if tip_release is not None and self._cup_grasped and not self._cup_tilted:
column = CUP_HANDS[prompt][1]
slow = _cup_tip_pivot(slow, column, tip_release, float(obs[_K[3]][column-14]))
self._cup_tilted = True
self._second_half = slow[8:24].copy()
return slow[:16].copy()
from copy import deepcopy
import dataclasses
import io
import os
import pathlib
import time
import numpy as np
from PIL import Image
from scipy.spatial.transform import Rotation
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
JPEG_QUALITY = 90
SIDES = ("left", "right")
POSE_KEY = "observation.pose.left_hand_root_to_right_hand_root.absolute"
_K = ("observation.image.left", "observation.image.right", POSE_KEY,
"observation.joint_states", "prompt")
CUP_PA1 = "Pick up the cup with your right hand and set it on the plate"
CUP_PA2 = "Pick up the cup with your left hand and return the cup to its original position"
CUP_HANDS = {CUP_PA1: ("right", 15), CUP_PA2: ("left", 14)}
CUP_ADVANCE = {CUP_PA1: 8, CUP_PA2: 4}
CUP_NOISE = {CUP_PA1: .15, CUP_PA2: .10}
PEN_PROMPTS = (
"Pick up a pencil on the table and put it into the pen holder",
"Pick up a ballpoint pen on the table and put it into the pen holder",
)
def _asset_id(ckpt):
found = sorted(p.parent.name for p in (ckpt / "assets").glob("*/norm_stats.json"))
if len(found) != 1:
raise FileNotFoundError(f"expected one assets/<id>/norm_stats.json in {ckpt}, found {found}")
return found[0]
def _m(r):
a, b = r[..., :3], r[..., 3:6]
c0 = a / np.linalg.norm(a, axis=-1, keepdims=True)
b = b - np.sum(c0 * b, axis=-1, keepdims=True) * c0
c1 = b / np.linalg.norm(b, axis=-1, keepdims=True)
return np.stack([c0, c1, np.cross(c0, c1)], axis=-1)
def _d(x):
rot = Rotation.from_matrix(_m(x[:, 3:9]))
pr = Rotation.concatenate([Rotation.identity(), rot[:-1]])
pp = np.concatenate([np.zeros((1, 3)), x[:-1, :3]])
return np.concatenate([pr.inv().apply(x[:, :3] - pp), (pr.inv() * rot).as_quat(canonical=True)], axis=1)
def _image(img):
from openpi.shared import image_tools
a = np.asarray(image_tools.resize_with_pad(np.asarray(img, dtype=np.uint8), 224, 224))
b = io.BytesIO(); Image.fromarray(a).save(b, format="JPEG", quality=JPEG_QUALITY)
return np.asarray(Image.open(io.BytesIO(b.getvalue())).convert("RGB"))
def _build_transforms():
from openpi import transforms as T
from openpi.training.config import DataConfigFactory, ModelTransformFactory
class _P(T.DataTransformFn):
def __call__(self, d):
return {k: d[k] for k in _K if k in d}
class _I(T.DataTransformFn):
n: int = 32
def __call__(self, d):
l, r = _image(d[_K[0]]), _image(d[_K[1]])
p = np.array(d[_K[2]], dtype=np.float64)
if p[6] < 0:
p[3:] *= -1
s = np.concatenate([p, np.asarray(d[_K[3]], dtype=np.float64)]).astype(np.float32)
o = {"state": T.pad_to_dim(s, self.n),
"image": {"base_0_rgb": np.zeros_like(l), "left_wrist_0_rgb": l, "right_wrist_0_rgb": r},
"image_mask": {"base_0_rgb": np.False_, "left_wrist_0_rgb": np.True_, "right_wrist_0_rgb": np.True_}}
if "prompt" in d:
o["prompt"] = d["prompt"]
return o
class _O(T.DataTransformFn):
def __call__(self, d):
a = np.asarray(d["actions"], dtype=np.float64)
return {"actions": np.concatenate([_d(a[:, 0:9]), _d(a[:, 9:18]), a[:, 18:20]], axis=1).astype(np.float32)}
class _C(DataConfigFactory):
def create(self, assets_dirs, model_config):
return dataclasses.replace(
self.create_base_config(assets_dirs, model_config),
repack_transforms=T.Group(inputs=[_P()]),
data_transforms=T.Group(inputs=[_I(model_config.action_dim)], outputs=[_O()]),