File size: 754 Bytes
dcc4594
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
# wandb
use_wandb: true
wandb_project: nano-genie-pipeline

# dataset
dataset: PICODOOM
fps: 30
preload_ratio: 0.05

# smaller shared model params
patch_size: 4
context_length: 4
frame_size: 64
latent_dim: 6 # for video tokenizer
num_bins: 2 # for video tokenizer
n_actions: 8
embed_dim: 32
num_heads: 4
hidden_dim: 64
num_blocks: 1

# simplest train settings
compile: false
tf32: true
amp: true
log_interval: 1
n_updates: 2
learning_rate: 0.0001
batch_size_per_gpu: 8
gradient_accumulation_steps: 1

# Stage configs
video_tokenizer_config: configs/video_tokenizer.yaml
latent_actions_config: configs/latent_actions.yaml
dynamics_config: configs/dynamics.yaml

# Which stages to run
run_video_tokenizer: true
run_latent_actions: true
run_dynamics: true