Fix checkpoint loading across months

Makefile

@@ -96,5 +96,5 @@ static-checks:
 
 test:
 # 	python -m pytest
-	@echo "Unit tests not implemented"
+	@python -m unittest discover -s sim/utils -p '*_test.py'
 .PHONY: test

sim/envs/__init__.py

@@ -18,7 +18,11 @@
 from sim.envs.humanoids.stompymicro_env import StompyMicroEnv
 from sim.envs.humanoids.stompymini_config import MiniCfg, MiniCfgPPO
 from sim.envs.humanoids.stompymini_env import MiniFreeEnv
-from sim.envs.humanoids.stompypro_config import StompyProCfg, StompyProStandingCfg, StompyProCfgPPO
+from sim.envs.humanoids.stompypro_config import (
+    StompyProCfg,
+    StompyProCfgPPO,
+    StompyProStandingCfg,
+)
 from sim.envs.humanoids.stompypro_env import StompyProFreeEnv
 from sim.envs.humanoids.xbot_config import XBotCfg, XBotCfgPPO
 from sim.envs.humanoids.xbot_env import XBotLFreeEnv

sim/envs/base/legged_robot.py

@@ -6,7 +6,7 @@
 from sim import ROOT_DIR
 from sim.envs.base.base_task import BaseTask
 from sim.utils.helpers import class_to_dict
-from sim.utils.math import get_euler_xyz_tensor, quat_apply_yaw, wrap_to_pi
+from sim.utils.math_sim import get_euler_xyz_tensor, quat_apply_yaw, wrap_to_pi
 
 # fmt: off
 from isaacgym import gymapi, gymtorch, gymutil # isort: skip

sim/envs/humanoids/stompymicro_config.py

@@ -12,7 +12,7 @@
 
 class StompyMicroCfg(LeggedRobotCfg):
     """Configuration class for the Legs humanoid robot."""
-    
+
     class env(LeggedRobotCfg.env):
         # change the observation dim
         frame_stack = 15
@@ -48,7 +48,7 @@ class asset(LeggedRobotCfg.asset):
             "DRIVING_ROTOR_PLATE_3",
             "DRIVING_ROTOR_PLATE_4",
             "DRIVING_ROTOR_PLATE_7",
-            "DRIVING_ROTOR_PLATE_8"
+            "DRIVING_ROTOR_PLATE_8",
         ]
 
         penalize_contacts_on = []
@@ -75,10 +75,10 @@ class terrain(LeggedRobotCfg.terrain):
         restitution = 0.0
 
     class noise:
-        add_noise = False # pfb30 bring it back
+        add_noise = False  # pfb30 bring it back
         noise_level = 0.6  # scales other values
 
-        class noise_scales: # pfb30 bring it back
+        class noise_scales:  # pfb30 bring it back
             dof_pos = 0.01
             dof_vel = 0.01
             ang_vel = 0.01
@@ -131,9 +131,9 @@ class domain_rand(LeggedRobotCfg.domain_rand):
         randomize_friction = False
         friction_range = [0.1, 2.0]
 
-        randomize_base_mass = False #True
+        randomize_base_mass = False  # True
         added_mass_range = [-1.0, 1.0]
-        push_robots = False #True
+        push_robots = False  # True
         push_interval_s = 4
         max_push_vel_xy = 0.2
         max_push_ang_vel = 0.4

sim/resources/stompymicro/joints.py

@@ -55,7 +55,7 @@ def __str__(self) -> str:
 class LeftArm(Node):
     shoulder_yaw = "left_shoulder_yaw"
     shoulder_pitch = "left_shoulder_pitch"
-    elbow_pitch = "left_elbow_yaw" # FIXME: yaw vs pitch
+    elbow_pitch = "left_elbow_yaw"  # FIXME: yaw vs pitch
 
 
 class RightArm(Node):
@@ -69,15 +69,15 @@ class LeftLeg(Node):
     hip_yaw = "left_hip_yaw"
     hip_pitch = "left_hip_pitch"
     knee_pitch = "left_knee_pitch"
-    ankle_pitch = "left_ankle_pitch" 
+    ankle_pitch = "left_ankle_pitch"
 
 
 class RightLeg(Node):
     hip_roll = "right_hip_roll"
     hip_yaw = "right_hip_yaw"
     hip_pitch = "right_hip_pitch"
     knee_pitch = "right_knee_pitch"
-    ankle_pitch = "right_ankle_pitch" 
+    ankle_pitch = "right_ankle_pitch"
 
 
 class Legs(Node):
@@ -264,6 +264,7 @@ def friction(cls) -> Dict[str, float]:
             "elbow_pitch": 0.0,
         }
 
+
 def print_joints() -> None:
     joints = Robot.all_joints()
     assert len(joints) == len(set(joints)), "Duplicate joint names found!"

sim/utils/helpers.py

@@ -108,20 +108,35 @@ def parse_sim_params(args, cfg):
 
 
 def get_load_path(root, load_run=-1, checkpoint=-1):
+    def month_to_number(month):
+        return datetime.datetime.strptime(month, "%b").month
+
     try:
         runs = os.listdir(root)
-        # TODO sort by date to handle change of month
-        runs.sort()
+        try:
+            # Handle both new format "YYYY_MMMDD_HH-MM-SS" and old format "MMMDD_HH-MM-SS"
+            def parse_run(run):
+                try:
+                    # New format
+                    return datetime.datetime.strptime(run[:20], "%Y_%b%d_%H-%M-%S")
+                except ValueError:
+                    # Old format
+                    return datetime.datetime.strptime(run[:15], "%b%d_%H-%M-%S")
+
+            runs.sort(key=parse_run)
+        except ValueError as e:
+            print("WARNING - Could not sort runs by date and time: " + str(e))
+            runs.sort()
         if "exported" in runs:
             runs.remove("exported")
         last_run = os.path.join(root, runs[-1])
-    except:
+    except Exception as e:
+        print(type(e).__name__, e)
         raise ValueError("No runs in this directory: " + root)
     if load_run == -1:
         load_run = last_run
     else:
         load_run = os.path.join(root, load_run)
-
     if checkpoint == -1:
         models = [file for file in os.listdir(load_run) if "model" in file]
         models.sort(key=lambda m: "{0:0>15}".format(m))

sim/utils/helpers_test.py

@@ -0,0 +1,143 @@
+import argparse
+import os
+import shutil
+import unittest
+from unittest.mock import MagicMock, mock_open, patch
+
+import numpy as np
+import torch
+
+from sim.utils.helpers import (
+    class_to_dict,
+    export_policy_as_jit,
+    export_policy_as_onnx,
+    get_load_path,
+    parse_sim_params,
+    set_seed,
+    update_cfg_from_args,
+    update_class_from_dict,
+)
+
+
+class TestHelpers(unittest.TestCase):
+
+    def test_class_to_dict(self) -> None:
+        class TestClass:
+            def __init__(self) -> None:
+                self.a = 1
+                self.b = [2, 3]
+                self._c = 4  # should be ignored
+
+        obj = TestClass()
+        result = class_to_dict(obj)
+        self.assertEqual(result, {"a": 1, "b": [2, 3]})
+
+    def test_update_class_from_dict(self) -> None:
+        class TestClass:
+            def __init__(self) -> None:
+                self.a = 1
+                self.b = 2
+
+        obj = TestClass()
+        update_class_from_dict(obj, {"a": 10, "b": 20})
+        self.assertEqual(obj.a, 10)
+        self.assertEqual(obj.b, 20)
+
+    def test_set_seed(self) -> None:
+        set_seed(42)
+        self.assertEqual(np.random.randint(0, 10000), 7270)
+        self.assertEqual(torch.randint(0, 10000, (1,)).item(), 5944)
+
+    def test_parse_sim_params(self) -> None:
+        args = argparse.Namespace(
+            physics_engine="SIM_PHYSX", use_gpu=True, subscenes=2, use_gpu_pipeline=True, num_threads=4
+        )
+        cfg = {}
+        sim_params = parse_sim_params(args, cfg)
+        self.assertTrue(sim_params.physx.use_gpu)
+        self.assertEqual(sim_params.physx.num_subscenes, 2)
+
+    @patch("os.makedirs")
+    @patch("os.path.exists")
+    @patch("builtins.open", new_callable=mock_open)
+    @patch("os.listdir")
+    def test_get_load_path(self, mock_listdir: MagicMock, mock_open: MagicMock, mock_exists: MagicMock, mock_makedirs: MagicMock) -> None:
+        root = "test_runs"
+        mock_exists.return_value = True
+        mock_listdir.return_value = ["2023_Jan01_00-00-00", "2023_Mar15_12-30-45", "2023_Dec31_23-59-59"]
+
+        # Mock the open function to simulate file existence
+        mock_open.side_effect = [
+            mock_open(read_data="test").return_value,
+            mock_open(read_data="test").return_value,
+            mock_open(read_data="test").return_value,
+            mock_open(read_data="test").return_value,
+        ]
+
+        # Test default behavior (latest run, latest model)
+        path = get_load_path(root)
+        self.assertTrue(path.endswith(os.path.join("2023_Dec31_23-59-59", "model_1.pt")))
+
+        # Test with specific run
+        path = get_load_path(root, load_run="2023_Mar15_12-30-45")
+        self.assertTrue(path.endswith(os.path.join("2023_Mar15_12-30-45", "model_1.pt")))
+
+        # Test with specific checkpoint
+        path = get_load_path(root, load_run="2023_Jan01_00-00-00", checkpoint=2)
+        self.assertTrue(path.endswith(os.path.join("2023_Jan01_00-00-00", "model_2.pt")))
+
+        # Test with non-existent run
+        with self.assertRaises(ValueError):
+            get_load_path(root, load_run="non_existent_run")
+
+    def test_update_cfg_from_args(self) -> None:
+        env_cfg = type("obj", (object,), {"env": type("obj", (object,), {"num_envs": 1})()})
+        cfg_train = type("obj", (object,), {"seed": 0, "runner": type("obj", (object,), {"max_iterations": 100})()})
+        args = argparse.Namespace(
+            num_envs=10,
+            seed=42,
+            max_iterations=200,
+            resume=True,
+            experiment_name="test",
+            run_name="run1",
+            load_run="run1",
+            checkpoint=1,
+        )
+        env_cfg, cfg_train = update_cfg_from_args(env_cfg, cfg_train, args)
+        self.assertEqual(env_cfg.env.num_envs, 10)
+        self.assertEqual(cfg_train.seed, 42)
+        self.assertEqual(cfg_train.runner.max_iterations, 200)
+
+    @patch("os.makedirs")
+    @patch("os.path.exists")
+    @patch("torch.jit.save")
+    def test_export_policy_as_jit(self, mock_jit_save: MagicMock, mock_exists: MagicMock, mock_makedirs: MagicMock) -> None:
+        class MockActorCritic:
+            def __init__(self) -> None:
+                self.actor = torch.nn.Linear(10, 2)
+
+        actor_critic = MockActorCritic()
+        path = "test_jit"
+        mock_exists.return_value = True
+        export_policy_as_jit(actor_critic, path)
+        mock_jit_save.assert_called_once()
+        mock_makedirs.assert_called_once_with(path, exist_ok=True)
+
+    @patch("os.makedirs")
+    @patch("os.path.exists")
+    @patch("torch.onnx.export")
+    def test_export_policy_as_onnx(self, mock_onnx_export: MagicMock, mock_exists: MagicMock, mock_makedirs: MagicMock) -> None:
+        class MockActorCritic:
+            def __init__(self) -> None:
+                self.actor = torch.nn.Linear(10, 2)
+
+        actor_critic = MockActorCritic()
+        path = "test_onnx"
+        mock_exists.return_value = True
+        export_policy_as_onnx(actor_critic, path)
+        mock_onnx_export.assert_called_once()
+        mock_makedirs.assert_called_once_with(path, exist_ok=True)
+
+
+if __name__ == "__main__":
+    unittest.main()

sim/utils/task_registry.py

@@ -155,11 +155,15 @@ def make_alg_runner(self, env, name=None, args=None, train_cfg=None, log_root="d
 
         if log_root == "default":
             log_root = os.path.join(ROOT_DIR, "logs", train_cfg.runner.experiment_name)
-            log_dir = os.path.join(log_root, datetime.now().strftime("%b%d_%H-%M-%S") + "_" + train_cfg.runner.run_name)
+            log_dir = os.path.join(
+                log_root, datetime.now().strftime("%Y_%b%d_%H-%M-%S") + "_" + train_cfg.runner.run_name
+            )
         elif log_root is None:
             log_dir = None
         else:
-            log_dir = os.path.join(log_root, datetime.now().strftime("%b%d_%H-%M-%S") + "_" + train_cfg.runner.run_name)
+            log_dir = os.path.join(
+                log_root, datetime.now().strftime("%Y_%b%d_%H-%M-%S") + "_" + train_cfg.runner.run_name
+            )
 
         train_cfg_dict = class_to_dict(train_cfg)
         env_cfg_dict = class_to_dict(self.env_cfg_for_wandb)