Revert "Revert "Override action queue defaults""

This reverts commit 5647d71c08.
This commit is contained in:
Simon Alibert 2024-05-03 11:05:59 +02:00
parent 5647d71c08
commit 75dd478b54
3 changed files with 14 additions and 6 deletions

View File

@ -89,9 +89,13 @@ def save_policy_to_safetensors(output_dir, env_name, policy_name, extra_override
if __name__ == "__main__": if __name__ == "__main__":
env_policies = [ env_policies = [
("xarm", "tdmpc", []), ("xarm", "tdmpc", ["policy.n_action_repeats=2"]),
("pusht", "diffusion", ["policy.num_inference_steps=10", "policy.down_dims=[128, 256, 512]"]), (
("aloha", "act", []), "pusht",
"diffusion",
["policy.n_action_steps=8", "policy.num_inference_steps=10", "policy.down_dims=[128, 256, 512]"],
),
("aloha", "act", ["policy.n_action_steps=10"]),
] ]
for env, policy, extra_overrides in env_policies: for env, policy, extra_overrides in env_policies:
save_policy_to_safetensors("tests/data/save_policy_to_safetensors", env, policy, extra_overrides) save_policy_to_safetensors("tests/data/save_policy_to_safetensors", env, policy, extra_overrides)

View File

@ -239,9 +239,13 @@ def test_normalize(insert_temporal_dim):
@pytest.mark.parametrize( @pytest.mark.parametrize(
"env_name, policy_name, extra_overrides", "env_name, policy_name, extra_overrides",
[ [
("xarm", "tdmpc", []), ("xarm", "tdmpc", ["policy.n_action_repeats=2"]),
("pusht", "diffusion", ["policy.num_inference_steps=10", "policy.down_dims=[128, 256, 512]"]), (
("aloha", "act", []), "pusht",
"diffusion",
["policy.n_action_steps=8", "policy.num_inference_steps=10", "policy.down_dims=[128, 256, 512]"],
),
("aloha", "act", ["policy.n_action_steps=10"]),
], ],
) )
# As artifacts have been generated on an x86_64 kernel, this test won't # As artifacts have been generated on an x86_64 kernel, this test won't