Fix: PickSequentialTask spawning

mani_skill2/envs/scenes/tasks/pick.py

@@ -170,8 +170,8 @@ def reconfigure(self):
  qpos = torch.tensor(
  self.agent.RESTING_QPOS[..., None].repeat(self.num_envs, axis=-1).transpose(1, 0)
  ).float()
- accept_spawn_loc_rots = [[]] * self.num_envs
- accept_dists = [[]] * self.num_envs
+ accept_spawn_loc_rots = [[] for _ in range(self.num_envs)]
+ accept_dists = [[] for _ in range(self.num_envs)]
  bounding_box_corners = [
  torch.tensor([dx, dy, 0]) for dx, dy in itertools.product([0.1, -0.1], [0.1, -0.1])
  ]
@@ -206,13 +206,12 @@ def reconfigure(self):
 
  for i in torch.where(slrs_within_range & (robot_force < 1e-3))[0]:
  accept_spawn_loc_rots[i].append(slrs[i].cpu().numpy().tolist())
- accept_dists[i].append(dists[i].cpu().numpy().tolist())
-
+ accept_dists[i].append(dists[i].item())
 
  self.num_spawn_loc_rots = torch.tensor([len(x) for x in accept_spawn_loc_rots])
  self.spawn_loc_rots = pad_sequence([
  torch.tensor(x) for x in accept_spawn_loc_rots
- ], batch_first=True, padding_value=0,)
+ ], batch_first=True, padding_value=0)
 
  self.closest_spawn_loc_rots = torch.stack([
  self.spawn_loc_rots[i][torch.argmin(torch.tensor(x))] for i, x in enumerate(accept_dists)
@@ -253,10 +252,9 @@ def _initialize_agent(self, env_idx):
  idxs = torch.tensor([
  torch.randint(max_idx.item(), (1,)) for max_idx in self.num_spawn_loc_rots
  ])
- loc_rot = self.spawn_loc_rots[torch.arange(self.num_envs), idxs].to(self.device)
+ loc_rot = self.spawn_loc_rots[torch.arange(self.num_envs), idxs]
  else:
- loc_rot = self.closest_spawn_loc_rots.to(self.device)
-
+ loc_rot = self.closest_spawn_loc_rots
  robot_pos = self.agent.robot.pose.p
  robot_pos[..., :2] = loc_rot[..., :2]
  self.agent.robot.set_pose(Pose.create_from_pq(p=robot_pos))
@@ -265,19 +263,19 @@ def _initialize_agent(self, env_idx):
  if self.randomize_base:
  # base pos
  robot_pos = self.agent.robot.pose.p
- robot_pos[..., :2] += torch.clip(torch.normal(
+ robot_pos[..., :2] += torch.clamp(torch.normal(
  0, 0.1, (b, len(robot_pos[0, :2]))
  ), -0.1, 0.1).to(self.device)
  self.agent.robot.set_pose(Pose.create_from_pq(p=robot_pos))
  # base rot
- qpos[..., 2:3] += torch.clip(torch.normal(
+ qpos[..., 2:3] += torch.clamp(torch.normal(
  0, 0.25, (b, len(qpos[0, 2:3]))
  ), -0.5, 0.5).to(self.device)
  if self.randomize_arm:
- qpos[..., 5:6] += torch.clip(torch.normal(
+ qpos[..., 5:6] += torch.clamp(torch.normal(
  0, 0.05, (b, len(qpos[0, 5:6]))
  ), -0.1, 0.1).to(self.device)
- qpos[..., 7:-2] += torch.clip(torch.normal(
+ qpos[..., 7:-2] += torch.clamp(torch.normal(
  0, 0.05, (b, len(qpos[0, 7:-2]))
  ), -0.1, 0.1).to(self.device)
  self.agent.reset(qpos)