workspace

workspace

Workspace for training and evaluating policies using PyTorch Lightning.

Workspace

Workspace(config, original_yaml_config)

Single workspace for training any policy using PyTorch Lightning.

This workspace handles: - Data loading and normalization - Policy instantiation and wrapping with Lightning - Trainer setup with callbacks - Policy training - WandB logging - Checkpointing

Initialize workspace.

Parameters:

Name	Type	Description	Default
config	DictConfig	Main configuration containing all settings, already instantiated	required
original_yaml_config	DictConfig	Original YAML config before instantiation, for saving	required

Source code in src/versatil/workspace.py

def __init__(self, config: DictConfig, original_yaml_config: DictConfig):
    """Initialize workspace.

    Args:
        config: Main configuration containing all settings, already instantiated
        original_yaml_config: Original YAML config before instantiation, for saving
    """
    self.config: MainConfig = config
    self.original_yaml_config = original_yaml_config
    self._normalize_quantization_workflow()
    hydra_cfg = HydraConfig.get()
    main_config_name = (
        hydra_cfg.job.config_name if hydra_cfg.job.config_name else "experiment"
    )
    additional_exp_name = config.experiment.name
    sweep_suffix = self._get_multirun_suffix(hydra_cfg)
    self.exp_name = f"{main_config_name}/{additional_exp_name}{sweep_suffix}"
    self.config.experiment.name = self.exp_name
    self.original_yaml_config.experiment.name = self.exp_name
    # Only the config basename goes on disk: checkpoint_folder already
    # carries the dataset, so the full Hydra config path would nest
    # checkpoints as <dataset>/end_to_end_training_runs/<dataset>/<model>.
    self.output_dir = (
        Path(config.experiment.checkpoint_folder)
        / Path(main_config_name).name
        / f"{additional_exp_name}{sweep_suffix}"
    )
    self.output_dir.mkdir(parents=True, exist_ok=True)
    self._set_seed()
    self.policy: Policy | None = None
    self.lightning_policy: LightningPolicy | None = None
    self.trainer: pl.Trainer | None = None
    self.train_loader: data.DataLoader | None = None
    self.val_loader: data.DataLoader | None = None
    self.normalizer: LinearNormalizer | None = None
    self.tokenizer: Tokenizer | None = None
    self.logger = None
    self.gripper_class_weights: torch.Tensor | None = None
    logging.info(f"Workspace initialized for experiment: {self.exp_name}")
    logging.info(f"Output directory: {self.output_dir}")
    self.save_config()

save_config

save_config()

Save configuration to YAML file in output directory.

The config is saved as 'config.yaml' and is required for inference and model explanation. This method should be called after workspace initialization to ensure the config is available for later use.

Source code in src/versatil/workspace.py

def save_config(self):
    """Save configuration to YAML file in output directory.

    The config is saved as 'config.yaml' and is required for inference
    and model explanation. This method should be called after workspace
    initialization to ensure the config is available for later use.
    """
    config_path = self.output_dir / "config.yaml"
    # Resolve all interpolations before saving
    resolved_config = OmegaConf.to_container(
        self.original_yaml_config, resolve=True
    )
    resolved_config_dict = OmegaConf.create(make_config_yaml_safe(resolved_config))
    OmegaConf.save(resolved_config_dict, config_path)
    logging.info(f"Config saved to {config_path}")

run

run()

Run the complete training workflow.

Source code in src/versatil/workspace.py

def run(self):
    """Run the complete training workflow."""
    self.logger = self._create_logger()
    self._setup_data()
    self._setup_policy()
    self.lightning_policy.set_dataloaders(
        train_loader=self.train_loader, val_loader=self.val_loader
    )
    self._setup_trainer()
    resume_checkpoint_path = None
    if self.config.experiment.resume_from is not None:
        checkpoint_path = Path(self.config.experiment.resume_from)
        if checkpoint_path.exists():
            logging.info(f"Resuming from checkpoint: {checkpoint_path}")
            resume_checkpoint_path = str(checkpoint_path)
        else:
            logging.warning(
                f"Checkpoint not found: {checkpoint_path}. Starting from scratch."
            )

    self._tune_hyperparameters()

    logging.info("Starting training...")
    if self.trainer is None:
        raise RuntimeError("Trainer should be initialized before training.")
    self.trainer.fit(
        model=self.lightning_policy,
        ckpt_path=resume_checkpoint_path,
        weights_only=False,
    )
    logging.info(f"Training completed. Best checkpoint saved to {self.output_dir}")