feat(vla): add SmolVLA conditioning and experiment artifacts
This commit is contained in:
@@ -12,10 +12,13 @@ from roboimi.vla.data.simpe_robot_dataset import SimpleRobotDataset
|
||||
|
||||
|
||||
class SimpleRobotDatasetImageLoadingTest(unittest.TestCase):
|
||||
def _write_episode(self, dataset_dir: Path) -> None:
|
||||
episode_path = dataset_dir / "episode_0.hdf5"
|
||||
def _write_episode(self, dataset_dir: Path, episode_idx: int = 0, action_offset: float = 0.0) -> None:
|
||||
episode_path = dataset_dir / f"episode_{episode_idx}.hdf5"
|
||||
with h5py.File(episode_path, "w") as root:
|
||||
root.create_dataset("action", data=np.arange(8, dtype=np.float32).reshape(4, 2))
|
||||
root.create_dataset(
|
||||
"action",
|
||||
data=(np.arange(8, dtype=np.float32).reshape(4, 2) + action_offset),
|
||||
)
|
||||
root.create_dataset(
|
||||
"observations/qpos",
|
||||
data=np.arange(16, dtype=np.float32).reshape(4, 4),
|
||||
@@ -79,3 +82,22 @@ class SimpleRobotDatasetImageLoadingTest(unittest.TestCase):
|
||||
|
||||
fake_cv2.resize.assert_not_called()
|
||||
self.assertEqual(tuple(sample["observation.front"].shape), (2, 3, 8, 8))
|
||||
|
||||
def test_dataset_can_filter_by_episode_indices_and_expose_available_episode_indices(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
dataset_dir = Path(tmpdir)
|
||||
self._write_episode(dataset_dir, episode_idx=3, action_offset=0.0)
|
||||
self._write_episode(dataset_dir, episode_idx=7, action_offset=100.0)
|
||||
|
||||
dataset = SimpleRobotDataset(
|
||||
dataset_dir,
|
||||
obs_horizon=2,
|
||||
pred_horizon=3,
|
||||
camera_names=["front"],
|
||||
episode_indices=[7],
|
||||
)
|
||||
sample = dataset[0]
|
||||
|
||||
self.assertEqual(dataset.available_episode_indices, [7])
|
||||
self.assertEqual(len(dataset), 4)
|
||||
self.assertEqual(sample["action"][0, 0].item(), 100.0)
|
||||
|
||||
Reference in New Issue
Block a user