# Copyright 2019 The dm_control Authors.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ============================================================================

"""Script for tuning Duplo stud sizes to give desired separation forces."""

import collections

import pprint

from absl import app

from absl import logging

from dm_control.entities.props import duplo

from dm_control.entities.props.duplo import utils

from scipy import optimize

# pylint: disable=protected-access,invalid-name

_StudSize = duplo._StudSize

ORIGINAL_STUD_SIZE_PARAMS = duplo._STUD_SIZE_PARAMS

# pylint: enable=protected-access,invalid-name

DESIRED_FORCES = _StudSize(minimum=6., lower_quartile=10., maximum=18.)

# The safety margin here is because the separation force isn't quite monotonic

# w.r.t. the stud radius. If we set the min and max radii according to the

# exact desired bounds on the separation force then we may occasionally sample

# stud radii that yield out-of-bounds forces.

SAFETY_MARGIN = 0.2

def get_separation_force_for_radius(radius, **duplo_kwargs):

"""Measures Duplo separation force as a function of stud radius."""

top_brick = duplo.Duplo(**duplo_kwargs)

bottom_brick = duplo.Duplo(**duplo_kwargs)

# Set the radius of the studs on the bottom brick (this would normally be done

# in `initialize_episode_mjcf`). Note: we also set the radius of the studs on

# the top brick, since this has a (tiny!) effect on its mass.

# pylint: disable=protected-access

top_brick._active_stud_dclass.geom.size[0] = radius

bottom_brick._active_stud_dclass.geom.size[0] = radius

# pylint: enable=protected-access

separation_force = utils.measure_separation_force(top_brick, bottom_brick)

logging.debug('Stud radius: %f\tseparation force: %f N',

radius, separation_force)

return separation_force

class _KeepBracketingSolutions:

"""Wraps objective func, keeps closest solutions bracketing the target."""

_solution = collections.namedtuple('_solution', ['x', 'residual'])

def __init__(self, func):

self._func = func

self.below = self._solution(x=None, residual=-float('inf'))

self.above = self._solution(x=None, residual=float('inf'))

def __call__(self, x):

residual = self._func(x)

if self.below.residual < residual <= 0:

self.below = self._solution(x=x, residual=residual)

elif 0 < residual < self.above.residual:

self.above = self._solution(x=x, residual=residual)

return residual

@property

def closest(self):

if abs(self.below.residual) < self.above.residual:

return self.below

else:

return self.above

def tune_stud_radius(desired_force,

min_radius=0.0045,

max_radius=0.005,

desired_places=6,

side='closest',

**duplo_kwargs):

"""Find a stud size that gives the desired separation force."""

@_KeepBracketingSolutions

def func(radius):

radius = round(radius, desired_places) # Round radius for aesthetics (!)

return (get_separation_force_for_radius(radius=radius, **duplo_kwargs)

- desired_force)

# Ensure that the min and max radii bracket the solution.

while func(min_radius) > 0:

min_radius = max(1e-3, min_radius - (max_radius - min_radius))

while func(max_radius) < 0:

max_radius += (max_radius - min_radius)

tolerance = 10**-(desired_places)

# Use bisection to refine the bounds on the optimal radius. Note: this assumes

# that separation force is monotonic w.r.t. stud radius, but this isn't

# exactly true in all cases.

optimize.bisect(func, a=min_radius, b=max_radius, xtol=tolerance, disp=True)

if side == 'below':

solution = func.below

elif side == 'above':

solution = func.above

else:

solution = func.closest

radius = round(solution.x, desired_places)

force = get_separation_force_for_radius(radius, **duplo_kwargs)

return radius, force

def main(argv):

if len(argv) > 1:

raise app.UsageError('Too many command-line arguments.')

tuned_stud_radii = {}

tuned_separation_forces = {}

for stud_params in sorted(ORIGINAL_STUD_SIZE_PARAMS):

duplo_kwargs = stud_params._asdict()

min_result = tune_stud_radius(

desired_force=DESIRED_FORCES.minimum + SAFETY_MARGIN,

variation=0.0, side='above', **duplo_kwargs)

lq_result = tune_stud_radius(

desired_force=DESIRED_FORCES.lower_quartile,

variation=0.0, side='closest', **duplo_kwargs)

max_result = tune_stud_radius(

desired_force=DESIRED_FORCES.maximum - SAFETY_MARGIN,

variation=0.0, side='below', **duplo_kwargs)

radii, forces = zip(*(min_result, lq_result, max_result))

logging.info('\nDuplo configuration: %s\nTuned radii: %s, forces: %s',

stud_params, radii, forces)

tuned_stud_radii[stud_params] = _StudSize(*radii)

tuned_separation_forces[stud_params] = _StudSize(*forces)

logging.info('%s\nNew Duplo parameters:\n%s\nSeparation forces:\n%s',

'-'*60,

pprint.pformat(tuned_stud_radii),

pprint.pformat(tuned_separation_forces))

if __name__ == '__main__':

app.run(main)