isaac-sim / IsaacGymEnvs

Isaac Gym Reinforcement Learning Environments
Other
1.77k stars 389 forks source link

Physical Dynamics Variation in Factory and IndustReal Environments #181

Closed patricknaughton01 closed 7 months ago

patricknaughton01 commented 8 months ago

Hello,

I've been using IsaacGym and in particular the IndustReal environments to learn peg insertion policies. I was hoping to extend these environments to simulate scenes that vary physical parameters (specifically the friction of the socket and plug, and the mass of the plug) between environment resets. Unfortunately, it seems that trying to adjust the friction of the socket (via set_actor_rigid_shape_properties​) in the reset_idx​ function causes the simulator to ignore collisions between the peg and socket. The problem seems to happen even if I simply read the sockets' properties and then set them back to the same values.

image (socket friction read and set back to same value).

Additionally, trying to read and reset the mass of the plug back to the same value seems to cause the simulator to miss a position update. Here's a video of a "normal" execution (without trying to write to the plug's mass):

https://github.com/NVIDIA-Omniverse/IsaacGymEnvs/assets/28968518/56f60ad5-0711-4b21-acf8-979d193e131f

And after trying to read and reset the plugs mass:

https://github.com/NVIDIA-Omniverse/IsaacGymEnvs/assets/28968518/5048611c-cf69-419f-a22d-7a87dc1c1eae

I was wondering if I'm maybe doing something wrong and I should be able to modify these parameters during reset_idx​ with the SDF geometries. I wasn't sure because I don't think the Factory or IndustReal project did this. Since the SDF representation seems to be relatively new, I was also wondering if there's something I could do on my end to integrate this functionality.

Thanks for your help!

bingjietang718 commented 8 months ago

Hi Patrick,

Can you try add self.simulate_and_refresh() every time right after you modify physics parameters? This function will essentially simulate one step and apply the changes you made to the physical environment.

If this does not work, can you share a snippet of your code in reset_idx() so I can replicate the issue locally?

Best, Bingjie

patricknaughton01 commented 8 months ago

Hi Bingjie,

Quick note is that I did try to call self.simulate_and_refresh after modifying the parameters, it fixed one of the issues but introduced a new one:

I did some more poking around to characterize the issue a bit more. For all of the below scenarios, I ran this command: python train.py task=IndustRealTaskPegsInsert task.rl.curriculum_height_bound=[0.005,0.01]

with this reset_idx (keeping gravity off so I can see the peg fall into the socket):

def reset_idx(self, env_ids):
      """Reset specified environments."""

      self._reset_franka()

      # Close gripper onto plug
      # self.disable_gravity()  # to prevent plug from falling
      self._reset_object()
      self._move_gripper_to_grasp_pose(
          sim_steps=self.cfg_task.env.num_gripper_move_sim_steps
      )
      self.close_gripper(sim_steps=self.cfg_task.env.num_gripper_close_sim_steps)
      self.enable_gravity()

      # Get plug SDF in goal pose for SDF-based reward
      self.plug_goal_sdfs = algo_utils.get_plug_goal_sdfs(
          wp_plug_meshes=self.wp_plug_meshes,
          asset_indices=self.asset_indices,
          socket_pos=self.socket_pos,
          socket_quat=self.socket_quat,
          wp_device=self.wp_device,
      )

      self._reset_buffers()

Socket modifications

I modified the _reset_socket to three different versions, adding the following lines in different places (which I think should just read, and then set the friction parameters back to the same values):

for i, ptr in enumerate(self.env_ptrs):
    shape_props = self.gym.get_actor_rigid_shape_properties(ptr, self.socket_handles[i])
    self.gym.set_actor_rigid_shape_properties(ptr, self.socket_handles[i], shape_props)

To reproduce the issue of the peg passing through the socket, I used this version:

    def _reset_socket(self):
        """Reset root state of socket."""

        # Randomize socket pos
        socket_noise_xy = 2 * (
            torch.rand((self.num_envs, 2), dtype=torch.float32, device=self.device)
            - 0.5
        )
        socket_noise_xy = socket_noise_xy @ torch.diag(
            torch.tensor(
                self.cfg_task.randomize.socket_pos_xy_noise,
                dtype=torch.float32,
                device=self.device,
            )
        )
        socket_noise_z = torch.zeros(
            (self.num_envs), dtype=torch.float32, device=self.device
        )
        socket_noise_z_mag = (
            self.cfg_task.randomize.socket_pos_z_noise_bounds[1]
            - self.cfg_task.randomize.socket_pos_z_noise_bounds[0]
        )
        socket_noise_z = (
            socket_noise_z_mag
            * torch.rand((self.num_envs), dtype=torch.float32, device=self.device)
            + self.cfg_task.randomize.socket_pos_z_noise_bounds[0]
        )

        self.socket_pos[:, 0] = (
            self.robot_base_pos[:, 0]
            + self.cfg_task.randomize.socket_pos_xy_initial[0]
            + socket_noise_xy[:, 0]
        )
        self.socket_pos[:, 1] = (
            self.robot_base_pos[:, 1]
            + self.cfg_task.randomize.socket_pos_xy_initial[1]
            + socket_noise_xy[:, 1]
        )
        self.socket_pos[:, 2] = self.cfg_base.env.table_height + socket_noise_z

        # Randomize socket rot
        socket_rot_noise = 2 * (
            torch.rand((self.num_envs, 3), dtype=torch.float32, device=self.device)
            - 0.5
        )
        socket_rot_noise = socket_rot_noise @ torch.diag(
            torch.tensor(
                self.cfg_task.randomize.socket_rot_noise,
                dtype=torch.float32,
                device=self.device,
            )
        )
        socket_rot_euler = (
            torch.zeros((self.num_envs, 3), dtype=torch.float32, device=self.device)
            + socket_rot_noise
        )
        socket_rot_quat = torch_utils.quat_from_euler_xyz(
            socket_rot_euler[:, 0], socket_rot_euler[:, 1], socket_rot_euler[:, 2]
        )
        self.socket_quat[:, :] = socket_rot_quat.clone()

        for i, ptr in enumerate(self.env_ptrs):
            shape_props = self.gym.get_actor_rigid_shape_properties(ptr, self.socket_handles[i])
            self.gym.set_actor_rigid_shape_properties(ptr, self.socket_handles[i], shape_props)

        # Stabilize socket
        self.socket_linvel[:, :] = 0.0
        self.socket_angvel[:, :] = 0.0

        # Set socket root state
        socket_actor_ids_sim = self.socket_actor_ids_sim.clone().to(dtype=torch.int32)
        self.gym.set_actor_root_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.root_state),
            gymtorch.unwrap_tensor(socket_actor_ids_sim),
            len(socket_actor_ids_sim),
        )

        # Simulate one step to apply changes
        self.simulate_and_refresh()

Adding self.simulate_and_refresh stops the peg from passing through the socket, but resets the socket back to some default pose (removing the randomization):

    def _reset_socket(self):
        """Reset root state of socket."""

        # Randomize socket pos
        socket_noise_xy = 2 * (
            torch.rand((self.num_envs, 2), dtype=torch.float32, device=self.device)
            - 0.5
        )
        socket_noise_xy = socket_noise_xy @ torch.diag(
            torch.tensor(
                self.cfg_task.randomize.socket_pos_xy_noise,
                dtype=torch.float32,
                device=self.device,
            )
        )
        socket_noise_z = torch.zeros(
            (self.num_envs), dtype=torch.float32, device=self.device
        )
        socket_noise_z_mag = (
            self.cfg_task.randomize.socket_pos_z_noise_bounds[1]
            - self.cfg_task.randomize.socket_pos_z_noise_bounds[0]
        )
        socket_noise_z = (
            socket_noise_z_mag
            * torch.rand((self.num_envs), dtype=torch.float32, device=self.device)
            + self.cfg_task.randomize.socket_pos_z_noise_bounds[0]
        )

        self.socket_pos[:, 0] = (
            self.robot_base_pos[:, 0]
            + self.cfg_task.randomize.socket_pos_xy_initial[0]
            + socket_noise_xy[:, 0]
        )
        self.socket_pos[:, 1] = (
            self.robot_base_pos[:, 1]
            + self.cfg_task.randomize.socket_pos_xy_initial[1]
            + socket_noise_xy[:, 1]
        )
        self.socket_pos[:, 2] = self.cfg_base.env.table_height + socket_noise_z

        # Randomize socket rot
        socket_rot_noise = 2 * (
            torch.rand((self.num_envs, 3), dtype=torch.float32, device=self.device)
            - 0.5
        )
        socket_rot_noise = socket_rot_noise @ torch.diag(
            torch.tensor(
                self.cfg_task.randomize.socket_rot_noise,
                dtype=torch.float32,
                device=self.device,
            )
        )
        socket_rot_euler = (
            torch.zeros((self.num_envs, 3), dtype=torch.float32, device=self.device)
            + socket_rot_noise
        )
        socket_rot_quat = torch_utils.quat_from_euler_xyz(
            socket_rot_euler[:, 0], socket_rot_euler[:, 1], socket_rot_euler[:, 2]
        )
        self.socket_quat[:, :] = socket_rot_quat.clone()

        for i, ptr in enumerate(self.env_ptrs):
            shape_props = self.gym.get_actor_rigid_shape_properties(ptr, self.socket_handles[i])
            self.gym.set_actor_rigid_shape_properties(ptr, self.socket_handles[i], shape_props)
        self.simulate_and_refresh()

        # Stabilize socket
        self.socket_linvel[:, :] = 0.0
        self.socket_angvel[:, :] = 0.0

        # Set socket root state
        socket_actor_ids_sim = self.socket_actor_ids_sim.clone().to(dtype=torch.int32)
        self.gym.set_actor_root_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.root_state),
            gymtorch.unwrap_tensor(socket_actor_ids_sim),
            len(socket_actor_ids_sim),
        )

        # Simulate one step to apply changes
        self.simulate_and_refresh()

Finally, to my eye, this version seems to work correctly (moving the friction randomization to the end):

    def _reset_socket(self):
        """Reset root state of socket."""

        # Randomize socket pos
        socket_noise_xy = 2 * (
            torch.rand((self.num_envs, 2), dtype=torch.float32, device=self.device)
            - 0.5
        )
        socket_noise_xy = socket_noise_xy @ torch.diag(
            torch.tensor(
                self.cfg_task.randomize.socket_pos_xy_noise,
                dtype=torch.float32,
                device=self.device,
            )
        )
        socket_noise_z = torch.zeros(
            (self.num_envs), dtype=torch.float32, device=self.device
        )
        socket_noise_z_mag = (
            self.cfg_task.randomize.socket_pos_z_noise_bounds[1]
            - self.cfg_task.randomize.socket_pos_z_noise_bounds[0]
        )
        socket_noise_z = (
            socket_noise_z_mag
            * torch.rand((self.num_envs), dtype=torch.float32, device=self.device)
            + self.cfg_task.randomize.socket_pos_z_noise_bounds[0]
        )

        self.socket_pos[:, 0] = (
            self.robot_base_pos[:, 0]
            + self.cfg_task.randomize.socket_pos_xy_initial[0]
            + socket_noise_xy[:, 0]
        )
        self.socket_pos[:, 1] = (
            self.robot_base_pos[:, 1]
            + self.cfg_task.randomize.socket_pos_xy_initial[1]
            + socket_noise_xy[:, 1]
        )
        self.socket_pos[:, 2] = self.cfg_base.env.table_height + socket_noise_z

        # Randomize socket rot
        socket_rot_noise = 2 * (
            torch.rand((self.num_envs, 3), dtype=torch.float32, device=self.device)
            - 0.5
        )
        socket_rot_noise = socket_rot_noise @ torch.diag(
            torch.tensor(
                self.cfg_task.randomize.socket_rot_noise,
                dtype=torch.float32,
                device=self.device,
            )
        )
        socket_rot_euler = (
            torch.zeros((self.num_envs, 3), dtype=torch.float32, device=self.device)
            + socket_rot_noise
        )
        socket_rot_quat = torch_utils.quat_from_euler_xyz(
            socket_rot_euler[:, 0], socket_rot_euler[:, 1], socket_rot_euler[:, 2]
        )
        self.socket_quat[:, :] = socket_rot_quat.clone()

        # Stabilize socket
        self.socket_linvel[:, :] = 0.0
        self.socket_angvel[:, :] = 0.0

        # Set socket root state
        socket_actor_ids_sim = self.socket_actor_ids_sim.clone().to(dtype=torch.int32)
        self.gym.set_actor_root_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.root_state),
            gymtorch.unwrap_tensor(socket_actor_ids_sim),
            len(socket_actor_ids_sim),
        )

        # Simulate one step to apply changes
        self.simulate_and_refresh()

        for i, ptr in enumerate(self.env_ptrs):
            shape_props = self.gym.get_actor_rigid_shape_properties(ptr, self.socket_handles[i])
            self.gym.set_actor_rigid_shape_properties(ptr, self.socket_handles[i], shape_props)

To my eye, the behavior of version 3 of this code seems to be correct (no resetting of the socket pose and the peg and socket clearly can collide), but I don't understand why changing the order like this makes a difference, which makes me a bit uneasy (is this introducing bugs that I'm not detecting that will crop up later down the line?).

Plug modifications

Additionally, if I try the same thing to modify the mass of the plug, it resets the position of the plug:

    def _reset_plug(self, before_move_to_grasp):
        """Reset root state of plug."""

        if before_move_to_grasp:
            # Generate randomized downward displacement based on curriculum
            curr_curriculum_disp_range = (
                self.curr_max_disp - self.cfg_task.rl.curriculum_height_bound[0]
            )
            self.curriculum_disp = self.cfg_task.rl.curriculum_height_bound[
                0
            ] + curr_curriculum_disp_range * (
                torch.rand((self.num_envs,), dtype=torch.float32, device=self.device)
            )

            # Generate plug pos noise
            self.plug_pos_xy_noise = 2 * (
                torch.rand((self.num_envs, 2), dtype=torch.float32, device=self.device)
                - 0.5
            )
            self.plug_pos_xy_noise = self.plug_pos_xy_noise @ torch.diag(
                torch.tensor(
                    self.cfg_task.randomize.plug_pos_xy_noise,
                    dtype=torch.float32,
                    device=self.device,
                )
            )

        # Set plug pos to assembled state, but offset plug Z-coordinate by height of socket,
        # minus curriculum displacement
        self.plug_pos[:, :] = self.socket_pos.clone()
        self.plug_pos[:, 2] += self.socket_heights
        self.plug_pos[:, 2] -= self.curriculum_disp

        # Apply XY noise to plugs not partially inserted into sockets
        socket_top_height = self.socket_pos[:, 2] + self.socket_heights
        plug_partial_insert_idx = np.argwhere(
            self.plug_pos[:, 2].cpu().numpy() > socket_top_height.cpu().numpy()
        ).squeeze()
        self.plug_pos[plug_partial_insert_idx, :2] += self.plug_pos_xy_noise[
            plug_partial_insert_idx
        ]

        self.plug_quat[:, :] = self.identity_quat.clone()

        # Stabilize plug
        self.plug_linvel[:, :] = 0.0
        self.plug_angvel[:, :] = 0.0

        # Set plug root state
        plug_actor_ids_sim = self.plug_actor_ids_sim.clone().to(dtype=torch.int32)
        self.gym.set_actor_root_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.root_state),
            gymtorch.unwrap_tensor(plug_actor_ids_sim),
            len(plug_actor_ids_sim),
        )

        # Simulate one step to apply changes
        self.simulate_and_refresh()

        for i, ptr in enumerate(self.env_ptrs):
            body_props = self.gym.get_actor_rigid_body_properties(ptr, self.plug_handles[i])
            self.gym.set_actor_rigid_body_properties(ptr, self.plug_handles[i], body_props)

And if I move it to the beginning of the _reset_plug function, the plug's pose is never updated:

    def _reset_plug(self, before_move_to_grasp):
        """Reset root state of plug."""

        if before_move_to_grasp:
            for i, ptr in enumerate(self.env_ptrs):
                body_props = self.gym.get_actor_rigid_body_properties(ptr, self.plug_handles[i])
                self.gym.set_actor_rigid_body_properties(ptr, self.plug_handles[i], body_props)
            # Generate randomized downward displacement based on curriculum
            curr_curriculum_disp_range = (
                self.curr_max_disp - self.cfg_task.rl.curriculum_height_bound[0]
            )
            self.curriculum_disp = self.cfg_task.rl.curriculum_height_bound[
                0
            ] + curr_curriculum_disp_range * (
                torch.rand((self.num_envs,), dtype=torch.float32, device=self.device)
            )

            # Generate plug pos noise
            self.plug_pos_xy_noise = 2 * (
                torch.rand((self.num_envs, 2), dtype=torch.float32, device=self.device)
                - 0.5
            )
            self.plug_pos_xy_noise = self.plug_pos_xy_noise @ torch.diag(
                torch.tensor(
                    self.cfg_task.randomize.plug_pos_xy_noise,
                    dtype=torch.float32,
                    device=self.device,
                )
            )

        # Set plug pos to assembled state, but offset plug Z-coordinate by height of socket,
        # minus curriculum displacement
        self.plug_pos[:, :] = self.socket_pos.clone()
        self.plug_pos[:, 2] += self.socket_heights
        self.plug_pos[:, 2] -= self.curriculum_disp

        # Apply XY noise to plugs not partially inserted into sockets
        socket_top_height = self.socket_pos[:, 2] + self.socket_heights
        plug_partial_insert_idx = np.argwhere(
            self.plug_pos[:, 2].cpu().numpy() > socket_top_height.cpu().numpy()
        ).squeeze()
        self.plug_pos[plug_partial_insert_idx, :2] += self.plug_pos_xy_noise[
            plug_partial_insert_idx
        ]

        self.plug_quat[:, :] = self.identity_quat.clone()

        # Stabilize plug
        self.plug_linvel[:, :] = 0.0
        self.plug_angvel[:, :] = 0.0

        # Set plug root state
        plug_actor_ids_sim = self.plug_actor_ids_sim.clone().to(dtype=torch.int32)
        self.gym.set_actor_root_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.root_state),
            gymtorch.unwrap_tensor(plug_actor_ids_sim),
            len(plug_actor_ids_sim),
        )

        # Simulate one step to apply changes
        self.simulate_and_refresh()

Remaining questions

Thanks for your help!

Patrick

bingjietang718 commented 8 months ago

Hi Patrick,

In general, gym setter functions (e.g., self.gym.set_actor_root_state_tensor_indexed and self.gym.set_actor_rigid_body_properties) should only be called once in each simulation step (see this post for reference). So what I would suggest is:

Also, maybe worth checking out this paper, they have tried physics parameter randomization and their code released as well.

Best, Bingjie

patricknaughton01 commented 8 months ago

Hi Bingjie,

I see, unfortunately, the simple fix of pushing self.simulate_and_refresh() into the for loop seems to cause some position updates to the plug to be ignored. Thanks for the pointer to the code, I'll check that out.

Patrick