updated pybullet example costs and tested fix for collisions

examples/panda_robot_with_obstacles.py

@@ -42,10 +42,23 @@ def compute_cost(self, sim):
 class EndEffectorGoalObjective(object):
     def __init__(self, cfg, device):
         self.nav_goal = torch.tensor(cfg.goal, device=cfg.mppi.device)
+        self.ort_goal = torch.tensor([1, 0, 0, 0], device=device)
+        self.w_coll = 1.
+        self.w_pos = 1.5
+        self.w_ort = 0.
 
     def compute_cost(self, sim):
         pos = sim.rigid_body_state[:, sim.robot_rigid_body_ee_idx, :3]
-        return 10 * torch.linalg.norm(pos - self.nav_goal, axis=1)
+        ort = sim.rigid_body_state[:, sim.robot_rigid_body_ee_idx, 3:7]
+
+        reach_cost = torch.linalg.norm(pos - self.nav_goal, axis=1)
+        align_cost = torch.linalg.norm(ort - self.ort_goal, axis=1)
+
+        # Collision avoidance with contact forces
+        xyz_contatcs = torch.sum(torch.abs(torch.cat((sim.net_cf[:, 0].unsqueeze(1), sim.net_cf[:, 1].unsqueeze(1), sim.net_cf[:, 2].unsqueeze(1)), 1)),1)
+        coll_cost = torch.sum(xyz_contatcs.reshape([sim.num_envs, int(xyz_contatcs.size(dim=0)/sim.num_envs)])[:, 1:sim.num_bodies], 1) # skip the first, it is the robot
+
+        return reach_cost * self.w_pos + align_cost * self.w_ort + coll_cost * self.w_coll
 
 
 def initalize_environment(cfg):

examples/point_robot_with_obstacle.py

@@ -26,29 +26,28 @@ class Objective(object):
     def __init__(self, cfg, device):
         self.nav_goal = torch.tensor(cfg.goal, device=cfg.mppi.device)
 
-        self.w_nav = 1.0
-        self.w_obs = 0.5
-
+        self.w_nav = 1.0 # 5.0
+        self.w_obs = 1.0  
+        self.w_coll = 0.1 # 0.01 
+
     def compute_cost(self, sim: IsaacGymWrapper):
         dof_state = sim.dof_state
         pos = torch.cat((dof_state[:, 0].unsqueeze(1), dof_state[:, 2].unsqueeze(1)), 1)
         obs_positions = sim.obstacle_positions
 
-        nav_cost = torch.clamp(
-            torch.linalg.norm(pos - self.nav_goal, axis=1) - 0.05, min=0, max=1999
-        )
-
-        # sim.gym.refresh_net_contact_force_tensor(sim.sim)
-        # sim.net_cf
+        nav_cost = torch.linalg.norm(pos - self.nav_goal, axis=1)
 
-        # This can cause steady state error if the goal is close to an obstacle, better use contact forces later on
+        # Coll avoidance with distance
         obs_cost = torch.sum(
             1 / torch.linalg.norm(obs_positions[:, :, :2] - pos.unsqueeze(1), axis=2),
             axis=1,
         )
 
-        return nav_cost * self.w_nav + obs_cost * self.w_obs
+        # Collision avoidance with contact forces
+        xy_contatcs = torch.sum(torch.abs(torch.cat((sim.net_cf[:, 0].unsqueeze(1), sim.net_cf[:, 1].unsqueeze(1)), 1)),1)
+        coll = torch.sum(xy_contatcs.reshape([sim.num_envs, int(xy_contatcs.size(dim=0)/sim.num_envs)])[:, 1:sim.num_bodies], 1) # skip the first, it is the robot
 
+        return nav_cost * self.w_nav + coll * self.w_coll # + obs_cost * self.w_obs
 
 def initalize_environment(cfg) -> UrdfEnv:
     """
@@ -154,7 +153,7 @@ def run_point_robot(cfg: ExampleConfig):
     """
     # Note: Workaround to trigger the dataclasses __post_init__ method
     cfg = OmegaConf.to_object(cfg)
-    OmegaConf.save(config=cfg, f='test.yaml')
+    # OmegaConf.save(configs=cfg, f='test.yaml')
 
     env = initalize_environment(cfg)
     planner = set_planner(cfg)