Merge pull request #4174 from sippeyxp/substep_rand

Fix incorrect application of disturbance push.

3 files changed, 50 insertions, 9 deletions

diff --git a/examples/pybullet/gym/pybullet_envs/minitaur/envs/env_randomizer_base.py b/examples/pybullet/gym/pybullet_envs/minitaur/envs/env_randomizer_base.py
index 8ab49b103..e07c8d505 100644
--- a/examples/pybullet/gym/pybullet_envs/minitaur/envs/env_randomizer_base.py
+++ b/examples/pybullet/gym/pybullet_envs/minitaur/envs/env_randomizer_base.py
@@ -25,11 +25,28 @@ class EnvRandomizerBase(object):
     pass
 
   def randomize_step(self, env):
-    """Randomize simulation steps.
+    """Randomize environment steps.
 
-    Will be called at every timestep. May add random forces/torques to Minitaur.
+    Will be called at every environment step.
+
+    It is NOT recommended to use this for force / torque disturbance because
+    pybullet applyExternalForce/Torque only persist for single simulation step
+    not the entire env step which can contain multiple simulation steps.
+
+    Args:
+      env: The Minitaur gym environment to be randomized.
+    """
+    pass
+
+  def randomize_sub_step(self, env, sub_step_index, num_sub_steps):
+    """Randomize simulation sub steps.
+
+    Will be called at every simulation step. This is the correct place to add
+    random forces/torques.
 
     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.
     """
     pass
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)
diff --git a/examples/pybullet/gym/pybullet_envs/minitaur/envs_v2/locomotion_gym_env.py b/examples/pybullet/gym/pybullet_envs/minitaur/envs_v2/locomotion_gym_env.py
index 0c30e5b8e..b98d3b005 100644
--- a/examples/pybullet/gym/pybullet_envs/minitaur/envs_v2/locomotion_gym_env.py
+++ b/examples/pybullet/gym/pybullet_envs/minitaur/envs_v2/locomotion_gym_env.py
@@ -457,6 +457,8 @@ class LocomotionGymEnv(gym.Env):
     for obj in self._dynamic_objects():
       obj.pre_control_step(autonomous_object.AUTONOMOUS_ACTION)
     for _ in range(self._num_action_repeat):
+      for env_randomizer in self._env_randomizers:
+        env_randomizer.randomize_sub_step(self, i, self._num_action_repeat)
       self._robot.apply_action(action)
       for obj in self._dynamic_objects():
         obj.update(self.get_time_since_reset(), self._observation_dict)