all the fixes

montrealrobotics · Sep 21, 2023 · 96e0238 · 96e0238
1 parent 167ca8e
commit 96e0238
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cleanrl/ppo_continuous_action_wandb.py b/cleanrl/ppo_continuous_action_wandb.py
@@ -366,7 +366,7 @@ def train_risk(cfg, model, data, criterion, opt, device):
     dataloader = DataLoader(dataset, batch_size=cfg.risk_batch_size, shuffle=True, num_workers=10, generator=torch.Generator(device='cpu'))
     net_loss = 0
     for batch in dataloader:
-        pred = model(batch[0].to(device))
+        pred = model(get_risk_obs(cfg, batch[0]).to(device))
         if cfg.model_type == "mlp":
             loss = criterion(pred, batch[1].squeeze().to(device))
         else: