| """YOR-only waypoint filtering and lossless interleaved temporal subsampling.""" |
| import numpy as np |
| from datasets.core import TrajectoryDataset |
|
|
| DEFAULT_THRESHOLDS = dict(left_xyz_m=.005,right_xyz_m=.005,left_rotation_deg=3., |
| right_rotation_deg=3.,left_hand_rad=.08,right_hand_rad=.08,lift_m=.005, |
| base_xy_m=.01,base_yaw_rad=.03) |
| GROUPS = dict(left_xyz_m=slice(4,7),right_xyz_m=slice(11,14), |
| left_hand_rad=slice(14,30),right_hand_rad=slice(34,50),lift_m=slice(54,55), |
| base_xy_m=slice(55,57),base_yaw_rad=slice(57,58)) |
|
|
| def validate_thresholds(thresholds): |
| unknown=set(thresholds)-set(DEFAULT_THRESHOLDS) |
| if unknown: raise ValueError(f'Unknown motion thresholds: {unknown}') |
| for key,t in thresholds.items(): |
| if t is None: continue |
| t=np.asarray(t,float) |
| width=GROUPS[key].stop-GROUPS[key].start if key in GROUPS else 1 |
| if t.ndim>1 or t.size not in (1,width) or not np.isfinite(t).all() or (t<=0).any(): |
| raise ValueError(f'{key}: use null to disable, or positive scalar/{width}-vector') |
|
|
| def motion_scores(actions,reference,thresholds): |
| """Largest per-axis normalized change. Quaternion change is sign-invariant.""" |
| a=np.asarray(actions,float);b=np.asarray(reference,float);scores=[] |
| for key,t in thresholds.items(): |
| if t is None: continue |
| if key in GROUPS: |
| d=np.abs(a[...,GROUPS[key]]-b[GROUPS[key]]) |
| if key=='base_yaw_rad':d=np.abs((a[...,57:58]-b[57:58]+np.pi)%(2*np.pi)-np.pi) |
| scores.append(np.max(d/np.asarray(t),axis=-1)) |
| else: |
| sl=slice(0,4) if key=='left_rotation_deg' else slice(7,11) |
| q=a[...,sl];r=b[sl];den=np.linalg.norm(q,axis=-1)*np.linalg.norm(r) |
| if (den<1e-9).any():raise ValueError('Invalid quaternion in motion selection') |
| angle=np.degrees(2*np.arccos(np.clip(np.abs(np.sum(q*r,axis=-1))/den,0,1))) |
| scores.append(angle/float(np.asarray(t).reshape(-1)[0])) |
| return np.max(np.stack(scores),axis=0) if scores else np.zeros(a.shape[:-1]) |
|
|
| def select_waypoints(actions,enabled=False,thresholds=None,max_gap_frames=30): |
| """Keep first/last, threshold crossings vs last kept, and optional hold anchors. |
| |
| max_gap_frames=0 removes the hold-anchor rule. It does not enforce minimum |
| motion on anchors or endpoints. Source data and source indices are unchanged. |
| """ |
| a=np.asarray(actions);n=len(a) |
| if a.ndim!=2 or a.shape[1]!=58 or not np.isfinite(a).all():raise ValueError('Expected finite [T,58] actions') |
| if int(max_gap_frames)!=max_gap_frames or max_gap_frames<0:raise ValueError('max_gap_frames must be a nonnegative integer') |
| t=dict(DEFAULT_THRESHOLDS if thresholds is None else thresholds);validate_thresholds(t) |
| if not enabled or n<2:return np.arange(n,dtype=np.int64) |
| kept=[0] |
| for i in range(1,n): |
| if motion_scores(a[i],a[kept[-1]],t)>=1 or (max_gap_frames and i-kept[-1]>=max_gap_frames):kept.append(i) |
| if kept[-1]!=n-1:kept.append(n-1) |
| return np.asarray(kept,dtype=np.int64) |
|
|
| def phase_indices(indices,frame_skip): |
| if int(frame_skip)!=frame_skip or frame_skip<1:raise ValueError('frame_skip must be a positive integer') |
| return [np.asarray(indices)[p::int(frame_skip)] for p in range(min(int(frame_skip),len(indices)))] |
|
|
| class YorTemporalDataset(TrajectoryDataset): |
| """Wrap AFTER episode splitting and embedding; phases cannot leak across splits.""" |
| def __init__(self,dataset,frame_skip=1,motion_filter=None): |
| self.dataset=dataset;self.frame_skip=frame_skip;self.motion_filter=dict(motion_filter or {}) |
| self.sequences=[] |
| for i in range(len(dataset)): |
| a=dataset.get_frames(i,range(dataset.get_seq_length(i)))[1] |
| if hasattr(a,'detach'):a=a.detach().cpu().numpy() |
| ids=select_waypoints(a,**self.motion_filter) |
| self.sequences.extend((i,phase,frames) for phase,frames in enumerate(phase_indices(ids,frame_skip))) |
| def __len__(self):return len(self.sequences) |
| def get_seq_length(self,idx):return len(self.sequences[idx][2]) |
| def get_frames(self,idx,frames): |
| i,phase,ids=self.sequences[idx] |
| return self.dataset.get_frames(i,ids[np.asarray(list(frames),dtype=int)].tolist()) |
| def __getitem__(self,idx):return self.get_frames(idx,range(self.get_seq_length(idx))) |
| def get_all_actions(self): |
| import torch |
| return torch.cat([self[i][1] for i in range(len(self))]) |
|
|
|
|
| from datasets.core import TrajectorySlicerDataset |
|
|
| class YorTemporalSlicer(TrajectorySlicerDataset): |
| """Pad short phases and optionally predict actions after each observation. |
| |
| action_offset is measured in sampled sequence rows. With offset 1 the |
| current observation at t predicts t+1 onward. Exclude anchors that have no |
| next target, but repeat the last target when a chunk reaches the end. |
| """ |
| def __init__(self,dataset,action_offset=0,**kwargs): |
| if int(action_offset)!=action_offset or action_offset<0: |
| raise ValueError('action_offset must be a nonnegative integer') |
| self.action_offset=int(action_offset) |
| super().__init__(dataset,**kwargs) |
| self.slices=[(i,t) for i in range(len(dataset)) |
| for t in range(max(0,dataset.get_seq_length(i)-self.action_offset)) |
| if self.pad_seq_length or t+self.action_offset+self.action_window<=dataset.get_seq_length(i)] |
| def source_indices(self,idx): |
| i,t=self.slices[idx];T=self.dataset.get_seq_length(i) |
| obs=np.clip(np.arange(t-self.window+1,t+1),0,T-1) |
| start=t if self.vqbet_get_future_action_chunk else t-self.window+1 |
| act=np.clip(np.arange(start,t+self.action_window)+self.action_offset,self.action_offset,T-1) |
| return i,obs,act |
| def __getitem__(self,idx): |
| i,oi,ai=self.source_indices(idx) |
| obs,act,*others=self.dataset[i] |
| values=(obs[oi],act[ai],*[other[oi] for other in others]) |
| return self.transform(values) if self.transform else values |
|
|
