Enable tests for TD-MPC (#160)

2024-05-09 13:42:12 +01:00
parent 7bb5b15f4c
commit e89521dfa0
7 changed files with 5 additions and 6 deletions
--- a/7
+++ b/7
@@ -22,9 +22,8 @@ test-end-to-end:
 	${MAKE} test-act-ete-eval
 	${MAKE} test-diffusion-ete-train
 	${MAKE} test-diffusion-ete-eval
-	# TODO(rcadene, alexander-soare): enable end-to-end tests for tdmpc
-	# ${MAKE} test-tdmpc-ete-train
-	# ${MAKE} test-tdmpc-ete-eval
+	${MAKE} test-tdmpc-ete-train
+	${MAKE} test-tdmpc-ete-eval
 	${MAKE} test-default-ete-eval

 test-act-ete-train:
@@ -80,7 +79,7 @@ test-tdmpc-ete-train:
 		policy=tdmpc \
 		env=xarm \
 		env.task=XarmLift-v0 \
-		dataset_repo_id=lerobot/xarm_lift_medium_replay \
+		dataset_repo_id=lerobot/xarm_lift_medium \
 		wandb.enable=False \
 		training.offline_steps=2 \
 		training.online_steps=2 \
--- a/lerobot/configs/policy/tdmpc.yaml
+++ b/lerobot/configs/policy/tdmpc.yaml
@@ -1,7 +1,7 @@
 # @package _global_

 seed: 1
-dataset_repo_id: lerobot/xarm_lift_medium_replay
+dataset_repo_id: lerobot/xarm_lift_medium

 training:
  offline_steps: 25000
--- a/tests/data/save_policy_to_safetensors/xarm_tdmpc/actions.safetensors
+++ b/tests/data/save_policy_to_safetensors/xarm_tdmpc/actions.safetensors
--- a/tests/data/save_policy_to_safetensors/xarm_tdmpc/grad_stats.safetensors
+++ b/tests/data/save_policy_to_safetensors/xarm_tdmpc/grad_stats.safetensors
--- a/tests/data/save_policy_to_safetensors/xarm_tdmpc/output_dict.safetensors
+++ b/tests/data/save_policy_to_safetensors/xarm_tdmpc/output_dict.safetensors
--- a/tests/data/save_policy_to_safetensors/xarm_tdmpc/param_stats.safetensors
+++ b/tests/data/save_policy_to_safetensors/xarm_tdmpc/param_stats.safetensors
--- a/tests/test_policies.py
+++ b/tests/test_policies.py
@@ -236,7 +236,7 @@ def test_normalize(insert_temporal_dim):
@pytest.mark.parametrize(
    "env_name, policy_name, extra_overrides",
    [
-        # ("xarm", "tdmpc", ["policy.n_action_repeats=2"]),
+        ("xarm", "tdmpc", []),
        (
            "pusht",
            "diffusion",