dodox/dppo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

v0.6 (#18)

Browse Source 
* Sampling over both env and denoising steps in DPPO updates (#13)

* sample one from each chain

* full random sampling

* Add Proficient Human (PH) Configs and Pipeline (#16)

* fix missing cfg

* add ph config

* fix how terminated flags are added to buffer in ibrl

* add ph config

* offline calql for 1M gradient updates

* bug fix: number of calql online gradient steps is the number of new transitions collected

* add sample config for DPPO with ta=1

* Sampling over both env and denoising steps in DPPO updates (#13)

* sample one from each chain

* full random sampling

* fix diffusion loss when predicting initial noise

* fix dppo inds

* fix typo

* remove print statement

---------

Co-authored-by: Justin M. Lidard <jlidard@neuronic.cs.princeton.edu>
Co-authored-by: allenzren <allen.ren@princeton.edu>

* update robomimic configs

* better calql formulation

* optimize calql and ibrl training

* optimize data transfer in ppo agents

* add kitchen configs

* re-organize config folders, rerun calql and rlpd

* add scratch gym locomotion configs

* add kitchen installation dependencies

* use truncated for termination in furniture env

* update furniture and gym configs

* update README and dependencies with kitchen

* add url for new data and checkpoints

* update demo RL configs

* update batch sizes for furniture unet configs

* raise error about dropout in residual mlp

* fix observation bug in bc loss

---------

Co-authored-by: Justin Lidard <60638575+jlidard@users.noreply.github.com>
Co-authored-by: Justin M. Lidard <jlidard@neuronic.cs.princeton.edu>
This commit is contained in:
Allen Z. Ren
2024-10-30 19:58:06 -04:00
committed by GitHub
co-authored by Justin M. Lidard Justin Lidard
parent 7b10df690d
commit dc8e0c9edc
126 changed files with 4614 additions and 553 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agent/finetune/train_calql_agent.py
+9 -7
View File
@@ -82,8 +82,6 @@ class TrainCalQLAgent(TrainAgent):
if self.train_online:
# number of episode to colect per epoch for training
self.n_episode_per_epoch = cfg.train.n_episode_per_epoch
# UTD ratio
self.online_utd_ratio = cfg.train.online_utd_ratio
# Eval episodes
self.n_eval_episode = cfg.train.n_eval_episode
@@ -204,9 +202,13 @@ class TrainCalQLAgent(TrainAgent):
action_venv = samples[:, : self.act_steps]
# Apply multi-step action
obs_venv, reward_venv, terminated_venv, truncated_venv, info_venv = (
self.venv.step(action_venv)
)
(
obs_venv,
reward_venv,
terminated_venv,
truncated_venv,
info_venv,
) = self.venv.step(action_venv)
done_venv = terminated_venv | truncated_venv
reward_trajs[step] = reward_venv
firsts_trajs[step + 1] = done_venv
@@ -308,7 +310,8 @@ class TrainCalQLAgent(TrainAgent):
# override num_update
if self.train_online:
num_update = len(reward_trajs) # assume one env!
# the amount of new transitions(single env)
num_update = len(reward_trajs_split[0])
else:
num_update = self.num_update
for _ in range(num_update):
@@ -413,7 +416,6 @@ class TrainCalQLAgent(TrainAgent):
reward_to_go_b,
terminated_b,
self.gamma,
alpha,
)
self.critic_optimizer.zero_grad()
loss_critic.backward()
agent/finetune/train_ibrl_agent.py
+8 -6
View File
@@ -145,7 +145,6 @@ class TrainIBRLAgent(TrainAgent):
# Collect a set of trajectories from env
cnt_episode = 0
for step in range(n_steps):
# Select action
with torch.no_grad():
cond = {
@@ -164,9 +163,13 @@ class TrainIBRLAgent(TrainAgent):
action_venv = samples[:, : self.act_steps]
# Apply multi-step action
obs_venv, reward_venv, terminated_venv, truncated_venv, info_venv = (
self.venv.step(action_venv)
)
(
obs_venv,
reward_venv,
terminated_venv,
truncated_venv,
info_venv,
) = self.venv.step(action_venv)
done_venv = terminated_venv | truncated_venv
reward_trajs[step] = reward_venv
firsts_trajs[step + 1] = done_venv
@@ -177,14 +180,13 @@ class TrainIBRLAgent(TrainAgent):
obs_buffer.append(prev_obs_venv["state"][i])
if "final_obs" in info_venv[i]: # truncated
next_obs_buffer.append(info_venv[i]["final_obs"]["state"])
terminated_venv[i] = False
else: # first obs in new episode
next_obs_buffer.append(obs_venv["state"][i])
action_buffer.append(action_venv[i])
reward_buffer.extend(
(reward_venv * self.scale_reward_factor).tolist()
)
terminated_buffer.append(terminated_venv.tolist())
terminated_buffer.extend(terminated_venv.tolist())
# update for next step
prev_obs_venv = obs_venv
agent/finetune/train_ppo_diffusion_agent.py
+30 -22
View File
@@ -19,7 +19,6 @@ from util.scheduler import CosineAnnealingWarmupRestarts
class TrainPPODiffusionAgent(TrainPPOAgent):
def __init__(self, cfg):
super().__init__(cfg)
@@ -46,7 +45,6 @@ class TrainPPODiffusionAgent(TrainPPOAgent):
)
def run(self):
# Start training loop
timer = Timer()
run_results = []
@@ -54,7 +52,6 @@ class TrainPPODiffusionAgent(TrainPPOAgent):
last_itr_eval = False
done_venv = np.zeros((1, self.n_envs))
while self.itr < self.n_train_itr:
# Prepare video paths for each envs --- only applies for the first set of episodes if allowing reset within iteration and each iteration has multiple episodes from one env
options_venv = [{} for _ in range(self.n_envs)]
if self.itr % self.render_freq == 0 and self.render_video:
@@ -126,9 +123,13 @@ class TrainPPODiffusionAgent(TrainPPOAgent):
action_venv = output_venv[:, : self.act_steps]
# Apply multi-step action
obs_venv, reward_venv, terminated_venv, truncated_venv, info_venv = (
self.venv.step(action_venv)
)
(
obs_venv,
reward_venv,
terminated_venv,
truncated_venv,
info_venv,
) = self.venv.step(action_venv)
done_venv = terminated_venv | truncated_venv
if self.save_full_observations: # state-only
obs_full_venv = np.array(
@@ -285,40 +286,45 @@ class TrainPPODiffusionAgent(TrainPPOAgent):
)
}
chains_k = einops.rearrange(
torch.tensor(chains_trajs).float().to(self.device),
torch.tensor(chains_trajs, device=self.device).float(),
"s e t h d -> (s e) t h d",
)
returns_k = (
torch.tensor(returns_trajs).float().to(self.device).reshape(-1)
torch.tensor(returns_trajs, device=self.device).float().reshape(-1)
)
values_k = (
torch.tensor(values_trajs).float().to(self.device).reshape(-1)
torch.tensor(values_trajs, device=self.device).float().reshape(-1)
)
advantages_k = (
torch.tensor(advantages_trajs).float().to(self.device).reshape(-1)
torch.tensor(advantages_trajs, device=self.device)
.float()
.reshape(-1)
)
logprobs_k = torch.tensor(logprobs_trajs).float().to(self.device)
logprobs_k = torch.tensor(logprobs_trajs, device=self.device).float()
# Update policy and critic
total_steps = self.n_steps * self.n_envs
inds_k = np.arange(total_steps)
total_steps = self.n_steps * self.n_envs * self.model.ft_denoising_steps
clipfracs = []
for update_epoch in range(self.update_epochs):
# for each epoch, go through all data in batches
flag_break = False
np.random.shuffle(inds_k)
inds_k = torch.randperm(total_steps, device=self.device)
num_batch = max(1, total_steps // self.batch_size) # skip last ones
for batch in range(num_batch):
start = batch * self.batch_size
end = start + self.batch_size
inds_b = inds_k[start:end] # b for batch
obs_b = {"state": obs_k["state"][inds_b]}
chains_b = chains_k[inds_b]
returns_b = returns_k[inds_b]
values_b = values_k[inds_b]
advantages_b = advantages_k[inds_b]
logprobs_b = logprobs_k[inds_b]
batch_inds_b, denoising_inds_b = torch.unravel_index(
inds_b,
(self.n_steps * self.n_envs, self.model.ft_denoising_steps),
)
obs_b = {"state": obs_k["state"][batch_inds_b]}
chains_prev_b = chains_k[batch_inds_b, denoising_inds_b]
chains_next_b = chains_k[batch_inds_b, denoising_inds_b + 1]
returns_b = returns_k[batch_inds_b]
values_b = values_k[batch_inds_b]
advantages_b = advantages_k[batch_inds_b]
logprobs_b = logprobs_k[batch_inds_b, denoising_inds_b]
# get loss
(
@@ -332,7 +338,9 @@ class TrainPPODiffusionAgent(TrainPPOAgent):
eta,
) = self.model.loss(
obs_b,
chains_b,
chains_prev_b,
chains_next_b,
denoising_inds_b,
returns_b,
values_b,
advantages_b,
agent/finetune/train_ppo_diffusion_img_agent.py
+21 -15
View File
@@ -283,40 +283,44 @@ class TrainPPOImgDiffusionAgent(TrainPPODiffusionAgent):
for k in obs_trajs
}
chains_k = einops.rearrange(
torch.tensor(chains_trajs).float().to(self.device),
torch.tensor(chains_trajs, device=self.device).float(),
"s e t h d -> (s e) t h d",
)
returns_k = (
torch.tensor(returns_trajs).float().to(self.device).reshape(-1)
torch.tensor(returns_trajs, device=self.device).float().reshape(-1)
)
values_k = (
torch.tensor(values_trajs).float().to(self.device).reshape(-1)
torch.tensor(values_trajs, device=self.device).float().reshape(-1)
)
advantages_k = (
torch.tensor(advantages_trajs).float().to(self.device).reshape(-1)
torch.tensor(advantages_trajs, device=self.device).float().reshape(-1)
)
logprobs_k = torch.tensor(logprobs_trajs).float().to(self.device)
logprobs_k = torch.tensor(logprobs_trajs, device=self.device).float()
# Update policy and critic
total_steps = self.n_steps * self.n_envs
inds_k = np.arange(total_steps)
total_steps = self.n_steps * self.n_envs * self.model.ft_denoising_steps
clipfracs = []
for update_epoch in range(self.update_epochs):
# for each epoch, go through all data in batches
flag_break = False
np.random.shuffle(inds_k)
inds_k = torch.randperm(total_steps, device=self.device)
num_batch = max(1, total_steps // self.batch_size) # skip last ones
for batch in range(num_batch):
start = batch * self.batch_size
end = start + self.batch_size
inds_b = inds_k[start:end] # b for batch
obs_b = {k: obs_k[k][inds_b] for k in obs_k}
chains_b = chains_k[inds_b]
returns_b = returns_k[inds_b]
values_b = values_k[inds_b]
advantages_b = advantages_k[inds_b]
logprobs_b = logprobs_k[inds_b]
batch_inds_b, denoising_inds_b = torch.unravel_index(
inds_b,
(self.n_steps * self.n_envs, self.model.ft_denoising_steps),
)
obs_b = {k: obs_k[k][batch_inds_b] for k in obs_k}
chains_prev_b = chains_k[batch_inds_b, denoising_inds_b]
chains_next_b = chains_k[batch_inds_b, denoising_inds_b + 1]
returns_b = returns_k[batch_inds_b]
values_b = values_k[batch_inds_b]
advantages_b = advantages_k[batch_inds_b]
logprobs_b = logprobs_k[batch_inds_b, denoising_inds_b]
# get loss
(
@@ -330,7 +334,9 @@ class TrainPPOImgDiffusionAgent(TrainPPODiffusionAgent):
eta,
) = self.model.loss(
obs_b,
chains_b,
chains_prev_b,
chains_next_b,
denoising_inds_b,
returns_b,
values_b,
advantages_b,
agent/finetune/train_ppo_exact_diffusion_agent.py
+6 -7
View File
@@ -249,29 +249,28 @@ class TrainPPOExactDiffusionAgent(TrainPPODiffusionAgent):
)
}
samples_k = einops.rearrange(
torch.tensor(samples_trajs).float().to(self.device),
torch.tensor(samples_trajs, device=self.device).float(),
"s e h d -> (s e) h d",
)
returns_k = (
torch.tensor(returns_trajs).float().to(self.device).reshape(-1)
torch.tensor(returns_trajs, device=self.device).float().reshape(-1)
)
values_k = (
torch.tensor(values_trajs).float().to(self.device).reshape(-1)
torch.tensor(values_trajs, device=self.device).float().reshape(-1)
)
advantages_k = (
torch.tensor(advantages_trajs).float().to(self.device).reshape(-1)
torch.tensor(advantages_trajs, device=self.device).float().reshape(-1)
)
logprobs_k = torch.tensor(logprobs_trajs).float().to(self.device)
logprobs_k = torch.tensor(logprobs_trajs, device=self.device).float()
# Update policy and critic
total_steps = self.n_steps * self.n_envs
inds_k = np.arange(total_steps)
clipfracs = []
for update_epoch in range(self.update_epochs):
# for each epoch, go through all data in batches
flag_break = False
np.random.shuffle(inds_k)
inds_k = torch.randperm(total_steps, device=self.device)
num_batch = max(1, total_steps // self.batch_size) # skip last ones
for batch in range(num_batch):
start = batch * self.batch_size
agent/finetune/train_ppo_gaussian_agent.py
+7 -10
View File
@@ -210,7 +210,7 @@ class TrainPPOGaussianAgent(TrainPPOAgent):
)
reward_trajs = reward_trajs_transpose.T
# bootstrap value with GAE if not done - apply reward scaling with constant if specified
# bootstrap value with GAE if not terminal - apply reward scaling with constant if specified
obs_venv_ts = {
"state": torch.from_numpy(obs_venv["state"])
.float()
@@ -250,31 +250,28 @@ class TrainPPOGaussianAgent(TrainPPOAgent):
)
}
samples_k = einops.rearrange(
torch.tensor(samples_trajs).float().to(self.device),
torch.tensor(samples_trajs, device=self.device).float(),
"s e h d -> (s e) h d",
)
returns_k = (
torch.tensor(returns_trajs).float().to(self.device).reshape(-1)
torch.tensor(returns_trajs, device=self.device).float().reshape(-1)
)
values_k = (
torch.tensor(values_trajs).float().to(self.device).reshape(-1)
torch.tensor(values_trajs, device=self.device).float().reshape(-1)
)
advantages_k = (
torch.tensor(advantages_trajs).float().to(self.device).reshape(-1)
)
logprobs_k = (
torch.tensor(logprobs_trajs).float().to(self.device).reshape(-1)
torch.tensor(advantages_trajs, device=self.device).float().reshape(-1)
)
logprobs_k = torch.tensor(logprobs_trajs, device=self.device).float()
# Update policy and critic
total_steps = self.n_steps * self.n_envs
inds_k = np.arange(total_steps)
clipfracs = []
for update_epoch in range(self.update_epochs):
# for each epoch, go through all data in batches
flag_break = False
np.random.shuffle(inds_k)
inds_k = torch.randperm(total_steps, device=self.device)
num_batch = max(1, total_steps // self.batch_size) # skip last ones
for batch in range(num_batch):
start = batch * self.batch_size
agent/finetune/train_ppo_gaussian_img_agent.py
+7 -8
View File
@@ -231,7 +231,7 @@ class TrainPPOImgGaussianAgent(TrainPPOGaussianAgent):
)
reward_trajs = reward_trajs_transpose.T
# bootstrap value with GAE if not done - apply reward scaling with constant if specified
# bootstrap value with GAE if not terminal - apply reward scaling with constant if specified
obs_venv_ts = {
key: torch.from_numpy(obs_venv[key]).float().to(self.device)
for key in self.obs_dims
@@ -271,29 +271,28 @@ class TrainPPOImgGaussianAgent(TrainPPOGaussianAgent):
for k in obs_trajs
}
samples_k = einops.rearrange(
torch.tensor(samples_trajs).float().to(self.device),
torch.tensor(samples_trajs, device=self.device).float(),
"s e h d -> (s e) h d",
)
returns_k = (
torch.tensor(returns_trajs).float().to(self.device).reshape(-1)
torch.tensor(returns_trajs, device=self.device).float().reshape(-1)
)
values_k = (
torch.tensor(values_trajs).float().to(self.device).reshape(-1)
torch.tensor(values_trajs, device=self.device).float().reshape(-1)
)
advantages_k = (
torch.tensor(advantages_trajs).float().to(self.device).reshape(-1)
torch.tensor(advantages_trajs, device=self.device).float().reshape(-1)
)
logprobs_k = torch.tensor(logprobs_trajs).float().to(self.device)
logprobs_k = torch.tensor(logprobs_trajs, device=self.device).float()
# Update policy and critic
total_steps = self.n_steps * self.n_envs
inds_k = np.arange(total_steps)
clipfracs = []
for update_epoch in range(self.update_epochs):
# for each epoch, go through all data in batches
flag_break = False
np.random.shuffle(inds_k)
inds_k = torch.randperm(total_steps, device=self.device)
num_batch = max(1, total_steps // self.batch_size) # skip last ones
for batch in range(num_batch):
start = batch * self.batch_size