Update push randomizer to use the new randomize_sub_step hook. This fixes issue that causes application of push force/torque is only valid for single simulation step.

1 files changed, 29 insertions, 7 deletions

diff --git a/examples/pybullet/gym/pybullet_envs/minitaur/envs/env_randomizers/minitaur_push_randomizer.py b/examples/pybullet/gym/pybullet_envs/minitaur/envs/env_randomizers/minitaur_push_randomizer.py
index 8d9aa8e64..aa4adb7d4 100644
--- a/examples/pybullet/gym/pybullet_envs/minitaur/envs/env_randomizers/minitaur_push_randomizer.py
+++ b/examples/pybullet/gym/pybullet_envs/minitaur/envs/env_randomizers/minitaur_push_randomizer.py
@@ -10,7 +10,9 @@ parentdir = os.path.dirname(os.path.dirname(currentdir))
 parentdir = os.path.dirname(os.path.dirname(parentdir))
 os.sys.path.insert(0, parentdir)
 
+import functools
 import math
+import gin
 import numpy as np
 from pybullet_envs.minitaur.envs import env_randomizer_base
 
@@ -23,6 +25,7 @@ _VERTICAL_FORCE_UPPER_BOUND = 300
 _VERTICAL_FORCE_LOWER_BOUND = 500
 
 
+@gin.configurable
 class MinitaurPushRandomizer(env_randomizer_base.EnvRandomizerBase):
   """Applies a random impulse to the base of Minitaur."""
 
@@ -54,6 +57,7 @@ class MinitaurPushRandomizer(env_randomizer_base.EnvRandomizerBase):
                                     [_HORIZONTAL_FORCE_LOWER_BOUND, _HORIZONTAL_FORCE_UPPER_BOUND])
     self._vertical_force_bound = (vertical_force_bound if vertical_force_bound else
                                   [_VERTICAL_FORCE_LOWER_BOUND, _VERTICAL_FORCE_UPPER_BOUND])
+    self._perturbation_parameter_dict = None
 
   def randomize_env(self, env):
     """Randomizes the simulation environment.
@@ -64,9 +68,10 @@ class MinitaurPushRandomizer(env_randomizer_base.EnvRandomizerBase):
     pass
 
   def randomize_step(self, env):
-    """Randomizes simulation steps.
+    """Randomizes env steps.
 
-    Will be called at every timestep. May add random forces/torques to Minitaur.
+    Will be called at every env step. Called to generate randomized  force and
+    torque to apply. Application of forces are done in randomize_sub_step.
 
     Args:
       env: The Minitaur gym environment to be randomized.
@@ -85,8 +90,25 @@ class MinitaurPushRandomizer(env_randomizer_base.EnvRandomizerBase):
     if (env.env_step_counter % self._perturbation_interval_steps <
         self._perturbation_duration_steps) and (env.env_step_counter >=
                                                 self._perturbation_start_step):
-      env.pybullet_client.applyExternalForce(objectUniqueId=env.minitaur.quadruped,
-                                             linkIndex=self._applied_link_id,
-                                             forceObj=self._applied_force,
-                                             posObj=[0.0, 0.0, 0.0],
-                                             flags=env.pybullet_client.LINK_FRAME)
+      # Parameter of pybullet_client.applyExternalForce()
+      self._perturbation_parameter_dict = dict(objectUniqueId=env.minitaur.quadruped,
+                                               linkIndex=self._applied_link_id,
+                                               forceObj=self._applied_force,
+                                               posObj=[0.0, 0.0, 0.0],
+                                               flags=env.pybullet_client.LINK_FRAME)
+    else:
+      self._perturbation_parameter_dict = None
+
+  def randomize_sub_step(self, env, sub_step_index, num_sub_steps):
+    """Randomize simulation steps per sub steps (simulation step).
+
+    Will be called at every simulation step. This is the correct place to add
+    random forces/torques to Minitaur.
+
+    Args:
+      env: The Minitaur gym environment to be randomized.
+      sub_step_index: Index of sub step, from 0 to N-1. N is the action repeat.
+      num_sub_steps: Number of sub steps, equals to action repeat.
+    """
+    if self._perturbation_parameter_dict is not None:
+      env.pybullet_client.applyExternalForce(**self._perturbation_parameter_dict)