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

Fix convergence of sac, multiple torch compile on the same model caused divergence

This commit is contained in:
AdilZouitine 2025-03-31 13:54:21 +00:00
parent 8494634d48
commit 026ad463a9
3 changed files with 1 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
lerobot/common/policies/sac/modeling_sac.py
View File

@ -112,7 +112,6 @@ class SACPolicy(
self.critic_ensemble = torch.compile(self.critic_ensemble) self.critic_ensemble = torch.compile(self.critic_ensemble)
self.critic_target = torch.compile(self.critic_target) self.critic_target = torch.compile(self.critic_target)
self.actor = Policy( self.actor = Policy(
encoder=encoder_actor, encoder=encoder_actor,
network=MLP(input_dim=encoder_actor.output_dim, **asdict(config.actor_network_kwargs)), network=MLP(input_dim=encoder_actor.output_dim, **asdict(config.actor_network_kwargs)),

1
lerobot/scripts/server/actor_server.py
View File

@ -231,7 +231,6 @@ def act_with_policy(
cfg=cfg.policy, cfg=cfg.policy,
env_cfg=cfg.env, env_cfg=cfg.env,
) )
policy = torch.compile(policy)
assert isinstance(policy, nn.Module) assert isinstance(policy, nn.Module)
obs, info = online_env.reset() obs, info = online_env.reset()

4
lerobot/scripts/server/learner_server.py
View File

@ -285,9 +285,7 @@ def add_actor_information_and_train(
# ds_meta=cfg.dataset, # ds_meta=cfg.dataset,
env_cfg=cfg.env, env_cfg=cfg.env,
) )
# compile policy
policy = torch.compile(policy)
assert isinstance(policy, nn.Module) assert isinstance(policy, nn.Module)
policy.train() policy.train()