[fix] fix pose viz bug when isaacgym set_actor_scale function is calling for several times

yzqin · yzqin · commit 0012bc20094d · 2023-08-29T15:58:22.000-07:00
diff --git a/sim_web_visualizer/__init__.py b/sim_web_visualizer/__init__.py
@@ -6,7 +6,7 @@
 
 from .base_visualizer_client import MeshCatVisualizerBase
 
-__version__ = "0.4.0"
+__version__ = "0.5.0"
 
 try:
     import sapien.core as sapien
diff --git a/sim_web_visualizer/isaac_visualizer_client.py b/sim_web_visualizer/isaac_visualizer_client.py
@@ -72,6 +72,7 @@ def __init__(self, port: Optional[int] = None, host="localhost", keep_default_vi
         # Cache
         self.asset_resource_map: Dict[int, AssetResource] = {}
         self.asset_geom_map = {}
+        self.actor_asset_map = {}
         self.env_map = {}
         self.actor_rigid_body_name_map: Dict[int, List[str]] = {}
 
@@ -206,6 +207,8 @@ def create_actor(env: gymapi.Env, asset: gymapi.Asset, pose: gymapi.Transform, n
                     geom_path = f"{robot_tree_path}/{rigid_body_name}"
                     self.viz[geom_path].set_object(geom)
 
+                self.actor_asset_map[f"{num_env}/{num_actor}"] = id(asset)
+
             return actor
 
         self.new_gym.add_method("create_actor", create_actor)
@@ -279,8 +282,15 @@ def set_actor_scale(arg0: gymapi.Env, actor_id: int, scale: float):
             if arg0 in self.env_map:
                 env_num = self.env_map[arg0]
                 actor_path = f"/Sim/env:{env_num}/actor:{actor_id}"
-                scaling_transform = np.eye(4) * scale
-                self.viz[actor_path].set_transform(scaling_transform)
+                asset_id = self.actor_asset_map[f"{env_num}/{actor_id}"]
+                resource = self.asset_resource_map[asset_id]
+
+                for path, data in resource.visual_data.items():
+                    pose = resource.pose_data[path]
+                    scaling_pose = pose.copy()
+                    scaling_pose[:3] *= scale
+                    print(scaling_pose)
+                    self.viz[f"/{actor_path}/{path}"].set_transform(scaling_pose)
 
         self.new_gym.add_method("set_rigid_body_color", set_rigid_body_color)
         self.new_gym.add_method("set_actor_scale", set_actor_scale)
