fix(nyz): fix mappo adv compute bug (#812)

opendilab · Jul 1, 2024 · 35ec39e · 35ec39e
1 parent b4ab08a
commit 35ec39e
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/ding/framework/middleware/functional/advantage_estimator.py b/ding/framework/middleware/functional/advantage_estimator.py
@@ -31,11 +31,12 @@ def gae_estimator(cfg: EasyDict, policy: Policy, buffer_: Optional[Buffer] = Non
         return task.void()
 
     model = policy.get_attribute('model')
-    # Unify the shape of obs and action
-    obs_shape = cfg['policy']['model']['obs_shape']
-    obs_shape = torch.Size(torch.tensor(obs_shape)) if isinstance(obs_shape, list) \
-        else ttorch.size.Size(convert_easy_dict_to_dict(obs_shape)) if isinstance(obs_shape, dict) \
-        else torch.Size(torch.tensor(obs_shape).unsqueeze(0))
+    if buffer_ is not None:
+        # Unify the shape of obs and action
+        obs_shape = cfg['policy']['model']['obs_shape']
+        obs_shape = torch.Size(torch.tensor(obs_shape)) if isinstance(obs_shape, list) \
+            else ttorch.size.Size(convert_easy_dict_to_dict(obs_shape)) if isinstance(obs_shape, dict) \
+            else torch.Size(torch.tensor(obs_shape).unsqueeze(0))
 
     def _gae(ctx: "OnlineRLContext"):
         """