File size: 4,146 Bytes
fb90b4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """
Adapted from https://github.com/jacarvalho/mpd-public
"""
import os
import torch
from experiment_launcher import single_experiment_yaml, run_experiment
from mdoc import trainer
from mdoc.models import UNET_DIM_MULTS, TemporalUnet
from mdoc.trainer import get_dataset, get_model, get_loss, get_summary
from mdoc.trainer.trainer import get_num_epochs
from torch_robotics.torch_utils.seed import fix_random_seed
from torch_robotics.torch_utils.torch_utils import get_torch_device
os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE"
@single_experiment_yaml
def experiment(
########################################################################
# Dataset
dataset_subdir: str = 'EnvSimple2D-RobotPointMass',
# dataset_subdir: str = 'EnvSpheres3D-RobotPanda',
include_velocity: bool = True,
########################################################################
# Diffusion Model
diffusion_model_class: str = 'GaussianDiffusionModel',
variance_schedule: str = 'exponential', # cosine
n_diffusion_steps: int = 25,
predict_epsilon: bool = True,
# Unet
unet_input_dim: int = 32,
unet_dim_mults_option: int = 1,
########################################################################
# Loss
loss_class: str = 'GaussianDiffusionLoss',
# Training parameters
batch_size: int = 32,
lr: float = 1e-4,
num_train_steps: int = 500000,
use_ema: bool = True,
use_amp: bool = False,
# Summary parameters
steps_til_summary: int = 10,
summary_class: str = 'SummaryTrajectoryGeneration',
steps_til_ckpt: int = 50000,
########################################################################
device: str = 'cuda',
debug: bool = True,
########################################################################
# MANDATORY
seed: int = 0,
results_dir: str = 'logs',
########################################################################
# WandB
wandb_mode: str = 'disabled', # "online", "offline" or "disabled"
wandb_entity: str = 'scoreplan',
wandb_project: str = 'test_train',
**kwargs
):
fix_random_seed(seed)
device = get_torch_device(device=device)
tensor_args = {'device': device, 'dtype': torch.float32}
# Dataset
train_subset, train_dataloader, val_subset, val_dataloader = get_dataset(
dataset_class='TrajectoryDataset',
include_velocity=include_velocity,
dataset_subdir=dataset_subdir,
batch_size=batch_size,
results_dir=results_dir,
save_indices=True,
tensor_args=tensor_args
)
dataset = train_subset.dataset
# Model
diffusion_configs = dict(
variance_schedule=variance_schedule,
n_diffusion_steps=n_diffusion_steps,
predict_epsilon=predict_epsilon,
)
unet_configs = dict(
state_dim=dataset.state_dim,
n_support_points=dataset.n_support_points,
unet_input_dim=unet_input_dim,
dim_mults=UNET_DIM_MULTS[unet_dim_mults_option],
)
model = get_model(
model_class=diffusion_model_class,
model=TemporalUnet(**unet_configs),
tensor_args=tensor_args,
**diffusion_configs,
**unet_configs
)
# Loss
loss_fn = val_loss_fn = get_loss(
loss_class=loss_class
)
# Summary
summary_fn = get_summary(
summary_class=summary_class,
)
# Train
trainer.train(
model=model,
train_dataloader=train_dataloader,
train_subset=train_subset,
val_dataloader=val_dataloader,
val_subset=train_subset,
epochs=get_num_epochs(num_train_steps, batch_size, len(dataset)),
model_dir=results_dir,
summary_fn=summary_fn,
lr=lr,
loss_fn=loss_fn,
val_loss_fn=val_loss_fn,
steps_til_summary=steps_til_summary,
steps_til_checkpoint=steps_til_ckpt,
clip_grad=True,
use_ema=use_ema,
use_amp=use_amp,
debug=debug,
tensor_args=tensor_args
)
if __name__ == '__main__':
# Leave unchanged
run_experiment(experiment)
|