95 lines
2.8 KiB
Python
95 lines
2.8 KiB
Python
import random
|
|
from typing import Callable
|
|
from uuid import uuid4
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from datasets import Dataset
|
|
|
|
from lerobot.common.datasets.utils import (
|
|
calculate_episode_data_index,
|
|
hf_transform_to_torch,
|
|
reset_episode_index,
|
|
)
|
|
from lerobot.common.utils.utils import (
|
|
get_global_random_state,
|
|
init_hydra_config,
|
|
seeded_context,
|
|
set_global_random_state,
|
|
set_global_seed,
|
|
)
|
|
|
|
# Random generation functions for testing the seeding and random state get/set.
|
|
rand_fns = [
|
|
random.random,
|
|
np.random.random,
|
|
lambda: torch.rand(1).item(),
|
|
]
|
|
if torch.cuda.is_available():
|
|
rand_fns.append(lambda: torch.rand(1, device="cuda"))
|
|
|
|
|
|
@pytest.mark.parametrize("rand_fn", rand_fns)
|
|
def test_seeding(rand_fn: Callable[[], int]):
|
|
set_global_seed(0)
|
|
a = rand_fn()
|
|
with seeded_context(1337):
|
|
c = rand_fn()
|
|
b = rand_fn()
|
|
set_global_seed(0)
|
|
a_ = rand_fn()
|
|
b_ = rand_fn()
|
|
# Check that `set_global_seed` lets us reproduce a and b.
|
|
assert a_ == a
|
|
# Additionally, check that the `seeded_context` didn't interrupt the global RNG.
|
|
assert b_ == b
|
|
set_global_seed(1337)
|
|
c_ = rand_fn()
|
|
# Check that `seeded_context` and `global_seed` give the same reproducibility.
|
|
assert c_ == c
|
|
|
|
|
|
def test_get_set_random_state():
|
|
"""Check that getting the random state, then setting it results in the same random number generation."""
|
|
random_state_dict = get_global_random_state()
|
|
rand_numbers = [rand_fn() for rand_fn in rand_fns]
|
|
set_global_random_state(random_state_dict)
|
|
rand_numbers_ = [rand_fn() for rand_fn in rand_fns]
|
|
assert rand_numbers_ == rand_numbers
|
|
|
|
|
|
def test_calculate_episode_data_index():
|
|
dataset = Dataset.from_dict(
|
|
{
|
|
"timestamp": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6],
|
|
"index": [0, 1, 2, 3, 4, 5],
|
|
"episode_index": [0, 0, 1, 2, 2, 2],
|
|
},
|
|
)
|
|
dataset.set_transform(hf_transform_to_torch)
|
|
episode_data_index = calculate_episode_data_index(dataset)
|
|
assert torch.equal(episode_data_index["from"], torch.tensor([0, 2, 3]))
|
|
assert torch.equal(episode_data_index["to"], torch.tensor([2, 3, 6]))
|
|
|
|
|
|
def test_reset_episode_index():
|
|
dataset = Dataset.from_dict(
|
|
{
|
|
"timestamp": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6],
|
|
"index": [0, 1, 2, 3, 4, 5],
|
|
"episode_index": [10, 10, 11, 12, 12, 12],
|
|
},
|
|
)
|
|
dataset.set_transform(hf_transform_to_torch)
|
|
correct_episode_index = [0, 0, 1, 2, 2, 2]
|
|
dataset = reset_episode_index(dataset)
|
|
assert dataset["episode_index"] == correct_episode_index
|
|
|
|
|
|
def test_init_hydra_config_empty():
|
|
test_file = f"/tmp/test_init_hydra_config_empty_{uuid4().hex}.yaml"
|
|
with open(test_file, "w") as f:
|
|
f.write("\n")
|
|
init_hydra_config(test_file)
|