JC6123/aloha
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

代码可以跑起来了

Browse Source
This commit is contained in:
JC6123
2026-02-19 15:32:28 +08:00
parent b701d939c2
commit 88d14221ae
11 changed files with 503 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

230
utils.py
View File

@@ -2,21 +2,88 @@ import numpy as np
import torch
import os
import h5py
import re
from torch.utils.data import TensorDataset, DataLoader
import IPython
e = IPython.embed
class EpisodicDataset(torch.utils.data.Dataset):
def __init__(self, episode_ids, dataset_dir, camera_names, norm_stats):
def __init__(self, episode_ids, dataset_dir, camera_names, norm_stats,
use_text_instruction=False,
instruction_mode='timestep-level',
use_cached_text_features=True,
text_feature_dim=768,
text_tokenizer_name='distilbert-base-uncased',
text_max_length=32):
super(EpisodicDataset).__init__()
self.episode_ids = episode_ids
self.dataset_dir = dataset_dir
self.camera_names = camera_names
self.norm_stats = norm_stats
self.use_text_instruction = use_text_instruction
self.instruction_mode = instruction_mode
self.use_cached_text_features = use_cached_text_features
self.text_feature_dim = text_feature_dim
self.text_max_length = text_max_length
self.is_sim = None
self.max_episode_len = None
self.action_dim = None
self.text_tokenizer = None
if self.use_text_instruction:
try:
from transformers import DistilBertTokenizerFast
except ImportError as exc:
raise ImportError(
'transformers is required for text instruction loading. '
'Install it with: pip install transformers'
) from exc
self.text_tokenizer = DistilBertTokenizerFast.from_pretrained(text_tokenizer_name)
self._init_episode_shapes()
self.__getitem__(0) # initialize self.is_sim
def _init_episode_shapes(self):
max_len = 0
action_dim = None
for episode_id in self.episode_ids:
dataset_path = os.path.join(self.dataset_dir, f'episode_{episode_id}.hdf5')
with h5py.File(dataset_path, 'r') as root:
shape = root['/action'].shape
if len(shape) != 2:
raise ValueError(f'Expected /action to have shape [T, D], got {shape} in {dataset_path}')
max_len = max(max_len, int(shape[0]))
if action_dim is None:
action_dim = int(shape[1])
elif int(shape[1]) != action_dim:
raise ValueError(
f'Inconsistent action dim in dataset. Expected {action_dim}, got {shape[1]} in {dataset_path}'
)
if max_len <= 0 or action_dim is None:
raise ValueError(f'Invalid dataset metadata in {self.dataset_dir}')
self.max_episode_len = max_len
self.action_dim = action_dim
@staticmethod
def _decode_instruction(raw_value):
if raw_value is None:
return ''
if isinstance(raw_value, bytes):
return raw_value.decode('utf-8')
if isinstance(raw_value, np.bytes_):
return raw_value.tobytes().decode('utf-8')
if isinstance(raw_value, np.ndarray):
if raw_value.shape == ():
return EpisodicDataset._decode_instruction(raw_value.item())
if raw_value.size == 0:
return ''
return EpisodicDataset._decode_instruction(raw_value.reshape(-1)[0])
return str(raw_value)
def __len__(self):
return len(self.episode_ids)
@@ -26,7 +93,7 @@ class EpisodicDataset(torch.utils.data.Dataset):
episode_id = self.episode_ids[index]
dataset_path = os.path.join(self.dataset_dir, f'episode_{episode_id}.hdf5')
with h5py.File(dataset_path, 'r') as root:
is_sim = root.attrs['sim']
is_sim = bool(root.attrs.get('sim', False))
original_action_shape = root['/action'].shape
episode_len = original_action_shape[0]
if sample_full_episode:
@@ -35,10 +102,40 @@ class EpisodicDataset(torch.utils.data.Dataset):
start_ts = np.random.choice(episode_len)
# get observation at start_ts only
qpos = root['/observations/qpos'][start_ts]
qvel = root['/observations/qvel'][start_ts]
image_dict = dict()
for cam_name in self.camera_names:
image_dict[cam_name] = root[f'/observations/images/{cam_name}'][start_ts]
instruction = ''
text_feature = None
if self.use_text_instruction:
effective_mode = self.instruction_mode
if effective_mode == 'timestep-level' and '/instruction_timestep' in root:
instruction = self._decode_instruction(root['/instruction_timestep'][start_ts])
elif '/instruction' in root:
instruction_node = root['/instruction']
if getattr(instruction_node, 'shape', ()) == ():
instruction = self._decode_instruction(instruction_node[()])
else:
if len(instruction_node.shape) == 1 and instruction_node.shape[0] == episode_len:
instruction = self._decode_instruction(instruction_node[start_ts])
else:
instruction = self._decode_instruction(instruction_node[0])
if self.use_cached_text_features:
if effective_mode == 'timestep-level' and '/instruction_features_timestep' in root:
text_feature = root['/instruction_features_timestep'][start_ts]
elif '/instruction_features' in root:
feat_node = root['/instruction_features']
if getattr(feat_node, 'shape', ()) == ():
text_feature = np.array(feat_node[()])
elif len(feat_node.shape) == 1:
text_feature = feat_node[()]
elif len(feat_node.shape) == 2 and feat_node.shape[0] == episode_len:
text_feature = feat_node[start_ts]
else:
text_feature = feat_node[0]
# get all actions after and including start_ts
if is_sim:
action = root['/action'][start_ts:]
@@ -48,10 +145,10 @@ class EpisodicDataset(torch.utils.data.Dataset):
action_len = episode_len - max(0, start_ts - 1) # hack, to make timesteps more aligned
self.is_sim = is_sim
padded_action = np.zeros(original_action_shape, dtype=np.float32)
padded_action = np.zeros((self.max_episode_len, self.action_dim), dtype=np.float32)
padded_action[:action_len] = action
is_pad = np.zeros(episode_len)
is_pad[action_len:] = 1
is_pad = np.ones(self.max_episode_len)
is_pad[:action_len] = 0
# new axis for different cameras
all_cam_images = []
@@ -73,55 +170,132 @@ class EpisodicDataset(torch.utils.data.Dataset):
action_data = (action_data - self.norm_stats["action_mean"]) / self.norm_stats["action_std"]
qpos_data = (qpos_data - self.norm_stats["qpos_mean"]) / self.norm_stats["qpos_std"]
return image_data, qpos_data, action_data, is_pad
if self.use_text_instruction and text_feature is not None:
text_feature_data = torch.from_numpy(np.array(text_feature)).float()
text_feature_valid = torch.tensor(True, dtype=torch.bool)
text_input_ids = torch.zeros(1, dtype=torch.long)
text_attention_mask = torch.zeros(1, dtype=torch.long)
elif self.use_text_instruction:
tokenized = self.text_tokenizer(
instruction,
padding='max_length',
truncation=True,
max_length=self.text_max_length,
return_tensors='pt',
)
text_input_ids = tokenized['input_ids'].squeeze(0).long()
text_attention_mask = tokenized['attention_mask'].squeeze(0).long()
text_feature_data = torch.zeros(self.text_feature_dim, dtype=torch.float32)
text_feature_valid = torch.tensor(False, dtype=torch.bool)
else:
text_input_ids = torch.zeros(1, dtype=torch.long)
text_attention_mask = torch.zeros(1, dtype=torch.long)
text_feature_data = torch.zeros(self.text_feature_dim, dtype=torch.float32)
text_feature_valid = torch.tensor(False, dtype=torch.bool)
return image_data, qpos_data, action_data, is_pad, text_input_ids, text_attention_mask, text_feature_data, text_feature_valid
def get_norm_stats(dataset_dir, num_episodes):
def _discover_episode_ids(dataset_dir, num_episodes=None):
pattern = re.compile(r'^episode_(\d+)\.hdf5$')
episode_ids = []
for fname in os.listdir(dataset_dir):
m = pattern.match(fname)
if m:
episode_ids.append(int(m.group(1)))
episode_ids.sort()
if num_episodes is not None:
episode_ids = episode_ids[:num_episodes]
return episode_ids
def get_norm_stats(dataset_dir, episode_ids):
all_qpos_data = []
all_action_data = []
for episode_idx in range(num_episodes):
example_qpos = None
for episode_idx in episode_ids:
dataset_path = os.path.join(dataset_dir, f'episode_{episode_idx}.hdf5')
with h5py.File(dataset_path, 'r') as root:
qpos = root['/observations/qpos'][()]
qvel = root['/observations/qvel'][()]
action = root['/action'][()]
all_qpos_data.append(torch.from_numpy(qpos))
all_action_data.append(torch.from_numpy(action))
all_qpos_data = torch.stack(all_qpos_data)
all_action_data = torch.stack(all_action_data)
all_action_data = all_action_data
qpos_t = torch.from_numpy(qpos)
action_t = torch.from_numpy(action)
all_qpos_data.append(qpos_t)
all_action_data.append(action_t)
if example_qpos is None and len(qpos) > 0:
example_qpos = qpos[0]
# Episodes may have different lengths; concatenate over time axis.
all_qpos_data = torch.cat(all_qpos_data, dim=0)
all_action_data = torch.cat(all_action_data, dim=0)
# normalize action data
action_mean = all_action_data.mean(dim=[0, 1], keepdim=True)
action_std = all_action_data.std(dim=[0, 1], keepdim=True)
action_mean = all_action_data.mean(dim=0, keepdim=True)
action_std = all_action_data.std(dim=0, keepdim=True)
action_std = torch.clip(action_std, 1e-2, np.inf) # clipping
# normalize qpos data
qpos_mean = all_qpos_data.mean(dim=[0, 1], keepdim=True)
qpos_std = all_qpos_data.std(dim=[0, 1], keepdim=True)
qpos_mean = all_qpos_data.mean(dim=0, keepdim=True)
qpos_std = all_qpos_data.std(dim=0, keepdim=True)
qpos_std = torch.clip(qpos_std, 1e-2, np.inf) # clipping
stats = {"action_mean": action_mean.numpy().squeeze(), "action_std": action_std.numpy().squeeze(),
"qpos_mean": qpos_mean.numpy().squeeze(), "qpos_std": qpos_std.numpy().squeeze(),
"example_qpos": qpos}
"example_qpos": example_qpos}
return stats
def load_data(dataset_dir, num_episodes, camera_names, batch_size_train, batch_size_val):
def load_data(dataset_dir, num_episodes, camera_names, batch_size_train, batch_size_val,
use_text_instruction=False,
instruction_mode='timestep-level',
use_cached_text_features=True,
text_feature_dim=768,
text_tokenizer_name='distilbert-base-uncased',
text_max_length=32):
print(f'\nData from: {dataset_dir}\n')
episode_ids = _discover_episode_ids(dataset_dir, num_episodes)
if len(episode_ids) == 0:
raise FileNotFoundError(f'No episode_*.hdf5 found in {dataset_dir}')
if len(episode_ids) < 2:
raise ValueError(f'Need at least 2 episodes for train/val split, found {len(episode_ids)} in {dataset_dir}')
# obtain train test split
train_ratio = 0.8
shuffled_indices = np.random.permutation(num_episodes)
train_indices = shuffled_indices[:int(train_ratio * num_episodes)]
val_indices = shuffled_indices[int(train_ratio * num_episodes):]
shuffled_indices = np.random.permutation(len(episode_ids))
train_count = int(train_ratio * len(episode_ids))
train_indices = shuffled_indices[:train_count]
val_indices = shuffled_indices[train_count:]
train_episode_ids = np.array(episode_ids)[train_indices]
val_episode_ids = np.array(episode_ids)[val_indices]
# obtain normalization stats for qpos and action
norm_stats = get_norm_stats(dataset_dir, num_episodes)
norm_stats = get_norm_stats(dataset_dir, episode_ids)
# construct dataset and dataloader
train_dataset = EpisodicDataset(train_indices, dataset_dir, camera_names, norm_stats)
val_dataset = EpisodicDataset(val_indices, dataset_dir, camera_names, norm_stats)
train_dataset = EpisodicDataset(
train_episode_ids,
dataset_dir,
camera_names,
norm_stats,
use_text_instruction=use_text_instruction,
instruction_mode=instruction_mode,
use_cached_text_features=use_cached_text_features,
text_feature_dim=text_feature_dim,
text_tokenizer_name=text_tokenizer_name,
text_max_length=text_max_length,
)
val_dataset = EpisodicDataset(
val_episode_ids,
dataset_dir,
camera_names,
norm_stats,
use_text_instruction=use_text_instruction,
instruction_mode=instruction_mode,
use_cached_text_features=use_cached_text_features,
text_feature_dim=text_feature_dim,
text_tokenizer_name=text_tokenizer_name,
text_max_length=text_max_length,
)
train_dataloader = DataLoader(train_dataset, batch_size=batch_size_train, shuffle=True, pin_memory=True, num_workers=1, prefetch_factor=1)
val_dataloader = DataLoader(val_dataset, batch_size=batch_size_val, shuffle=True, pin_memory=True, num_workers=1, prefetch_factor=1)