🚀 [RofuncRL] Fix the mistake in ASE latent space construction

Skylark0924 · Skylark0924 · commit 69949ed83bb5 · 2023-06-26T01:06:58.000+08:00
diff --git a/examples/learning_rl/example_HumanoidASE_RofuncRL.py b/examples/learning_rl/example_HumanoidASE_RofuncRL.py
@@ -101,9 +101,9 @@ def inference(custom_args):
     # HumanoidASEReachSwordShield -> reallusion_sword_shield/RL_Avatar_Idle_Ready_Motion.npy
     # HumanoidASELocationSwordShield -> reallusion_sword_shield/RL_Avatar_Idle_Ready_Motion.npy
     # HumanoidASEStrikeSwordShield -> reallusion_sword_shield/RL_Avatar_Idle_Ready_Motion.npy
-    parser.add_argument("--task", type=str, default="HumanoidASEReachSwordShield")
+    parser.add_argument("--task", type=str, default="HumanoidASEGetupSwordShield")
     parser.add_argument("--motion_file", type=str,
-                        default="reallusion_sword_shield/RL_Avatar_Idle_Ready_Motion.npy")
+                        default="reallusion_sword_shield/dataset_reallusion_sword_shield.yaml")
     parser.add_argument("--agent", type=str, default="ase")  # Available agent: ase
     parser.add_argument("--num_envs", type=int, default=4096)
     parser.add_argument("--sim_device", type=str, default="cuda:{}".format(gpu_id))
diff --git a/rofunc/config/learning/rl/train/HumanoidASEGetupSwordShieldASERofuncRL.yaml b/rofunc/config/learning/rl/train/HumanoidASEGetupSwordShieldASERofuncRL.yaml
@@ -52,6 +52,9 @@ Agent:
   discriminator_weight_decay_scale: 0.0001
 
   ase_latent_dim: 64
+  ase_latent_steps_min: 1
+  ase_latent_steps_max: 150
+
   enc_reward_scale: 1
   enc_weight_decay_scale: 0
   enc_gradient_penalty_scale: 0
diff --git a/rofunc/config/learning/rl/train/HumanoidASEHeadingSwordShieldASERofuncRL.yaml b/rofunc/config/learning/rl/train/HumanoidASEHeadingSwordShieldASERofuncRL.yaml
@@ -45,6 +45,8 @@ Agent:
   use_gae: True                     # If true, use generalized advantage estimation.
 
   ase_latent_dim: 64
+  ase_latent_steps_min: 1
+  ase_latent_steps_max: 150
 
   task_reward_weight: 0.9
   style_reward_weight: 0.1
diff --git a/rofunc/config/learning/rl/train/HumanoidASELocationSwordShieldASERofuncRL.yaml b/rofunc/config/learning/rl/train/HumanoidASELocationSwordShieldASERofuncRL.yaml
@@ -45,6 +45,8 @@ Agent:
   use_gae: True                     # If true, use generalized advantage estimation.
 
   ase_latent_dim: 64
+  ase_latent_steps_min: 1
+  ase_latent_steps_max: 150
 
   task_reward_weight: 0.9
   style_reward_weight: 0.1
diff --git a/rofunc/config/learning/rl/train/HumanoidASEPerturbSwordShieldASERofuncRL.yaml b/rofunc/config/learning/rl/train/HumanoidASEPerturbSwordShieldASERofuncRL.yaml
@@ -52,6 +52,9 @@ Agent:
   discriminator_weight_decay_scale: 0.0001
 
   ase_latent_dim: 64
+  ase_latent_steps_min: 1
+  ase_latent_steps_max: 150
+
   enc_reward_scale: 1
   enc_weight_decay_scale: 0
   enc_gradient_penalty_scale: 0
diff --git a/rofunc/config/learning/rl/train/HumanoidASEReachSwordShieldASERofuncRL.yaml b/rofunc/config/learning/rl/train/HumanoidASEReachSwordShieldASERofuncRL.yaml
@@ -45,6 +45,8 @@ Agent:
   use_gae: True                     # If true, use generalized advantage estimation.
 
   ase_latent_dim: 64
+  ase_latent_steps_min: 1
+  ase_latent_steps_max: 150
 
   task_reward_weight: 0.9
   style_reward_weight: 0.1
diff --git a/rofunc/config/learning/rl/train/HumanoidASEStrikeSwordShieldASERofuncRL.yaml b/rofunc/config/learning/rl/train/HumanoidASEStrikeSwordShieldASERofuncRL.yaml
@@ -45,6 +45,8 @@ Agent:
   use_gae: True                     # If true, use generalized advantage estimation.
 
   ase_latent_dim: 64
+  ase_latent_steps_min: 1
+  ase_latent_steps_max: 150
 
   task_reward_weight: 0.9
   style_reward_weight: 0.1
diff --git a/rofunc/learning/RofuncRL/agents/mixline/ase_agent.py b/rofunc/learning/RofuncRL/agents/mixline/ase_agent.py
@@ -63,8 +63,6 @@ def __init__(self,
         """ASE specific parameters"""
         self._lr_e = cfg.Agent.lr_e
         self._ase_latent_dim = cfg.Agent.ase_latent_dim
-        # self._ase_latent_steps_min = self.cfg.Agent.ase_latent_steps_min
-        # self._ase_latent_steps_max = self.cfg.Agent.ase_latent_steps_max
         # self._amp_diversity_bonus = self.cfg.Agent.amp_diversity_bonus
         # self._amp_diversity_tar = self.cfg.Agent.amp_diversity_tar
         # self._enc_coef = self.cfg.Agent.enc_coef
@@ -92,24 +90,21 @@ def __init__(self,
         self.memory.create_tensor(name="ase_latents", size=self._ase_latent_dim, dtype=torch.float32)
         self._tensors_names.append("ase_latents")
 
+        self._ase_latents = torch.zeros((self.memory.num_envs, self._ase_latent_dim), dtype=torch.float32,
+                                        device=self.device)
+
     def _set_up(self):
         super()._set_up()
         self.optimizer_enc = torch.optim.Adam(self.encoder.parameters(), lr=self._lr_e, eps=self._adam_eps)
         if self._lr_scheduler is not None:
             self.scheduler_enc = self._lr_scheduler(self.optimizer_enc, **self._lr_scheduler_kwargs)
         self.checkpoint_modules["optimizer_enc"] = self.optimizer_enc
 
-    def _update_latents(self, num_envs: int):
-        # Equ. 11, provide the model with a latent space
-        z_bar = torch.normal(torch.zeros([num_envs, self._ase_latent_dim]))
-        self._ase_latents = z = torch.nn.functional.normalize(z_bar, dim=-1).to(self.device)
-
     def act(self, states: torch.Tensor, deterministic: bool = False, ase_latents: torch.Tensor = None):
         if self._current_states is not None:
             states = self._current_states
 
         if ase_latents is None:
-            self._update_latents(states.shape[0])
             ase_latents = self._ase_latents
 
         if not deterministic:
@@ -171,10 +166,10 @@ def update_net(self):
             amp_logits = self.discriminator(self._amp_state_preprocessor(amp_states))
             if self._least_square_discriminator:
                 style_rewards = torch.maximum(torch.tensor(1 - 0.25 * torch.square(1 - amp_logits)),
-                                             torch.tensor(0.0001, device=self.device))
+                                              torch.tensor(0.0001, device=self.device))
             else:
                 style_rewards = -torch.log(torch.maximum(torch.tensor(1 - 1 / (1 + torch.exp(-amp_logits))),
-                                                        torch.tensor(0.0001, device=self.device)))
+                                                         torch.tensor(0.0001, device=self.device)))
             style_rewards *= self._discriminator_reward_scale
 
             # Compute encoder reward
diff --git a/rofunc/learning/RofuncRL/agents/mixline/ase_hrl_agent.py b/rofunc/learning/RofuncRL/agents/mixline/ase_hrl_agent.py
@@ -127,7 +127,8 @@ def __init__(self,
         self._task_reward_weight = self.cfg.Agent.task_reward_weight
         self._style_reward_weight = self.cfg.Agent.style_reward_weight
         self._kl_threshold = self.cfg.Agent.kl_threshold
-        self._rewards_shaper = self.cfg.get("Agent", {}).get("rewards_shaper", lambda rewards: rewards * 0.01)
+        self._rewards_shaper = None
+        # self._rewards_shaper = self.cfg.get("Agent", {}).get("rewards_shaper", lambda rewards: rewards * 0.01)
         self._state_preprocessor = RunningStandardScaler
         self._state_preprocessor_kwargs = self.cfg.get("Agent", {}).get("state_preprocessor_kwargs",
                                                                         {"size": observation_space, "device": device})
diff --git a/rofunc/learning/RofuncRL/trainers/ase_trainer.py b/rofunc/learning/RofuncRL/trainers/ase_trainer.py
@@ -13,6 +13,7 @@
  See the License for the specific language governing permissions and
  limitations under the License.
  """
+import torch
 
 from rofunc.learning.RofuncRL.agents.mixline.ase_agent import ASEAgent
 from rofunc.learning.RofuncRL.agents.mixline.ase_hrl_agent import ASEHRLAgent
@@ -46,9 +47,42 @@ def __init__(self, cfg, env, device, env_name, hrl=False):
                                       num_samples))
         self.setup_wandb()
 
+        '''Misc variables'''
+        self._latent_reset_steps = torch.zeros(self.env.num_envs, dtype=torch.int32).to(self.device)
+        self._latent_steps_min = self.cfg.Agent.ase_latent_steps_min
+        self._latent_steps_max = self.cfg.Agent.ase_latent_steps_max
+
+    def _reset_latents(self, env_ids):
+        # Equ. 11, provide the model with a latent space
+        z_bar = torch.normal(torch.zeros([len(env_ids), self.agent._ase_latent_dim]))
+        self.agent._ase_latents[env_ids] = torch.nn.functional.normalize(z_bar, dim=-1).to(self.device)
+
+    def _reset_latent_step_count(self, env_ids):
+        self._latent_reset_steps[env_ids] = torch.randint_like(self._latent_reset_steps[env_ids],
+                                                               low=self._latent_steps_min,
+                                                               high=self._latent_steps_max)
+
+    def _update_latents(self):
+        new_latent_envs = self._latent_reset_steps <= self.env.progress_buf
+
+        need_update = torch.any(new_latent_envs)
+        if need_update:
+            new_latent_env_ids = new_latent_envs.nonzero(as_tuple=False).flatten()
+            self._reset_latents(new_latent_env_ids)
+            self._latent_reset_steps[new_latent_env_ids] += torch.randint_like(
+                self._latent_reset_steps[new_latent_env_ids],
+                low=self._latent_steps_min,
+                high=self._latent_steps_max)
+
     def pre_interaction(self):
-        if self.collect_observation is not None:
-            self.agent._current_states = self.collect_observation()
+        if self.collect_observation is not None:  # Reset failed envs
+            obs_dict, done_env_ids = self.env.reset_done()
+            self.agent._current_states = obs_dict["obs"]
+            if len(done_env_ids) > 0:
+                self._reset_latents(done_env_ids)
+                self._reset_latent_step_count(done_env_ids)
+
+        self._update_latents()
 
     def post_interaction(self):
         self._rollout += 1
