fix(hil-serl): discrete critic send through network (#1468)

Co-authored-by: Khalil Meftah <kmeftah.khalil@gmail.com>
Co-authored-by: jpizarrom <jpizarrom@gmail.com>
This commit is contained in:
Adil Zouitine
2025-07-09 16:22:40 +02:00
committed by GitHub
parent cf86b9300d
commit ce2b9724bf
6 changed files with 81 additions and 51 deletions

View File

@@ -1109,8 +1109,18 @@ def check_nan_in_transition(
def push_actor_policy_to_queue(parameters_queue: Queue, policy: nn.Module):
logging.debug("[LEARNER] Pushing actor policy to the queue")
state_dict = move_state_dict_to_device(policy.actor.state_dict(), device="cpu")
state_bytes = state_to_bytes(state_dict)
# Create a dictionary to hold all the state dicts
state_dicts = {"policy": move_state_dict_to_device(policy.actor.state_dict(), device="cpu")}
# Add discrete critic if it exists
if hasattr(policy, "discrete_critic") and policy.discrete_critic is not None:
state_dicts["discrete_critic"] = move_state_dict_to_device(
policy.discrete_critic.state_dict(), device="cpu"
)
logging.debug("[LEARNER] Including discrete critic in state dict push")
state_bytes = state_to_bytes(state_dicts)
parameters_queue.put(state_bytes)