fix DiscreteRLPolicy set_state bug

microsoft · May 15, 2023 · a9a6381 · a9a6381
1 parent ef2a358
commit a9a6381
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/maro/rl/policy/discrete_rl_policy.py b/maro/rl/policy/discrete_rl_policy.py
@@ -244,7 +244,7 @@ def get_state(self) -> dict:
         }
 
     def set_state(self, policy_state: dict) -> None:
-        self._q_net.set_state(policy_state)
+        self._q_net.set_state(policy_state["net"])
         self._warmup = policy_state["policy"]["warmup"]
         self._call_count = policy_state["policy"]["call_count"]