wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Cleanup, add output_dict

This commit is contained in:
Simon Alibert 2024-05-01 19:55:29 +02:00
parent 7ee03afabb
commit 7002813bbb
2 changed files with 16 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
tests/scripts/save_policy_to_safetensor.py
View File

@ -37,7 +37,6 @@ def get_policy_stats(env_name, policy_name):
batch = next(iter(dataloader))
output_dict = policy.forward(batch)
loss = output_dict["loss"]
# TODO Check output dict values
loss.backward()
grad_stats = {}
@ -63,12 +62,8 @@ def get_policy_stats(env_name, policy_name):
if k in ["observation.image", "observation.images.top", "observation.state"]
}
# TODO(aliberts): refacor `select_action` methods so that it expects `obs` instead of `batch`
# if policy_name == "diffusion":
# obs = {k: batch[k] for k in ["observation.image", "observation.state"]}
actions = {str(i): policy.select_action(obs).contiguous() for i in range(cfg.policy.n_action_steps)}
return grad_stats, param_stats, actions
return output_dict, grad_stats, param_stats, actions
def save_policy_to_safetensors(output_dir, env_name, policy_name):
@ -78,21 +73,18 @@ def save_policy_to_safetensors(output_dir, env_name, policy_name):
shutil.rmtree(env_policy_dir)
env_policy_dir.mkdir(parents=True, exist_ok=True)
grad_stats, param_stats, actions = get_policy_stats(env_name, policy_name)
save_file(grad_stats, env_policy_dir / "grad_stats")
save_file(param_stats, env_policy_dir / "param_stats")
save_file(actions, env_policy_dir / "actions")
output_dict, grad_stats, param_stats, actions = get_policy_stats(env_name, policy_name)
save_file(output_dict, env_policy_dir / "output_dict.safetensors")
save_file(grad_stats, env_policy_dir / "grad_stats.safetensors")
save_file(param_stats, env_policy_dir / "param_stats.safetensors")
save_file(actions, env_policy_dir / "actions.safetensors")
if __name__ == "__main__":
env_policies = [
# ("xarm", "tdmpc", ["policy.mpc=true"]),
# ("pusht", "tdmpc", ["policy.mpc=false"]),
(
"pusht",
"diffusion",
),
# ("aloha", "act"),
# ("xarm", "tdmpc"),
("pusht", "diffusion"),
("aloha", "act"),
]
for env, policy in env_policies:
save_policy_to_safetensors("tests/data/save_policy_to_safetensors", env, policy)

11
tests/test_regression.py
View File

@ -17,12 +17,15 @@ from tests.scripts.save_policy_to_safetensor import get_policy_stats
)
def test_backward_compatibility(env_name, policy_name):
env_policy_dir = Path("tests/data/save_policy_to_safetensors") / f"{env_name}_{policy_name}"
saved_grad_stats = load_file(env_policy_dir / "grad_stats")
saved_param_stats = load_file(env_policy_dir / "param_stats")
saved_actions = load_file(env_policy_dir / "actions")
saved_output_dict = load_file(env_policy_dir / "output_dict.safetensors")
saved_grad_stats = load_file(env_policy_dir / "grad_stats.safetensors")
saved_param_stats = load_file(env_policy_dir / "param_stats.safetensors")
saved_actions = load_file(env_policy_dir / "actions.safetensors")
grad_stats, param_stats, actions = get_policy_stats(env_name, policy_name)
output_dict, grad_stats, param_stats, actions = get_policy_stats(env_name, policy_name)
for key in saved_output_dict:
assert torch.isclose(output_dict[key], saved_output_dict[key]).all()
for key in saved_grad_stats:
assert torch.isclose(grad_stats[key], saved_grad_stats[key]).all()
for key in saved_param_stats: