1"""An environment randomizer that randomizes physical parameters from config."""
2
3from __future__ import absolute_import
4from __future__ import division
5from __future__ import print_function
6
7import functools
8import random
9
10import os, inspect
11currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
12parentdir = os.path.dirname(os.path.dirname(currentdir))
13parentdir = os.path.dirname(os.path.dirname(parentdir))
14os.sys.path.insert(0, parentdir)
15
16import numpy as np
17import tf.compat.v1 as tf
18from pybullet_envs.minitaur.envs import env_randomizer_base
19from pybullet_envs.minitaur.envs.env_randomizers import minitaur_env_randomizer_config
20
21SIMULATION_TIME_STEP = 0.001
22
23
24class MinitaurEnvRandomizerFromConfig(env_randomizer_base.EnvRandomizerBase):
25  """A randomizer that change the minitaur_gym_env during every reset."""
26
27  def __init__(self, config=None):
28    if config is None:
29      config = "all_params"
30    try:
31      config = getattr(minitaur_env_randomizer_config, config)
32    except AttributeError:
33      raise ValueError("Config {} is not found.".format(config))
34    self._randomization_param_dict = config()
35    tf.logging.info("Randomization config is: {}".format(self._randomization_param_dict))
36
37  def randomize_env(self, env):
38    """Randomize various physical properties of the environment.
39
40    It randomizes the physical parameters according to the input configuration.
41
42    Args:
43      env: A minitaur gym environment.
44    """
45    self._randomization_function_dict = self._build_randomization_function_dict(env)
46    for param_name, random_range in self._randomization_param_dict.iteritems():
47      self._randomization_function_dict[param_name](lower_bound=random_range[0],
48                                                    upper_bound=random_range[1])
49
50  def _build_randomization_function_dict(self, env):
51    func_dict = {}
52    func_dict["mass"] = functools.partial(self._randomize_masses, minitaur=env.minitaur)
53    func_dict["inertia"] = functools.partial(self._randomize_inertia, minitaur=env.minitaur)
54    func_dict["latency"] = functools.partial(self._randomize_latency, minitaur=env.minitaur)
55    func_dict["joint friction"] = functools.partial(self._randomize_joint_friction,
56                                                    minitaur=env.minitaur)
57    func_dict["motor friction"] = functools.partial(self._randomize_motor_friction,
58                                                    minitaur=env.minitaur)
59    func_dict["restitution"] = functools.partial(self._randomize_contact_restitution,
60                                                 minitaur=env.minitaur)
61    func_dict["lateral friction"] = functools.partial(self._randomize_contact_friction,
62                                                      minitaur=env.minitaur)
63    func_dict["battery"] = functools.partial(self._randomize_battery_level, minitaur=env.minitaur)
64    func_dict["motor strength"] = functools.partial(self._randomize_motor_strength,
65                                                    minitaur=env.minitaur)
66    # Settinmg control step needs access to the environment.
67    func_dict["control step"] = functools.partial(self._randomize_control_step, env=env)
68    return func_dict
69
70  def _randomize_control_step(self, env, lower_bound, upper_bound):
71    randomized_control_step = random.uniform(lower_bound, upper_bound)
72    env.set_time_step(randomized_control_step)
73    tf.logging.info("control step is: {}".format(randomized_control_step))
74
75  def _randomize_masses(self, minitaur, lower_bound, upper_bound):
76    base_mass = minitaur.GetBaseMassesFromURDF()
77    random_base_ratio = random.uniform(lower_bound, upper_bound)
78    randomized_base_mass = random_base_ratio * np.array(base_mass)
79    minitaur.SetBaseMasses(randomized_base_mass)
80    tf.logging.info("base mass is: {}".format(randomized_base_mass))
81
82    leg_masses = minitaur.GetLegMassesFromURDF()
83    random_leg_ratio = random.uniform(lower_bound, upper_bound)
84    randomized_leg_masses = random_leg_ratio * np.array(leg_masses)
85    minitaur.SetLegMasses(randomized_leg_masses)
86    tf.logging.info("leg mass is: {}".format(randomized_leg_masses))
87
88  def _randomize_inertia(self, minitaur, lower_bound, upper_bound):
89    base_inertia = minitaur.GetBaseInertiasFromURDF()
90    random_base_ratio = random.uniform(lower_bound, upper_bound)
91    randomized_base_inertia = random_base_ratio * np.array(base_inertia)
92    minitaur.SetBaseInertias(randomized_base_inertia)
93    tf.logging.info("base inertia is: {}".format(randomized_base_inertia))
94    leg_inertia = minitaur.GetLegInertiasFromURDF()
95    random_leg_ratio = random.uniform(lower_bound, upper_bound)
96    randomized_leg_inertia = random_leg_ratio * np.array(leg_inertia)
97    minitaur.SetLegInertias(randomized_leg_inertia)
98    tf.logging.info("leg inertia is: {}".format(randomized_leg_inertia))
99
100  def _randomize_latency(self, minitaur, lower_bound, upper_bound):
101    randomized_latency = random.uniform(lower_bound, upper_bound)
102    minitaur.SetControlLatency(randomized_latency)
103    tf.logging.info("control latency is: {}".format(randomized_latency))
104
105  def _randomize_joint_friction(self, minitaur, lower_bound, upper_bound):
106    num_knee_joints = minitaur.GetNumKneeJoints()
107    randomized_joint_frictions = np.random.uniform([lower_bound] * num_knee_joints,
108                                                   [upper_bound] * num_knee_joints)
109    minitaur.SetJointFriction(randomized_joint_frictions)
110    tf.logging.info("joint friction is: {}".format(randomized_joint_frictions))
111
112  def _randomize_motor_friction(self, minitaur, lower_bound, upper_bound):
113    randomized_motor_damping = random.uniform(lower_bound, upper_bound)
114    minitaur.SetMotorViscousDamping(randomized_motor_damping)
115    tf.logging.info("motor friction is: {}".format(randomized_motor_damping))
116
117  def _randomize_contact_restitution(self, minitaur, lower_bound, upper_bound):
118    randomized_restitution = random.uniform(lower_bound, upper_bound)
119    minitaur.SetFootRestitution(randomized_restitution)
120    tf.logging.info("foot restitution is: {}".format(randomized_restitution))
121
122  def _randomize_contact_friction(self, minitaur, lower_bound, upper_bound):
123    randomized_foot_friction = random.uniform(lower_bound, upper_bound)
124    minitaur.SetFootFriction(randomized_foot_friction)
125    tf.logging.info("foot friction is: {}".format(randomized_foot_friction))
126
127  def _randomize_battery_level(self, minitaur, lower_bound, upper_bound):
128    randomized_battery_voltage = random.uniform(lower_bound, upper_bound)
129    minitaur.SetBatteryVoltage(randomized_battery_voltage)
130    tf.logging.info("battery voltage is: {}".format(randomized_battery_voltage))
131
132  def _randomize_motor_strength(self, minitaur, lower_bound, upper_bound):
133    randomized_motor_strength_ratios = np.random.uniform([lower_bound] * minitaur.num_motors,
134                                                         [upper_bound] * minitaur.num_motors)
135    minitaur.SetMotorStrengthRatios(randomized_motor_strength_ratios)
136    tf.logging.info("motor strength is: {}".format(randomized_motor_strength_ratios))
137