hongshaorou/tdmpc2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix multitask model api conversion

Browse Source
This commit is contained in:
Nicklas Hansen
2025-02-27 16:25:21 -08:00
parent 55bde9745f
commit 0a914570dc
2 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
tdmpc2/common/layers.py
View File

@@ -212,6 +212,7 @@ def api_model_conversion(target_state_dict, source_state_dict):
# copy log_std_min and log_std_max from target_state_dict to new_state_dict
new_state_dict['log_std_min'] = target_state_dict['log_std_min']
new_state_dict['log_std_dif'] = target_state_dict['log_std_dif']
new_state_dict['_action_masks'] = target_state_dict['_action_masks']
# copy new_state_dict to source_state_dict
source_state_dict.update(new_state_dict)

5
tdmpc2/tdmpc2.py
View File

@@ -83,7 +83,10 @@ class TDMPC2(torch.nn.Module):
Args:
fp (str or dict): Filepath or state dict to load.
"""
state_dict = fp if isinstance(fp, dict) else torch.load(fp, map_location=torch.get_default_device())
if isinstance(fp, dict):
state_dict = fp
else:
state_dict = torch.load(fp, map_location=torch.get_default_device(), weights_only=False)
state_dict = state_dict["model"] if "model" in state_dict else state_dict
state_dict = api_model_conversion(self.model.state_dict(), state_dict)
self.model.load_state_dict(state_dict)