From 9a8f77a82765ae57f920ce71a9f0506238541f6c Mon Sep 17 00:00:00 2001 From: Fanbo Xiang Date: Mon, 15 Aug 2022 14:37:30 -0700 Subject: [PATCH] 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 | 4 +- mani_skill2/envs/mpm/base_env.py | 10 ++++- mani_skill2/envs/mpm/excavate_env.py | 4 +- mani_skill2/envs/mpm/fill_env.py | 2 +- mani_skill2/envs/mpm/hang_env.py | 2 +- mani_skill2/envs/mpm/pour_env.py | 3 +- warp_maniskill/mpm/mpm_integrator.py | 3 ++ warp_maniskill/mpm/mpm_model.py | 39 ++++++++++++-------- 8 files changed, 41 insertions(+), 26 deletions(-) diff --git a/mani_skill2/agents/configs/panda/variants.py b/mani_skill2/agents/configs/panda/variants.py index 90e1be8a2..1f10ad858 100644 --- a/mani_skill2/agents/configs/panda/variants.py +++ b/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, diff --git a/mani_skill2/envs/mpm/base_env.py b/mani_skill2/envs/mpm/base_env.py index 05bd21178..1b9cb3eee 100644 --- a/mani_skill2/envs/mpm/base_env.py +++ b/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() diff --git a/mani_skill2/envs/mpm/excavate_env.py b/mani_skill2/envs/mpm/excavate_env.py index 1811ece19..ef1a809c3 100644 --- a/mani_skill2/envs/mpm/excavate_env.py +++ b/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 diff --git a/mani_skill2/envs/mpm/fill_env.py b/mani_skill2/envs/mpm/fill_env.py index d3a09d09b..2e50b4a74 100644 --- a/mani_skill2/envs/mpm/fill_env.py +++ b/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 diff --git a/mani_skill2/envs/mpm/hang_env.py b/mani_skill2/envs/mpm/hang_env.py index 68f953d51..9fbd2f79c 100644 --- a/mani_skill2/envs/mpm/hang_env.py +++ b/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 diff --git a/mani_skill2/envs/mpm/pour_env.py b/mani_skill2/envs/mpm/pour_env.py index 1817463c8..3f03c6ede 100644 --- a/mani_skill2/envs/mpm/pour_env.py +++ b/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 diff --git a/warp_maniskill/mpm/mpm_integrator.py b/warp_maniskill/mpm/mpm_integrator.py index eda5c16a3..aca5fdb3a 100644 --- a/warp_maniskill/mpm/mpm_integrator.py +++ b/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), diff --git a/warp_maniskill/mpm/mpm_model.py b/warp_maniskill/mpm/mpm_model.py index 9ad38d161..a499b769e 100644 --- a/warp_maniskill/mpm/mpm_model.py +++ b/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):