generated from siddk/mjolnir
-
Notifications
You must be signed in to change notification settings - Fork 2
/
train.py
151 lines (129 loc) · 5.35 KB
/
train.py
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
149
150
151
"""
train.py
Core training script -- loads and preprocesses collected Panda demonstrations for a given task (or set of tasks),
instantiates a Lightning Module, and runs LILA or Imitation Learning training!
To be extended with support for multiple tasks, baselines, and eventually, support for language instructions!
Run with: `python train.py --config conf/lila-mturk-all-config.yaml`
"""
import os
from datetime import datetime
import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from quinine import QuinineArgumentParser
from torch.utils.data import DataLoader
from src.models import GCAE, FiLM, Imitation
from src.overwatch import MetricsLogger, get_overwatch
from src.preprocessing import get_imitation_dataset, get_la_dataset, get_lila_dataset
from src.util import create_paths
# Disable Tokenizers Parallelism to Avoid Deadlock
os.environ["TOKENIZERS_PARALLELISM"] = "false"
def train() -> None:
# Parse Quinfig (via Quinine Argparse Binding)
print("[*] LILA :: Training Latent Actions =>>>")
quinfig = QuinineArgumentParser().parse_quinfig()
# Create Unique Run Name (for Logging, Checkpointing) :: Initialize all Directories
run_id = quinfig.run_id
if run_id is None:
run_id = f"lila-{quinfig.mode}-{quinfig.arch}+{datetime.now().strftime('%Y-%m-%d-%H:%M')}"
paths = create_paths(run_id, quinfig.run_dir)
# Overwatch :: Setup & Configure Console/File Logger --> Handle Process 0 vs. other Process Logging!
overwatch = get_overwatch(paths["runs"], run_id, quinfig.log_level)
overwatch.info(f"Starting Run: {run_id}...")
# Set Randomness
overwatch.info(f"Setting Random Seed to {quinfig.seed}!")
np.random.seed(quinfig.seed)
torch.manual_seed(quinfig.seed)
# Build Dataset & DataLoaders --> may want to play with the default parameters to this function!
overwatch.info(f"Building Dataset using Demonstrations located at `{quinfig.demonstrations}`...")
# Obtain Demonstrations based on "modality" of architectures
if "lila" in quinfig.arch:
train, validation = get_lila_dataset(
quinfig.demonstrations,
paths["runs"],
augment=quinfig.augment,
augment_factor=quinfig.augmentation_factor,
seed=quinfig.seed,
)
elif "no-lang" in quinfig.arch:
train, validation = get_la_dataset(
quinfig.demonstrations, augment=quinfig.augment, augment_factor=quinfig.augment_factor, seed=quinfig.seed
)
elif "imitation" in quinfig.arch:
train, validation = get_imitation_dataset(
quinfig.demonstrations,
run_dir=paths["runs"],
augment=quinfig.augment,
augment_factor=quinfig.augmentation_factor,
seed=quinfig.seed,
)
else:
raise NotImplementedError(f"Dataset Loading for Architecture {quinfig.arch} isn't supported...")
# Create DataLoaders
train_loader = DataLoader(train, batch_size=quinfig.bsz, shuffle=True)
val_loader = DataLoader(validation, batch_size=quinfig.bsz, shuffle=False)
# Create Model (one of multiple architectures)
if quinfig.arch in ["lila"]:
overwatch.info("Initializing LILA :: FiLM-GeLU Conditional Autoencoder...")
nn = FiLM(
quinfig.state_dim,
quinfig.action_dim,
768,
quinfig.latent_dim,
hidden_dim=quinfig.hidden_dim,
lr=quinfig.lr,
lr_step_size=quinfig.lr_step_size,
lr_gamma=quinfig.lr_gamma,
zaug=True,
zaug_lambda=quinfig.zaug_lambda,
retrieval=True,
run_dir=paths["runs"],
)
elif quinfig.arch in ["no-lang"]:
overwatch.info("Initializing No-Language Latent Actions :: GeLU Conditional Autoencoder...")
nn = GCAE(
quinfig.state_dim,
quinfig.action_dim,
quinfig.latent_dim,
hidden_dim=quinfig.hidden_dim,
lr=quinfig.lr,
lr_step_size=quinfig.lr_step_size,
lr_gamme=quinfig.lr_gamma,
zaug=True,
zaug_lambda=quinfig.zaug_lambda,
run_dir=paths["runs"],
)
elif quinfig.arch in ["imitation"]:
overwatch.info("Initializing Imitation Learning :: FiLM-GELU Behavioral Cloning...")
nn = Imitation(
quinfig.state_dim,
quinfig.action_dim,
768,
hidden_dim=quinfig.hidden_dim,
lr=quinfig.lr,
lr_step_size=quinfig.lr_step_size,
lr_gamma=quinfig.lr_gamma,
unnatural=True,
run_dir=paths["runs"],
)
else:
raise NotImplementedError(f"Model `{quinfig.arch}` not implemented!")
# Create Callbacks
checkpoint_callback = ModelCheckpoint(
dirpath=paths["runs"],
filename=f"{run_id}+" + "{train_loss:.2f}-{val_loss:.2f}",
monitor="val_loss",
mode="min",
save_top_k=1,
)
early_stopping = EarlyStopping(monitor="val_loss")
logger = MetricsLogger(name=run_id, save_dir=paths["runs"])
# Train!
overwatch.info("Training...")
trainer = pl.Trainer(
max_epochs=quinfig.epochs, gpus=quinfig.gpus, logger=logger, callbacks=[checkpoint_callback, early_stopping]
)
trainer.fit(nn, train_loader, val_loader)
if __name__ == "__main__":
train()