Enabling fluid simulation

Since Isaac Lab's support for fluid simulation in reinforcement learning (RL) is currently limited, we need to make some modifications to enable fluid simulation during RL training. To the best of our knowledge (by 11/2024), no one has yet proposed a complete solution for conducting fluid simulation in Isaac Lab. The implementation steps are as follows:

• Run simulation with cpu mode using --device cpu. After extensive testing, we found that fluid motion is visible in the GUI only when Isaac Sim is running in CPU mode, regardless of whether the fabric extension is enabled. Additionally, it is possible to retrieve the positions and velocities of fluid particles using Python code in this configuration.
• When running RL in CPU mode, some parts of the rsl_rl code may not function correctly. To address this, you need to add the following line to line 114 of the on_policy_runner.py file located in rsl_rl\rsl_rl\runners:
  actions = self.alg.act(obs, critic_obs)
  obs, rewards, dones, infos = self.env.step(actions)
  
  # Add this line.
  obs = obs.to(self.device)
  obs = self.obs_normalizer(obs)
  

• When running Isaac Lab in headless mode, the simulator settings differ from those in GUI mode, which prevents us from retrieving the positions of fluid particles. To resolve this, modify lines 213-214 in the file IsaacLab\source\apps\isaaclab.python.headless.kit as follows:
  updateToUsd = true
  updateParticlesToUsd = true
  

• The Manager-based workflow does not support fluid creation and reset, so we use the Direct workflow instead. Examples of the Direct workflow can be found in the directory:
  IsaacLab\source\extensions\omni.isaac.lab_tasks\omni\isaac\lab_tasks\direct
  

The Code

This is a sample code of RL training with fluid simulation:

The Code Explained

This section covers the implementation of fluid simulation in Isaac Lab.

In Isaac Lab, when creating a scene, we typically only need to create one env, and the other envs are duplicated by Isaac Lab based on the first one. However, since Isaac Lab does not support duplicating particle objects, we need to create individual particle objects for each env.

        # Create fluid in each envs.
        for i, origin in enumerate(self.scene.env_origins):
            self._create_fluid(i, self.stage, Sdf.Path(f"/World"), (origin + torch.tensor([-0.02, -0.03, 0.15]).to(self.device)).tolist(), [4, 4, 8])
            # Record the initial state (positions, velocities) of each particle.
            self.particle_init_states.append(self._get_particles_state(i))

We use particleUtils.add_physx_particle_system to create the particle system and the CreateMdlMaterialPrim command to bind the physical and surface materials of the particles.

            particle_system = particleUtils.add_physx_particle_system(
                stage=stage,
                particle_system_path=particleSystemPath,
                simulation_owner="physicsScene",
                contact_offset=restOffset * 1.5 + 0.005,
                rest_offset=restOffset * 1.5,
                particle_contact_offset=particleContactOffset,
                solid_rest_offset=0.0,
                fluid_rest_offset=fluidRestOffset,
                solver_position_iterations=16,
            )

            omni.kit.commands.execute(
                "CreateMdlMaterialPrim",
                mtl_url="OmniSurfacePresets.mdl",
                mtl_name="OmniSurface_DeepWater",
                mtl_path='/World/Looks/OmniSurface_DeepWater',
                select_new_prim=False,
            )
            pbd_particle_material_path = '/World/Looks/OmniSurface_DeepWater'
            omni.kit.commands.execute(
                "BindMaterial", prim_path=particleSystemPath, material_path=pbd_particle_material_path
            )

Next, we configure the physical properties of the particle system.

            # Create a pbd particle material and set it on the particle system
            particleUtils.add_pbd_particle_material(
                stage,
                pbd_particle_material_path,
                cohesion=10,
                viscosity=0.91,
                surface_tension=0.74,
                friction=0.1,
            )
            physicsUtils.add_physics_material_to_prim(stage, particle_system.GetPrim(), pbd_particle_material_path)

            particle_system.CreateMaxVelocityAttr().Set(200)

We use particleUtils.add_physx_particleset_points to create a particle set, which belongs to the particle system and contains the actual particles. In an RL environment, only one particle system is needed, but a separate particle set must be created for each environment.

        particleUtils.add_physx_particleset_points(
            stage=stage,
            path=particlesPath,
            positions_list=Vt.Vec3fArray(positions),
            velocities_list=Vt.Vec3fArray(velocities),
            widths_list=widths,
            particle_system_path=particleSystemPath,
            self_collision=True,
            fluid=True,
            particle_group=0,
            particle_mass=1,
            density=1000,
        )

Currently, there doesn't seem to be a direct way to retrieve the positions and velocities of particles from Isaac Lab's backend. As a workaround, we use the UsdGeom Points API to obtain the particle states. This is also why it's necessary to enable UpdateToUsd and UpdateParticlesToUsd.

    def _get_particles_state(self, env_id)->tuple[Gf.Vec3f, Gf.Vec3f]:
        # Gets particles' positions and velocities 
        particles = UsdGeom.Points(self.stage.GetPrimAtPath(Sdf.Path(f"/World/particles{env_id}")))
        particles_pos = particles.GetPointsAttr().Get()
        particles_vel = particles.GetVelocitiesAttr().Get()

        return particles_pos, particles_vel

Directly converting particle states into tensors is very slow, significantly impacting RL training efficiency. However, converting them into NumPy arrays is much faster. Therefore, we convert the particle positions into a NumPy array and record the minimum value along the z-axis for each group of particles, as this is the only data we need.

        # data for fluid
        self.particle_min = []
        for idx in range(self.num_envs):
            poses = np.array(self._get_particles_position(idx))
            self.particle_min.append(np.min(poses[:, 2]))
        self.particle_min = torch.tensor(self.particle_min, device=self.device)