Organize test folders (#856)

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>

tests/optim/test_optimizers.py

@@ -0,0 +1,56 @@
+# 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 pytest
+import torch
+
+from lerobot.common.constants import (
+    OPTIMIZER_PARAM_GROUPS,
+    OPTIMIZER_STATE,
+)
+from lerobot.common.optim.optimizers import (
+    AdamConfig,
+    AdamWConfig,
+    SGDConfig,
+    load_optimizer_state,
+    save_optimizer_state,
+)
+
+
+@pytest.mark.parametrize(
+    "config_cls, expected_class",
+    [
+        (AdamConfig, torch.optim.Adam),
+        (AdamWConfig, torch.optim.AdamW),
+        (SGDConfig, torch.optim.SGD),
+    ],
+)
+def test_optimizer_build(config_cls, expected_class, model_params):
+    config = config_cls()
+    optimizer = config.build(model_params)
+    assert isinstance(optimizer, expected_class)
+    assert optimizer.defaults["lr"] == config.lr
+
+
+def test_save_optimizer_state(optimizer, tmp_path):
+    save_optimizer_state(optimizer, tmp_path)
+    assert (tmp_path / OPTIMIZER_STATE).is_file()
+    assert (tmp_path / OPTIMIZER_PARAM_GROUPS).is_file()
+
+
+def test_save_and_load_optimizer_state(model_params, optimizer, tmp_path):
+    save_optimizer_state(optimizer, tmp_path)
+    loaded_optimizer = AdamConfig().build(model_params)
+    loaded_optimizer = load_optimizer_state(loaded_optimizer, tmp_path)
+
+    torch.testing.assert_close(optimizer.state_dict(), loaded_optimizer.state_dict())

tests/optim/test_schedulers.py

@@ -0,0 +1,94 @@
+# 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.
+from torch.optim.lr_scheduler import LambdaLR
+
+from lerobot.common.constants import SCHEDULER_STATE
+from lerobot.common.optim.schedulers import (
+    CosineDecayWithWarmupSchedulerConfig,
+    DiffuserSchedulerConfig,
+    VQBeTSchedulerConfig,
+    load_scheduler_state,
+    save_scheduler_state,
+)
+
+
+def test_diffuser_scheduler(optimizer):
+    config = DiffuserSchedulerConfig(name="cosine", num_warmup_steps=5)
+    scheduler = config.build(optimizer, num_training_steps=100)
+    assert isinstance(scheduler, LambdaLR)
+
+    optimizer.step()  # so that we don't get torch warning
+    scheduler.step()
+    expected_state_dict = {
+        "_get_lr_called_within_step": False,
+        "_last_lr": [0.0002],
+        "_step_count": 2,
+        "base_lrs": [0.001],
+        "last_epoch": 1,
+        "lr_lambdas": [None],
+        "verbose": False,
+    }
+    assert scheduler.state_dict() == expected_state_dict
+
+
+def test_vqbet_scheduler(optimizer):
+    config = VQBeTSchedulerConfig(num_warmup_steps=10, num_vqvae_training_steps=20, num_cycles=0.5)
+    scheduler = config.build(optimizer, num_training_steps=100)
+    assert isinstance(scheduler, LambdaLR)
+
+    optimizer.step()
+    scheduler.step()
+    expected_state_dict = {
+        "_get_lr_called_within_step": False,
+        "_last_lr": [0.001],
+        "_step_count": 2,
+        "base_lrs": [0.001],
+        "last_epoch": 1,
+        "lr_lambdas": [None],
+        "verbose": False,
+    }
+    assert scheduler.state_dict() == expected_state_dict
+
+
+def test_cosine_decay_with_warmup_scheduler(optimizer):
+    config = CosineDecayWithWarmupSchedulerConfig(
+        num_warmup_steps=10, num_decay_steps=90, peak_lr=0.01, decay_lr=0.001
+    )
+    scheduler = config.build(optimizer, num_training_steps=100)
+    assert isinstance(scheduler, LambdaLR)
+
+    optimizer.step()
+    scheduler.step()
+    expected_state_dict = {
+        "_get_lr_called_within_step": False,
+        "_last_lr": [0.0001818181818181819],
+        "_step_count": 2,
+        "base_lrs": [0.001],
+        "last_epoch": 1,
+        "lr_lambdas": [None],
+        "verbose": False,
+    }
+    assert scheduler.state_dict() == expected_state_dict
+
+
+def test_save_scheduler_state(scheduler, tmp_path):
+    save_scheduler_state(scheduler, tmp_path)
+    assert (tmp_path / SCHEDULER_STATE).is_file()
+
+
+def test_save_load_scheduler_state(scheduler, tmp_path):
+    save_scheduler_state(scheduler, tmp_path)
+    loaded_scheduler = load_scheduler_state(scheduler, tmp_path)
+
+    assert scheduler.state_dict() == loaded_scheduler.state_dict()