Soft body patch (#12)

* fix Excavate max_particles too small

* fix recovery from a softbody sim crash

* fix Pour gripper initial position

* [BC] Fix camera for envs using excavation bucket

mani_skill2/agents/configs/panda/variants.py

@@ -59,8 +59,8 @@ def cameras(self):
         return dict(
             hand_camera=MountedCameraConfig(
                 mount_link="bucket",
-                mount_p=[0.0, 0.0, 0.0],
-                mount_q=[0.0, 0, 0.0, 0],
+                mount_p=[0.0, 0.08, 0.0],
+                mount_q=[0.5, -0.5, -0.5, -0.5],
                 hide_mount_link=False,
                 width=128,
                 height=128,

mani_skill2/envs/mpm/base_env.py

@@ -82,6 +82,8 @@ class ARGS:
         sapien.VulkanRenderer.set_camera_shader_dir(shader_dir)
         sapien.VulkanRenderer.set_viewer_shader_dir(shader_dir)
 
+        self.sim_crashed = False
+
         self._mpm_step_per_sapien_step = mpm_freq // sim_freq
         self._mpm_dt = 1 / mpm_freq
         self.max_particles = max_particles
@@ -338,7 +340,7 @@ def _initialize_mpm(self):
             color=(0.65237011, 0.14198029, 0.02201299),
             random_state=self._episode_rng,
         )
-        self.model_builder.init_model_state(self.mpm_model, self.mpm_states[0])
+        self.model_builder.init_model_state(self.mpm_model, self.mpm_states)
         self.mpm_model.struct.static_ke = 100.0
         self.mpm_model.struct.static_kd = 0.0
         self.mpm_model.struct.static_mu = 1.0
@@ -387,7 +389,11 @@ def get_obs(self):
         mpm_state = self.get_mpm_state()
 
         if self._obs_mode == "particles":
-            obs = OrderedDict(particles=mpm_state, agent=self._get_obs_agent())
+            obs = OrderedDict(
+                particles=mpm_state,
+                agent=self._get_obs_agent(),
+                extra=self._get_obs_extra(),
+            )
         else:
             obs = super().get_obs()
 

mani_skill2/envs/mpm/excavate_env.py

@@ -28,14 +28,12 @@ def __init__(
         *args,
         sim_freq=500,
         mpm_freq=2000,
-        max_particles=15000,
         **kwargs,
     ):
         super().__init__(
             *args,
             sim_freq=sim_freq,
             mpm_freq=mpm_freq,
-            max_particles=max_particles,
             **kwargs,
         )
 
@@ -78,7 +76,7 @@ def _initialize_mpm(self):
             random_state=self._episode_rng,
         )
 
-        self.model_builder.init_model_state(self.mpm_model, self.mpm_states[0])
+        self.model_builder.init_model_state(self.mpm_model, self.mpm_states)
         self.mpm_model.struct.static_ke = 100.0
         self.mpm_model.struct.static_kd = 0.0
         self.mpm_model.struct.static_mu = 1.0

mani_skill2/envs/mpm/fill_env.py

@@ -92,7 +92,7 @@ def _initialize_mpm(self):
             color=(1, 1, 0.5),
             random_state=self._episode_rng,
         )
-        self.model_builder.init_model_state(self.mpm_model, self.mpm_states[0])
+        self.model_builder.init_model_state(self.mpm_model, self.mpm_states)
         self.mpm_model.struct.static_ke = 100.0
         self.mpm_model.struct.static_kd = 0.0
         self.mpm_model.struct.static_mu = 1.0

mani_skill2/envs/mpm/hang_env.py

@@ -98,7 +98,7 @@ def _initialize_mpm(self):
             color=(1, 1, 0.5),
             random_state=self._episode_rng,
         )
-        self.model_builder.init_model_state(self.mpm_model, self.mpm_states[0])
+        self.model_builder.init_model_state(self.mpm_model, self.mpm_states)
         self.mpm_model.struct.static_ke = 100.0
         self.mpm_model.struct.static_kd = 0.0
         self.mpm_model.struct.static_mu = 1.0

mani_skill2/envs/mpm/pour_env.py

@@ -190,7 +190,7 @@ def _initialize_mpm(self):
         for i in range(len(self.model_builder.mpm_particle_volume)):
             self.model_builder.mpm_particle_volume[i] *= 1.2
 
-        self.model_builder.init_model_state(self.mpm_model, self.mpm_states[0])
+        self.model_builder.init_model_state(self.mpm_model, self.mpm_states)
         self.mpm_model.struct.static_ke = 100.0
         self.mpm_model.struct.static_kd = 0.0
         self.mpm_model.struct.static_mu = 1.0
@@ -327,7 +327,6 @@ def _determine_target_pos(self):
             if not success:
                 continue
 
-            # TODO: remove hardcode
             result[-2:] = 0.04
             self._init_qpos = result
             return

warp_maniskill/mpm/mpm_integrator.py

@@ -135,6 +135,7 @@ def zero_everything(
         if state.grid_m[grid_x, grid_y, grid_z] > 0.0:
             state.grid_m[grid_x, grid_y, grid_z] = 0.0
             state.grid_mv[grid_x, grid_y, grid_z] = wp.vec3(0.0)
+            state.grid_v[grid_x, grid_y, grid_z] = wp.vec3(0.0)
 
 
 @wp.kernel
@@ -340,6 +341,8 @@ def g2p(
 
     if wp.isnan(new_x[0]) or wp.isnan(new_x[1]) or wp.isnan(new_x[2]) or state_in.error[0] == 1:
         state_out.error[0] = 1
+    else:
+        state_out.error[0] = 0
 
     new_x = wp.vec3(
         wp.max(new_x[0], 3.0 * model.dx),

warp_maniskill/mpm/mpm_model.py

@@ -509,7 +509,7 @@ def finalize(self, device: str) -> MPMModel:
 
         return m
 
-    def init_model_state(self, model: MPMModel, state: MPMState):
+    def init_model_state(self, model: MPMModel, states):
         assert len(self.mpm_particle_q) <= model.max_n_particles
         model.struct.n_particles = len(self.mpm_particle_q)
         model.struct.particle_vol.assign(
@@ -524,23 +524,32 @@ def init_model_state(self, model: MPMModel, state: MPMState):
         model.struct.particle_friction_cohesion.assign(
             np.array(self.mpm_particle_friction_cohesion, dtype=np.float32)
         )
-        model.struct.particle_type.assign(
-            np.array(self.mpm_particle_type, dtype=int)
-        )
+        model.struct.particle_type.assign(np.array(self.mpm_particle_type, dtype=int))
         model.mpm_particle_colors = np.array(self.mpm_particle_colors, dtype=np.float32)
 
-        state.struct.particle_q.assign(np.array(self.mpm_particle_q, dtype=np.float32))
-        state.struct.particle_qd.assign(
-            np.array(self.mpm_particle_qd, dtype=np.float32)
-        )
+        for state in states:
+            state.struct.particle_q.assign(
+                np.array(self.mpm_particle_q, dtype=np.float32)
+            )
+            state.struct.particle_qd.assign(
+                np.array(self.mpm_particle_qd, dtype=np.float32)
+            )
 
-        eye = np.array([np.eye(3, dtype=np.float32)] * model.max_n_particles)
-        state.struct.particle_F.assign(eye)
-        state.struct.particle_volume_correction.zero_()
-        state.struct.particle_C.zero_()
-        state.struct.particle_vol.assign(
-            np.array(self.mpm_particle_volume, dtype=np.float32)
-        )
+            eye = np.array([np.eye(3, dtype=np.float32)] * model.max_n_particles)
+            state.struct.particle_F.assign(eye)
+            state.struct.particle_volume_correction.zero_()
+            state.struct.particle_C.zero_()
+            state.struct.particle_vol.assign(
+                np.array(self.mpm_particle_volume, dtype=np.float32)
+            )
+
+            state.struct.grid_lower.zero_()
+            state.struct.grid_upper.zero_()
+            state.struct.particle_f.zero_()
+            state.struct.grid_m.zero_()
+            state.struct.grid_mv.zero_()
+            state.struct.grid_v.zero_()
+            state.struct.error.zero_()
 
 
 def mpm_collide(model: MPMModel, state: MPMState):