- Refactor observation encoder in modeling_sac.py

- added `torch.compile` to the actor and learner servers.
- organized imports in `train_sac.py`
- optimized the parameters push by not sending the frozen pre-trained encoder.

Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com>
This commit is contained in:
Michel Aractingi
2025-01-31 16:45:52 +00:00
committed by AdilZouitine
parent faab32fe14
commit b29401e4e2
6 changed files with 199 additions and 85 deletions

View File

@@ -52,6 +52,8 @@ policy:
n_action_steps: 1
shared_encoder: true
# vision_encoder_name: null
freeze_vision_encoder: false
input_shapes:
# # TODO(rcadene, alexander-soare): add variables for height and width from the dataset/env?
observation.state: ["${env.state_dim}"]