diff --git a/examples/3_train_policy.py b/examples/3_train_policy.py index d3467562..7a7a7aaf 100644 --- a/examples/3_train_policy.py +++ b/examples/3_train_policy.py @@ -53,16 +53,10 @@ step = 0 done = False while not done: for batch in dataloader: - for k in batch: - batch[k] = batch[k].to(device, non_blocking=True) + batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()} info = policy.update(batch) if step % log_freq == 0: - num_samples = (step + 1) * cfg.batch_size - loss = info["loss"] - update_s = info["update_s"] - print( - f"step: {step} samples: {num_samples} loss: {loss:.3f} update_time: {update_s:.3f} (seconds)" - ) + print(f"step: {step} loss: {info['loss']:.3f} update_time: {info['update_s']:.3f} (seconds)") step += 1 if step >= training_steps: done = True diff --git a/lerobot/common/policies/act/modeling_act.py b/lerobot/common/policies/act/modeling_act.py index 567721cd..143d1459 100644 --- a/lerobot/common/policies/act/modeling_act.py +++ b/lerobot/common/policies/act/modeling_act.py @@ -65,12 +65,16 @@ class ActionChunkingTransformerPolicy(nn.Module): "ActionChunkingTransformerPolicy does not handle multiple observation steps." ) - def __init__(self, cfg: ActionChunkingTransformerConfig): + def __init__(self, cfg: ActionChunkingTransformerConfig | None = None): """ - TODO(alexander-soare): Add documentation for all parameters once we have model configs established. + Args: + cfg: Policy configuration class instance or None, in which case the default instantiation of the + configuration class is used. """ super().__init__() - if getattr(cfg, "n_obs_steps", 1) != 1: + if cfg is None: + cfg = ActionChunkingTransformerConfig() + if cfg.n_obs_steps != 1: raise ValueError(self._multiple_obs_steps_not_handled_msg) self.cfg = cfg @@ -161,8 +165,8 @@ class ActionChunkingTransformerPolicy(nn.Module): if self.cfg.n_action_steps is not None: self._action_queue = deque([], maxlen=self.cfg.n_action_steps) - def forward(self, batch: dict[str, Tensor], **_) -> Tensor: - return self.select_action(self, batch) + # def forward(self, batch: dict[str, Tensor], **_) -> Tensor: + # return self.select_action(self, batch) @torch.no_grad def select_action(self, batch: dict[str, Tensor], **_) -> Tensor: diff --git a/lerobot/common/policies/diffusion/modeling_diffusion.py b/lerobot/common/policies/diffusion/modeling_diffusion.py index 9a02c6a2..dfab9bb7 100644 --- a/lerobot/common/policies/diffusion/modeling_diffusion.py +++ b/lerobot/common/policies/diffusion/modeling_diffusion.py @@ -33,8 +33,6 @@ from lerobot.common.policies.utils import ( populate_queues, ) -logger = logging.getLogger(__name__) - class DiffusionPolicy(nn.Module): """ @@ -44,8 +42,17 @@ class DiffusionPolicy(nn.Module): name = "diffusion" - def __init__(self, cfg: DiffusionConfig, lr_scheduler_num_training_steps: int): + def __init__(self, cfg: DiffusionConfig | None, lr_scheduler_num_training_steps: int = 0): super().__init__() + """ + Args: + cfg: Policy configuration class instance or None, in which case the default instantiation of the + configuration class is used. + """ + # TODO(alexander-soare): LR scheduler will be removed. + assert lr_scheduler_num_training_steps > 0 + if cfg is None: + cfg = DiffusionConfig() self.cfg = cfg # queues are populated during rollout of the policy, they contain the n latest observations and actions diff --git a/tests/test_examples.py b/tests/test_examples.py index 6cab7a1a..c510eb1e 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -40,7 +40,8 @@ def test_examples_3_and_2(): ], ) - exec(file_contents) + # Pass empty globals to allow dictionary comprehension https://stackoverflow.com/a/32897127/4391249. + exec(file_contents, {}) for file_name in ["model.pt", "stats.pth", "config.yaml"]: assert Path(f"outputs/train/example_pusht_diffusion/{file_name}").exists()