logging state specs

HorizonRobotics · Sep 29, 2023 · 7746aef · 7746aef
1 parent fd26365
commit 7746aef
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/alf/trainers/policy_trainer.py b/alf/trainers/policy_trainer.py
@@ -556,6 +556,13 @@ def __init__(self, config: TrainerConfig, ddp_rank: int = -1):
             config=self._config,
             debug_summaries=self._debug_summaries)
 
+        logging.info("predict_state_spec=\n%s" % pformat_pycolor(
+            self._algorithm.predict_state_spec))
+        logging.info("rollout_state_spec=\n%s" % pformat_pycolor(
+            self._algorithm.rollout_state_spec))
+        logging.info("train_state_spec=\n%s" % pformat_pycolor(
+            self._algorithm.train_state_spec))
+
         # recover offline buffer
         self._algorithm.load_offline_replay_buffer(
             untransformed_observation_spec)