diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py index 1b869fd2b52..252078eebeb 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py @@ -66,19 +66,6 @@ def __init__(self, cfg: AssemblyEnvCfg, render_mode: str | None = None, **kwargs if self.cfg_task.if_logging_eval: self._init_eval_logging() - if self.cfg_task.sample_from != "rand": - self._init_eval_loading() - - def _init_eval_loading(self): - eval_held_asset_pose, eval_fixed_asset_pose, eval_success = automate_log.load_log_from_hdf5( - self.cfg_task.eval_filename - ) - - if self.cfg_task.sample_from == "gp": - self.gp = automate_algo.model_succ_w_gp(eval_held_asset_pose, eval_fixed_asset_pose, eval_success) - elif self.cfg_task.sample_from == "gmm": - self.gmm = automate_algo.model_succ_w_gmm(eval_held_asset_pose, eval_fixed_asset_pose, eval_success) - def _init_eval_logging(self): self.held_asset_pose_log = torch.empty( @@ -246,7 +233,7 @@ def _load_disassembly_data(self): # offset each trajectory to be relative to the goal eef_pos_traj.append(curr_ee_traj - curr_ee_goal) - self.eef_pos_traj = torch.tensor(eef_pos_traj, dtype=torch.float32, device=self.device).squeeze() + self.eef_pos_traj = torch.tensor(np.array(eef_pos_traj), dtype=torch.float32, device=self.device).squeeze() def _get_keypoint_offsets(self, num_keypoints): """Get uniformly-spaced keypoints along a line of unit length, centered at 0.""" @@ -804,28 +791,12 @@ def randomize_held_initial_state(self, env_ids, pre_grasp): torch.rand((self.num_envs,), dtype=torch.float32, device=self.device) ) - if self.cfg_task.sample_from == "rand": - - rand_sample = torch.rand((len(env_ids), 3), dtype=torch.float32, device=self.device) - held_pos_init_rand = 2 * (rand_sample - 0.5) # [-1, 1] - held_asset_init_pos_rand = torch.tensor( - self.cfg_task.held_asset_init_pos_noise, dtype=torch.float32, device=self.device - ) - self.held_pos_init_rand = held_pos_init_rand @ torch.diag(held_asset_init_pos_rand) - - if self.cfg_task.sample_from == "gp": - rand_sample = torch.rand((self.cfg_task.num_gp_candidates, 3), dtype=torch.float32, device=self.device) - held_pos_init_rand = 2 * (rand_sample - 0.5) # [-1, 1] - held_asset_init_pos_rand = torch.tensor( - self.cfg_task.held_asset_init_pos_noise, dtype=torch.float32, device=self.device - ) - held_asset_init_candidates = held_pos_init_rand @ torch.diag(held_asset_init_pos_rand) - self.held_pos_init_rand, _ = automate_algo.propose_failure_samples_batch_from_gp( - self.gp, held_asset_init_candidates.cpu().detach().numpy(), len(env_ids), self.device - ) - - if self.cfg_task.sample_from == "gmm": - self.held_pos_init_rand = automate_algo.sample_rel_pos_from_gmm(self.gmm, len(env_ids), self.device) + rand_sample = torch.rand((len(env_ids), 3), dtype=torch.float32, device=self.device) + held_pos_init_rand = 2 * (rand_sample - 0.5) # [-1, 1] + held_asset_init_pos_rand = torch.tensor( + self.cfg_task.held_asset_init_pos_noise, dtype=torch.float32, device=self.device + ) + self.held_pos_init_rand = held_pos_init_rand @ torch.diag(held_asset_init_pos_rand) # Set plug pos to assembled state, but offset plug Z-coordinate by height of socket, # minus curriculum displacement diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_tasks_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_tasks_cfg.py index 729402ccc82..f2100216010 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_tasks_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_tasks_cfg.py @@ -139,10 +139,6 @@ class AssemblyTask: num_eval_trials: int = 100 eval_filename: str = "evaluation_00015.h5" - # Fine-tuning - sample_from: str = "rand" # gp, gmm, idv, rand - num_gp_candidates: int = 1000 - @configclass class Peg8mm(HeldAssetCfg): diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_tasks_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_tasks_cfg.py index fe292d31b4d..9308f281491 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_tasks_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_tasks_cfg.py @@ -118,7 +118,7 @@ class Extraction(DisassemblyTask): assembly_id = "00015" assembly_dir = f"{ASSET_DIR}/{assembly_id}/" disassembly_dir = "disassembly_dir" - num_log_traj = 1000 + num_log_traj = 100 fixed_asset_cfg = Hole8mm() held_asset_cfg = Peg8mm() diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/run_w_id.py b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/run_w_id.py index 01329d8ab70..9b5be5d914e 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/run_w_id.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/run_w_id.py @@ -61,9 +61,9 @@ def main(): bash_command = None if sys.platform.startswith("win"): - bash_command = "isaaclab.bat -p" + bash_command = "NUMBA_CUDA_LOW_OCCUPANCY_WARNINGS=0 isaaclab.bat -p" elif sys.platform.startswith("linux"): - bash_command = "./isaaclab.sh -p" + bash_command = "NUMBA_CUDA_LOW_OCCUPANCY_WARNINGS=0 ./isaaclab.sh -p" if args.train: bash_command += " scripts/reinforcement_learning/rl_games/train.py --task=Isaac-AutoMate-Assembly-Direct-v0" bash_command += f" --seed={str(args.seed)} --max_iterations={str(args.max_iterations)}" diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/soft_dtw_cuda.py b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/soft_dtw_cuda.py index f319a90e008..f010348cd4c 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/soft_dtw_cuda.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/soft_dtw_cuda.py @@ -123,8 +123,8 @@ class _SoftDTWCUDA(Function): def forward(ctx, D, gamma, bandwidth): dev = D.device dtype = D.dtype - gamma = torch.cuda.FloatTensor([gamma]) - bandwidth = torch.cuda.FloatTensor([bandwidth]) + gamma = torch.tensor([gamma], dtype=torch.float, device="cuda") + bandwidth = torch.tensor([bandwidth], dtype=torch.float, device="cuda") B = D.shape[0] N = D.shape[1]