[BugFix] Add keys to GAE in PPO/A2C (#1618)

Vincent Moens · web-flow · commit 805918c8924b · 2023-10-09T20:31:45.000+01:00
diff --git a/examples/a2c/a2c_atari.py b/examples/a2c/a2c_atari.py
@@ -76,6 +76,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
     )
 
     # use end-of-life as done key
+    adv_module.set_keys(done="end-of-life", terminated="end-of-life")
     loss_module.set_keys(done="end-of-life", terminated="end-of-life")
 
     # Create optimizer
diff --git a/examples/ppo/ppo_atari.py b/examples/ppo/ppo_atari.py
@@ -79,6 +79,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
     )
 
     # use end-of-life as done key
+    adv_module.set_keys(done="end-of-life", terminated="end-of-life")
     loss_module.set_keys(done="end-of-life", terminated="end-of-life")
 
     # Create optimizer

Original file line number	Diff line number	Diff line change
`@@ -76,6 +76,7 @@ def main(cfg: "DictConfig"): # noqa: F821`
`76`	`76`	`)`
`77`	`77`
`78`	`78`	`# use end-of-life as done key`
	`79`	`+ adv_module.set_keys(done="end-of-life", terminated="end-of-life")`
`79`	`80`	`loss_module.set_keys(done="end-of-life", terminated="end-of-life")`
`80`	`81`
`81`	`82`	`# Create optimizer`
Original file line number	Diff line number	Diff line change
`@@ -79,6 +79,7 @@ def main(cfg: "DictConfig"): # noqa: F821`
`79`	`79`	`)`
`80`	`80`
`81`	`81`	`# use end-of-life as done key`
	`82`	`+ adv_module.set_keys(done="end-of-life", terminated="end-of-life")`
`82`	`83`	`loss_module.set_keys(done="end-of-life", terminated="end-of-life")`
`83`	`84`
`84`	`85`	`# Create optimizer`