From 9da961a2bf6d8fb8bcddd267d94f9eba9e542671 Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Wed, 10 Apr 2024 14:53:05 +0200 Subject: [PATCH 01/15] remove robomimic_eef policy --- imitation/config/policy/robomimic_eef.yaml | 12 ------------ 1 file changed, 12 deletions(-) delete mode 100644 imitation/config/policy/robomimic_eef.yaml diff --git a/imitation/config/policy/robomimic_eef.yaml b/imitation/config/policy/robomimic_eef.yaml deleted file mode 100644 index ef331d3..0000000 --- a/imitation/config/policy/robomimic_eef.yaml +++ /dev/null @@ -1,12 +0,0 @@ -_target_: imitation.policy.robomimic_lowdim_policy.RobomimicPretrainedWrapper -obs_dim: ${task.obs_dim} -action_dim: ${task.action_dim} -# only works with fixed horizons, pred_horizon = obs_horizon -algo_name: "bc_rnn" -obs_type: "low_dim" -task_name: ${task.task_name} -dataset_type: ${task.dataset_type} - -dataset: ${task.dataset} -ckpt_path: ./weights/lift_ph_low_dim_epoch_1000_succ_100.pth -lr: 0.0001 \ No newline at end of file From 8a2cf404d8e44d55170d736878669d5447ac4115 Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Tue, 2 Jan 2024 15:46:29 +0100 Subject: [PATCH 02/15] use EMA model --- conda_environment.yaml | 4 ++++ imitation/policy/diffusion_policy.py | 23 +++++++++++------------ 2 files changed, 15 insertions(+), 12 deletions(-) diff --git a/conda_environment.yaml b/conda_environment.yaml index 36eb619..fcf40dd 100644 --- a/conda_environment.yaml +++ b/conda_environment.yaml @@ -17,6 +17,10 @@ dependencies: - -e git+https://github.com/ARISE-Initiative/robomimic@main#egg=robomimic - diffusers - zarr + - h5py + - robomimic + - diffusers + - zarr - einops - tqdm - pybullet diff --git a/imitation/policy/diffusion_policy.py b/imitation/policy/diffusion_policy.py index 9ca75f5..7bdd3a2 100644 --- a/imitation/policy/diffusion_policy.py +++ b/imitation/policy/diffusion_policy.py @@ -7,6 +7,7 @@ from tqdm.auto import tqdm from diffusers.optimization import get_scheduler +from diffusers.training_utils import EMAModel from diffusers.schedulers.scheduling_ddpm import DDPMScheduler import wandb @@ -160,7 +161,7 @@ def get_action(self, obs_seq): # only take action_horizon number of actions action = action_pred[:self.action_horizon,:] # (action_horizon, action_dim) - return action # TODO limit this in runner + return action def validate(self, dataset=None, model_path=None): ''' @@ -189,10 +190,9 @@ def train(self, # accelerates training and improves stability # holds a copy of the model weights - # TODO use EMA - # ema = EMAModel( - # model=noise_pred_net, - # power=0.75) + ema = EMAModel( + parameters=self.noise_pred_net.parameters(), + power=0.75) # Standard ADAM optimizer # Note that EMA parameters are not optimized @@ -263,9 +263,8 @@ def train(self, # this is different from standard pytorch behavior lr_scheduler.step() - # TODO use EMA # update Exponential Moving Average of the model weights - # ema.step(noise_pred_net) + ema.step(self.noise_pred_net.parameters()) # logging @@ -274,10 +273,10 @@ def train(self, tepoch.set_postfix(loss=loss_cpu) tglobal.set_postfix(loss=np.mean(epoch_loss)) wandb.log({'epoch_loss': np.mean(epoch_loss)}) + # Weights of the EMA model are used for inference + ema_noise_pred_net = self.noise_pred_net + ema.copy_to(ema_noise_pred_net.parameters()) # save model checkpoint - torch.save(self.noise_pred_net.state_dict(), model_path) + torch.save(ema_noise_pred_net.state_dict(), model_path) - # Weights of the EMA model - # is used for inference - # ema_noise_pred_net = ema.averaged_model - self.ema_noise_pred_net = self.noise_pred_net \ No newline at end of file + \ No newline at end of file From f6842e6dc7b7f3be7f654ee1f7be514bd433577d Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Wed, 10 Apr 2024 15:58:28 +0200 Subject: [PATCH 03/15] fix lowdim_wrapper --- imitation/env/robomimic_lowdim_wrapper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/imitation/env/robomimic_lowdim_wrapper.py b/imitation/env/robomimic_lowdim_wrapper.py index 370f121..6a100cc 100644 --- a/imitation/env/robomimic_lowdim_wrapper.py +++ b/imitation/env/robomimic_lowdim_wrapper.py @@ -105,7 +105,7 @@ def _robosuite_obs_to_robomimic_obs(self, obs): def reset(self): - obs, _ = self.env.reset() + obs = self.env.reset() return self._robosuite_obs_to_robomimic_obs(obs) From 1132be00d8427b873ca1584c4ee7a7d756ea2891 Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Wed, 10 Apr 2024 15:58:44 +0200 Subject: [PATCH 04/15] add default steps to lowdim runner --- imitation/env_runner/robomimic_lowdim_runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/imitation/env_runner/robomimic_lowdim_runner.py b/imitation/env_runner/robomimic_lowdim_runner.py index e192ba9..3dec28c 100644 --- a/imitation/env_runner/robomimic_lowdim_runner.py +++ b/imitation/env_runner/robomimic_lowdim_runner.py @@ -57,7 +57,7 @@ def reset(self) -> None: self.obs_deque = collections.deque( [self.obs] * self.obs_horizon, maxlen=self.obs_horizon) - def run(self, agent: BaseAgent, n_steps: int) -> Dict: + def run(self, agent: BaseAgent, n_steps: int = 1000) -> Dict: log.info(f"Running agent {agent.__class__.__name__} for {n_steps} steps") if self.output_video: self.start_video() From fd77f11c8be5ac934eb120b3d4c2b2831b2efefa Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Wed, 10 Apr 2024 17:46:48 +0200 Subject: [PATCH 05/15] fix lowdim_wrapper --- imitation/env/robomimic_lowdim_wrapper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/imitation/env/robomimic_lowdim_wrapper.py b/imitation/env/robomimic_lowdim_wrapper.py index 6a100cc..370f121 100644 --- a/imitation/env/robomimic_lowdim_wrapper.py +++ b/imitation/env/robomimic_lowdim_wrapper.py @@ -105,7 +105,7 @@ def _robosuite_obs_to_robomimic_obs(self, obs): def reset(self): - obs = self.env.reset() + obs, _ = self.env.reset() return self._robosuite_obs_to_robomimic_obs(obs) From ea4ee7387b865ea98f9f6d904b2da7461d7ff1c1 Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Wed, 10 Apr 2024 18:26:09 +0200 Subject: [PATCH 06/15] add empty validate policy to diffusion_policy --- imitation/env_runner/robomimic_lowdim_runner.py | 2 +- imitation/policy/diffusion_policy.py | 6 ------ 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/imitation/env_runner/robomimic_lowdim_runner.py b/imitation/env_runner/robomimic_lowdim_runner.py index 3dec28c..f72a2cd 100644 --- a/imitation/env_runner/robomimic_lowdim_runner.py +++ b/imitation/env_runner/robomimic_lowdim_runner.py @@ -57,7 +57,7 @@ def reset(self) -> None: self.obs_deque = collections.deque( [self.obs] * self.obs_horizon, maxlen=self.obs_horizon) - def run(self, agent: BaseAgent, n_steps: int = 1000) -> Dict: + def run(self, agent: BaseAgent, n_steps: int = 500) -> Dict: log.info(f"Running agent {agent.__class__.__name__} for {n_steps} steps") if self.output_video: self.start_video() diff --git a/imitation/policy/diffusion_policy.py b/imitation/policy/diffusion_policy.py index 7bdd3a2..679e3d4 100644 --- a/imitation/policy/diffusion_policy.py +++ b/imitation/policy/diffusion_policy.py @@ -163,12 +163,6 @@ def get_action(self, obs_seq): # (action_horizon, action_dim) return action - def validate(self, dataset=None, model_path=None): - ''' - Calculate validation loss for noise prediction model in the given dataset - ''' - return None - def train(self, dataset=None, num_epochs=100, From 1914c2561979e16fc5c6728116073f0e2afd5fe4 Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Thu, 11 Apr 2024 19:51:30 +0200 Subject: [PATCH 07/15] step environment inside try, to avoid stepping after termination --- imitation/env_runner/robomimic_lowdim_runner.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/imitation/env_runner/robomimic_lowdim_runner.py b/imitation/env_runner/robomimic_lowdim_runner.py index f72a2cd..4d77387 100644 --- a/imitation/env_runner/robomimic_lowdim_runner.py +++ b/imitation/env_runner/robomimic_lowdim_runner.py @@ -80,12 +80,14 @@ def run(self, agent: BaseAgent, n_steps: int = 500) -> Dict: if self.output_video: self.end_video() return rewards, info - obs, reward, done, info = self.env.step(action) + try: + obs, reward, done, info = self.env.step(action) + except Exception as e: + print(e) self.obs_deque.append(obs) if self.render: self.env.render() - # time.sleep(1/self.fps) # TODO properly fix the rendering speed or not if self.output_video: # We need to directly grab full observations so we can get image data From 9648fb788bd4072fab371d80502c2437c6f90a7d Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Thu, 11 Apr 2024 20:20:56 +0200 Subject: [PATCH 08/15] change reset for inside runner.run() method --- eval.py | 1 - imitation/env_runner/robomimic_lowdim_runner.py | 9 ++++----- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/eval.py b/eval.py index 5efb01c..1274bfb 100644 --- a/eval.py +++ b/eval.py @@ -50,7 +50,6 @@ def eval_main(cfg): # run policy in environment success_count = 0 for i in range(cfg.num_episodes): - runner.reset() rewards, info = runner.run(agent, cfg.max_steps) assert "success" in info, "info['success'] not returned in info from runner" print(f"info: {info}") diff --git a/imitation/env_runner/robomimic_lowdim_runner.py b/imitation/env_runner/robomimic_lowdim_runner.py index 4d77387..16120f6 100644 --- a/imitation/env_runner/robomimic_lowdim_runner.py +++ b/imitation/env_runner/robomimic_lowdim_runner.py @@ -57,8 +57,9 @@ def reset(self) -> None: self.obs_deque = collections.deque( [self.obs] * self.obs_horizon, maxlen=self.obs_horizon) - def run(self, agent: BaseAgent, n_steps: int = 500) -> Dict: + def run(self, agent: BaseAgent, n_steps: int = 100) -> Dict: log.info(f"Running agent {agent.__class__.__name__} for {n_steps} steps") + self.reset() if self.output_video: self.start_video() done = False @@ -80,10 +81,8 @@ def run(self, agent: BaseAgent, n_steps: int = 500) -> Dict: if self.output_video: self.end_video() return rewards, info - try: - obs, reward, done, info = self.env.step(action) - except Exception as e: - print(e) + + obs, reward, done, info = self.env.step(action) self.obs_deque.append(obs) if self.render: From 82b4d60c5246654960f96f81c5aab1e7e1e40d7c Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Thu, 11 Apr 2024 20:39:20 +0200 Subject: [PATCH 09/15] separate files for eef and lowdim datasets --- imitation/dataset/robomimic_eef_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/imitation/dataset/robomimic_eef_dataset.py b/imitation/dataset/robomimic_eef_dataset.py index 2fd5aa7..0ea2330 100644 --- a/imitation/dataset/robomimic_eef_dataset.py +++ b/imitation/dataset/robomimic_eef_dataset.py @@ -37,7 +37,7 @@ def __init__(self, self.data_at_indices = [] # if indices file exists, load it index_file = dataset_path.replace(".hdf5", f"_indices_{obs_horizon}_{action_horizon}_{pred_horizon}.npy") - data_at_indices_file = dataset_path.replace(".hdf5", f"_data_at_indices_{obs_horizon}_{action_horizon}_{pred_horizon}.npy") + data_at_indices_file = dataset_path.replace(".hdf5", f"eef__data_at_indices_{obs_horizon}_{action_horizon}_{pred_horizon}.npy") if os.path.exists(index_file): self.indices = np.load(index_file) self.data_at_indices = np.load(data_at_indices_file, allow_pickle=True) From ae9e388a9217d75aa6cce9350e9111205c4d9723 Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Thu, 11 Apr 2024 20:40:08 +0200 Subject: [PATCH 10/15] change both index_file and data at indices file --- imitation/dataset/robomimic_eef_dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/imitation/dataset/robomimic_eef_dataset.py b/imitation/dataset/robomimic_eef_dataset.py index 0ea2330..d07af12 100644 --- a/imitation/dataset/robomimic_eef_dataset.py +++ b/imitation/dataset/robomimic_eef_dataset.py @@ -36,8 +36,8 @@ def __init__(self, self.indices = [] self.data_at_indices = [] # if indices file exists, load it - index_file = dataset_path.replace(".hdf5", f"_indices_{obs_horizon}_{action_horizon}_{pred_horizon}.npy") - data_at_indices_file = dataset_path.replace(".hdf5", f"eef__data_at_indices_{obs_horizon}_{action_horizon}_{pred_horizon}.npy") + index_file = dataset_path.replace(".hdf5", f"_eef_indices_{obs_horizon}_{action_horizon}_{pred_horizon}.npy") + data_at_indices_file = dataset_path.replace(".hdf5", f"_eef_data_at_indices_{obs_horizon}_{action_horizon}_{pred_horizon}.npy") if os.path.exists(index_file): self.indices = np.load(index_file) self.data_at_indices = np.load(data_at_indices_file, allow_pickle=True) From 58966ced21edf4b834a9314385fcd26b53b88000 Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Fri, 12 Apr 2024 12:33:17 +0200 Subject: [PATCH 11/15] add n_latency_steps = -1 to lowdim_dataset (joint_vel) --- imitation/dataset/robomimic_lowdim_dataset.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/imitation/dataset/robomimic_lowdim_dataset.py b/imitation/dataset/robomimic_lowdim_dataset.py index b41b4c7..e30b9c8 100644 --- a/imitation/dataset/robomimic_lowdim_dataset.py +++ b/imitation/dataset/robomimic_lowdim_dataset.py @@ -85,9 +85,10 @@ def create_sample_indices(self): |------------ pred_horizon -------------| ''' idx_global = 0 + n_latency_steps = -1 for key in tqdm(self.dataset_keys): episode_length = len(self.dataset_root[f"data/{key}/obs/{self.obs_keys[0]}"]) - for idx in range(episode_length - self.pred_horizon): + for idx in range(episode_length - self.pred_horizon + n_latency_steps): if idx - self.obs_horizon < 0: continue self.indices.append(idx_global + idx) @@ -99,7 +100,7 @@ def create_sample_indices(self): data_obs_keys.append(obs) data_action_keys = [] for action_key in self.action_keys: - action = self.dataset_root[f"data/{key}/obs/{action_key}"][idx:idx+self.pred_horizon, :] + action = self.dataset_root[f"data/{key}/obs/{action_key}"][idx + n_latency_steps:idx + n_latency_steps + self.pred_horizon, :] if "quat" in action_key: action = self.rotation_transformer.forward(action) data_action_keys.append(action) From 375c11a9d02947c40b6e054e6a0188379a779d4d Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Sun, 14 Apr 2024 19:15:02 +0200 Subject: [PATCH 12/15] remove joint positions from observations in lowdim --- imitation/env/robomimic_lowdim_wrapper.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/imitation/env/robomimic_lowdim_wrapper.py b/imitation/env/robomimic_lowdim_wrapper.py index 370f121..bd2b238 100644 --- a/imitation/env/robomimic_lowdim_wrapper.py +++ b/imitation/env/robomimic_lowdim_wrapper.py @@ -85,8 +85,8 @@ def _robosuite_obs_to_robomimic_obs(self, obs): for i in range(len(self.robots)): j = i*39 # 7 - sin of joint angles - robot_joint_pos = obs[j:j + 7] - # 7 - sin of joint angles + # robot_joint_pos = obs[j:j + 7] + # 7 - sin of joint angles # robot_joint_sin = obs[j + 7:j + 14] # 7 - cos of joint angles # robot_joint_cos = obs[j + 14:j + 21] @@ -97,7 +97,7 @@ def _robosuite_obs_to_robomimic_obs(self, obs): eef_6d = self.rotation_transformer.forward(eef_quat) gripper_pose = obs[j + 35:j + 37] # Skip 2 - gripper joint velocities - robot_i = [*robot_joint_pos, *robot_joint_vel, *eef_pose, *eef_6d, *gripper_pose] + robot_i = [*robot_joint_vel, *eef_pose, *eef_6d, *gripper_pose] final_obs = [*final_obs, *robot_i] objects = obs[39*len(self.robots):] From e070bbabe6ac2a28b2bc39d1c78fb88238618723 Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Sun, 14 Apr 2024 22:29:44 +0200 Subject: [PATCH 13/15] replace joint velocities for positions in observations --- imitation/env/robomimic_lowdim_wrapper.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/imitation/env/robomimic_lowdim_wrapper.py b/imitation/env/robomimic_lowdim_wrapper.py index bd2b238..14c0cf1 100644 --- a/imitation/env/robomimic_lowdim_wrapper.py +++ b/imitation/env/robomimic_lowdim_wrapper.py @@ -85,19 +85,19 @@ def _robosuite_obs_to_robomimic_obs(self, obs): for i in range(len(self.robots)): j = i*39 # 7 - sin of joint angles - # robot_joint_pos = obs[j:j + 7] + robot_joint_pos = obs[j:j + 7] # 7 - sin of joint angles # robot_joint_sin = obs[j + 7:j + 14] # 7 - cos of joint angles # robot_joint_cos = obs[j + 14:j + 21] # 7 - joint velocities - robot_joint_vel = obs[j + 21:j + 28] + # robot_joint_vel = obs[j + 21:j + 28] eef_pose = obs[j + 28:j + 31] eef_quat = obs[j + 31:j + 35] eef_6d = self.rotation_transformer.forward(eef_quat) gripper_pose = obs[j + 35:j + 37] # Skip 2 - gripper joint velocities - robot_i = [*robot_joint_vel, *eef_pose, *eef_6d, *gripper_pose] + robot_i = [*robot_joint_pos, *eef_pose, *eef_6d, *gripper_pose] final_obs = [*final_obs, *robot_i] objects = obs[39*len(self.robots):] From a903e394f0d63b3652f92519de30e653c215cc24 Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Wed, 17 Apr 2024 11:48:04 +0200 Subject: [PATCH 14/15] change n_latency_steps to 0 --- imitation/dataset/robomimic_lowdim_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/imitation/dataset/robomimic_lowdim_dataset.py b/imitation/dataset/robomimic_lowdim_dataset.py index e30b9c8..395e0c5 100644 --- a/imitation/dataset/robomimic_lowdim_dataset.py +++ b/imitation/dataset/robomimic_lowdim_dataset.py @@ -85,7 +85,7 @@ def create_sample_indices(self): |------------ pred_horizon -------------| ''' idx_global = 0 - n_latency_steps = -1 + n_latency_steps = 0 for key in tqdm(self.dataset_keys): episode_length = len(self.dataset_root[f"data/{key}/obs/{self.obs_keys[0]}"]) for idx in range(episode_length - self.pred_horizon + n_latency_steps): From 1f1b9d079f37ef089668355202a34cc38c5911d7 Mon Sep 17 00:00:00 2001 From: Caio Freitas Date: Sun, 28 Apr 2024 16:15:50 +0200 Subject: [PATCH 15/15] remove conflict by merge --- imitation/env/robomimic_lowdim_wrapper.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/imitation/env/robomimic_lowdim_wrapper.py b/imitation/env/robomimic_lowdim_wrapper.py index 7c62cfb..14c0cf1 100644 --- a/imitation/env/robomimic_lowdim_wrapper.py +++ b/imitation/env/robomimic_lowdim_wrapper.py @@ -91,10 +91,7 @@ def _robosuite_obs_to_robomimic_obs(self, obs): # 7 - cos of joint angles # robot_joint_cos = obs[j + 14:j + 21] # 7 - joint velocities -<<<<<<< HEAD # robot_joint_vel = obs[j + 21:j + 28] -======= ->>>>>>> cec97d0e06b1ae33c4bd568c2fb3e93a6d9f916d eef_pose = obs[j + 28:j + 31] eef_quat = obs[j + 31:j + 35] eef_6d = self.rotation_transformer.forward(eef_quat)