123 lines
3.7 KiB
Python
123 lines
3.7 KiB
Python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import random
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from lerobot.common.utils.random_utils import (
|
|
deserialize_numpy_rng_state,
|
|
deserialize_python_rng_state,
|
|
deserialize_rng_state,
|
|
deserialize_torch_rng_state,
|
|
get_rng_state,
|
|
seeded_context,
|
|
serialize_numpy_rng_state,
|
|
serialize_python_rng_state,
|
|
serialize_rng_state,
|
|
serialize_torch_rng_state,
|
|
set_rng_state,
|
|
set_seed,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def fixed_seed():
|
|
"""Fixture to set a consistent initial seed for each test."""
|
|
set_seed(12345)
|
|
yield
|
|
|
|
|
|
def test_serialize_deserialize_python_rng(fixed_seed):
|
|
# Save state after generating val1
|
|
_ = random.random()
|
|
st = serialize_python_rng_state()
|
|
# Next random is val2
|
|
val2 = random.random()
|
|
# Restore the state, so the next random should match val2
|
|
deserialize_python_rng_state(st)
|
|
val3 = random.random()
|
|
assert val2 == val3
|
|
|
|
|
|
def test_serialize_deserialize_numpy_rng(fixed_seed):
|
|
_ = np.random.rand()
|
|
st = serialize_numpy_rng_state()
|
|
val2 = np.random.rand()
|
|
deserialize_numpy_rng_state(st)
|
|
val3 = np.random.rand()
|
|
assert val2 == val3
|
|
|
|
|
|
def test_serialize_deserialize_torch_rng(fixed_seed):
|
|
_ = torch.rand(1).item()
|
|
st = serialize_torch_rng_state()
|
|
val2 = torch.rand(1).item()
|
|
deserialize_torch_rng_state(st)
|
|
val3 = torch.rand(1).item()
|
|
assert val2 == val3
|
|
|
|
|
|
def test_serialize_deserialize_rng(fixed_seed):
|
|
# Generate one from each library
|
|
_ = random.random()
|
|
_ = np.random.rand()
|
|
_ = torch.rand(1).item()
|
|
# Serialize
|
|
st = serialize_rng_state()
|
|
# Generate second set
|
|
val_py2 = random.random()
|
|
val_np2 = np.random.rand()
|
|
val_th2 = torch.rand(1).item()
|
|
# Restore, so the next draws should match val_py2, val_np2, val_th2
|
|
deserialize_rng_state(st)
|
|
assert random.random() == val_py2
|
|
assert np.random.rand() == val_np2
|
|
assert torch.rand(1).item() == val_th2
|
|
|
|
|
|
def test_get_set_rng_state(fixed_seed):
|
|
st = get_rng_state()
|
|
val1 = (random.random(), np.random.rand(), torch.rand(1).item())
|
|
# Change states
|
|
random.random()
|
|
np.random.rand()
|
|
torch.rand(1)
|
|
# Restore
|
|
set_rng_state(st)
|
|
val2 = (random.random(), np.random.rand(), torch.rand(1).item())
|
|
assert val1 == val2
|
|
|
|
|
|
def test_set_seed():
|
|
set_seed(1337)
|
|
val1 = (random.random(), np.random.rand(), torch.rand(1).item())
|
|
set_seed(1337)
|
|
val2 = (random.random(), np.random.rand(), torch.rand(1).item())
|
|
assert val1 == val2
|
|
|
|
|
|
def test_seeded_context(fixed_seed):
|
|
val1 = (random.random(), np.random.rand(), torch.rand(1).item())
|
|
with seeded_context(1337):
|
|
seeded_val1 = (random.random(), np.random.rand(), torch.rand(1).item())
|
|
val2 = (random.random(), np.random.rand(), torch.rand(1).item())
|
|
with seeded_context(1337):
|
|
seeded_val2 = (random.random(), np.random.rand(), torch.rand(1).item())
|
|
|
|
assert seeded_val1 == seeded_val2
|
|
assert all(a != b for a, b in zip(val1, seeded_val1, strict=True)) # changed inside the context
|
|
assert all(a != b for a, b in zip(val2, seeded_val2, strict=True)) # changed again after exiting
|