-
Notifications
You must be signed in to change notification settings - Fork 0
/
trainer.py
50 lines (35 loc) · 1.78 KB
/
trainer.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
import torch
from o3f import O3F
from config import o3f_config
from modules import Actor, EnsembledCritic
from dataset import ReplayBuffer
import wandb
from tqdm import tqdm
class O3FTrainer:
def __init__(self,
cfg=o3f_config) -> None:
self.cfg = cfg
self.device = cfg.device
self.state_dim = 17
self.action_dim = 6
self.batch_size = cfg.batch_size
actor = Actor(self.state_dim, self.action_dim, cfg.hidden_dim, edac_init=True).to(self.device)
critic = EnsembledCritic(self.state_dim, self.action_dim, cfg.hidden_dim, num_critics=cfg.num_critics).to(self.device)
self.o3f = O3F(cfg,
actor,
critic)
self.buffer = ReplayBuffer(self.state_dim, self.action_dim, cfg.buffer_size)
self.buffer.from_json(cfg.dataset_name)
def fit(self):
print(f"Training starts on {self.cfg.device} 🚀")
with wandb.init(project=self.cfg.project, entity="zzmtsvv", group=self.cfg.group, name=self.cfg.name):
wandb.config.update({k: v for k, v in self.cfg.__dict__.items() if not k.startswith("__")})
for t in tqdm(range(self.cfg.max_timesteps), desc="O3F steps"):
batch = self.buffer.sample(self.batch_size)
states, actions, rewards, next_states, dones = [x.to(self.device) for x in batch]
logging_dict = self.o3f.train(states,
actions,
rewards,
next_states,
dones)
wandb.log(logging_dict, step=self.o3f.total_iterations)