| """Cached LeRobot v3 batches with explicit episode-safe windows.""" |
| from concurrent.futures import ThreadPoolExecutor |
| import json |
| from pathlib import Path |
| import numpy as np |
| from prepare import DELTA, normalize |
|
|
| IMAGE_NAMES=('base_0_rgb','left_wrist_0_rgb','right_wrist_0_rgb') |
|
|
| class Dataset: |
| def __init__(self,root,workers=32): |
| self.root=Path(root);self.cache=self.root/'cache/prepared' |
| self.norm=json.loads((self.root/'config/normalization.json').read_text()) |
| assert json.loads((self.root/'config/preparation_receipt.json').read_text())['status']=='passed' |
| for name in ['states_raw','states','actions_raw','episode_ids','frame_ids','tokens','token_masks']: |
| setattr(self,name,np.load(self.cache/f'{name}.npy',mmap_mode='r')) |
| self.rgb={i:np.load(self.cache/f'rgb_{i:03d}.npy',mmap_mode='r') for i in range(259)} |
| self.pool=ThreadPoolExecutor(max_workers=workers) |
| self.prefetch=ThreadPoolExecutor(max_workers=1) |
|
|
| def starts(self,split):return np.load(self.cache/f'{split}_starts.npy') |
|
|
| def _image(self,i):return self.rgb[int(self.episode_ids[i])][int(self.frame_ids[i])] |
|
|
| def batch(self,starts,batch_size): |
| n=len(starts);idx=np.pad(starts,(0,batch_size-n),mode='edge') |
| images=np.stack(list(self.pool.map(self._image,idx))) |
| actions=np.array(self.actions_raw[idx[:,None]+np.arange(30)],copy=True) |
| assert np.all(self.episode_ids[idx[:,None]+np.arange(30)]==self.episode_ids[idx,None]) |
| actions[:,:,DELTA]-=self.states_raw[idx][:,None,DELTA] |
| actions=np.pad(normalize(actions,self.norm['actions']),((0,0),(0,0),(0,9))) |
| batch={'images':{name:images[:,k] for k,name in enumerate(IMAGE_NAMES)},'state':np.asarray(self.states[idx]), |
| 'tokens':np.asarray(self.tokens[idx]),'token_mask':np.asarray(self.token_masks[idx]),'actions':actions, |
| 'ids':idx.astype(np.int32)} |
| return batch,{'n':n,'episode_ids':np.asarray(self.episode_ids[starts]),'starts':starts} |
|
|
| def eval_batches(self,split,batch_size): |
| starts=self.starts(split) |
| chunks=[starts[i:i+batch_size] for i in range(0,len(starts),batch_size)] |
| future=self.prefetch.submit(self.batch,chunks[0],batch_size) |
| for i in range(len(chunks)): |
| result=future.result() |
| if i+1<len(chunks):future=self.prefetch.submit(self.batch,chunks[i+1],batch_size) |
| yield result |
|
|
| def train_batches(self,batch_size,seed,start_step=0): |
| starts=self.starts('train');per_epoch=len(starts)//batch_size;step=start_step |
| while True: |
| epoch=step//per_epoch;offset=step%per_epoch |
| perm=np.random.default_rng(np.random.SeedSequence([seed,epoch])).permutation(starts) |
| for b in range(offset,per_epoch): |
| yield self.batch(perm[b*batch_size:(b+1)*batch_size],batch_size)[0] |
| step+=1 |
|
|
| def inverse_actions(self,pred,starts): |
| n=self.norm['actions'];pred=np.asarray(pred)[...,:23] |
| out=(pred+1)/2*(np.array(n['q99'])-np.array(n['q01'])+1e-6)+np.array(n['q01']) |
| out[:,:,DELTA]+=self.states_raw[starts][:,None,DELTA] |
| return out |
|
|
| def close(self):self.prefetch.shutdown();self.pool.shutdown() |
|
|
