hangX/lerobot
1
0
Fork 0
forked from tangger/lerobot
Code Pull Requests Activity

Remove EMA model from Diffusion Policy (#134)

Browse Source
This commit is contained in:
Alexander Soare
2024-05-05 11:26:12 +01:00
committed by GitHub
parent d747195c57
commit f3bba0270d
11 changed files with 21 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tests/scripts/save_policy_to_safetensor.py
View File

@@ -88,14 +88,8 @@ def save_policy_to_safetensors(output_dir, env_name, policy_name, extra_override
if __name__ == "__main__":
env_policies = [
# ("xarm", "tdmpc", ["policy.n_action_repeats=2"]),
(
"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"]),
]
# Instructions: include the policies that you want to save artifacts for here. Please make sure to revert
# your changes when you are done.
env_policies = []
for env, policy, extra_overrides in env_policies:
save_policy_to_safetensors("tests/data/save_policy_to_safetensors", env, policy, extra_overrides)