HI
openpi RoboCasa365 fork + pi0.5 port, single-task PickPlaceCounterToCabinet configs
44aecac verified | import os | |
| import pytest | |
| import numpy as np | |
| import torch | |
| from openpi.training.groot_lerobot_dataset import GrootMultiDataset | |
| def test_groot_multi_sampling_weights(): | |
| dir_a = os.environ["TEST_GROOT_DIR_A"] | |
| dir_b = os.environ["TEST_GROOT_DIR_B"] | |
| weights = [0.9, 0.1] | |
| # Wrap the original class to annotate samples with dataset_index for counting | |
| class AnnotatedMulti(GrootMultiDataset): | |
| def __getitem__(self, index): | |
| sample = super().__getitem__(index) | |
| # Add testing-only field | |
| sample["dataset_index"] = int(self.sampling_indices[index][0]) | |
| return sample | |
| multi = AnnotatedMulti( | |
| data_dirs=[dir_a, dir_b], | |
| weights=weights, | |
| action_horizon=4, | |
| shuffle=True, | |
| action_dim=12, | |
| ) | |
| # Iterate over a DataLoader and accumulate counts to reach ~1000 samples | |
| loader = torch.utils.data.DataLoader(multi, batch_size=16, shuffle=True) | |
| target_samples = 1000 | |
| counts = np.zeros(2, dtype=int) | |
| seen = 0 | |
| for batch in loader: | |
| # dataset_index will be a tensor of shape (B,) | |
| di = batch["dataset_index"].numpy() | |
| # Accumulate counts from this batch | |
| binc = np.bincount(di, minlength=2) | |
| counts[: len(binc)] += binc | |
| seen += di.shape[0] | |
| if seen >= target_samples: | |
| break | |
| empirical = counts / counts.sum() | |
| target = np.array(weights) / np.sum(weights) | |
| assert np.allclose(empirical, target, atol=0.05), f"Empirical {empirical.tolist()} vs target {target.tolist()}" |