Adding new randomization hook for sim steps (sub steps). This is necessary for correctly implement force / torque randomization.

2 files changed, 21 insertions, 2 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_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)