def sample(sample_type="const",
a=None,
b=None,
*,
random_state: np.random.RandomState = None):
if sample_type == "const":
return a
elif sample_type == "randint":
return random_state.randint(a, b)
elif sample_type == "random":
return random_state.random() * (b - a) + a
else:
raise TypeError("unknown sample type: %s" % sample_type)
def move_one_object_to_the_air_with_restrictions(
current_placement: np.ndarray,
height_range: Tuple[float, float],
object_size: float,
random_state: np.random.RandomState,
pickup_proba: float = 0.0,
stacking_proba: float = 0.0,
goal_distance_ratio: float = 1.0,
) -> np.ndarray:
"""
Modify current_placement to move one object to the air.
Set height w.r.t goal_distance_ratio too.
With probability `pickup_proba`, we will randomly move one object up into the air.
Otherwise do nothing.
The pickup height is discounted by `goal_distance_ratio`.
"""
assert 0.0 <= pickup_proba + stacking_proba <= 1.0
if pickup_proba + stacking_proba == 0:
return current_placement
logger.info(
f"Move one object to the air with restrictions "
f"(goal_distance_ratio={goal_distance_ratio} "
f"pickup_proba={pickup_proba} stacking_proba={stacking_proba})")
p = random_state.random()
if p > pickup_proba + stacking_proba:
# rearrange task: do nothing.
return current_placement
elif p < pickup_proba:
# pick up task
min_h, max_h = height_range
height = random_state.uniform(low=min_h, high=max_h)
n_objects = current_placement.shape[0]
target_i = random_state.randint(n_objects)
# The height of a pick-and-place goal is set w.r.t the goal distance too.
current_placement[target_i, -1] += height * goal_distance_ratio
return current_placement
else:
# stacking task
n_objects = current_placement.shape[0]
if n_objects >= 2:
# Randomly pick block indices for building a tower of size [2, n_objects].
tower_size = random_state.randint(2, n_objects + 1)
obj_indices = np.random.choice(list(range(n_objects)),
size=tower_size,
replace=False)
assert len(obj_indices) > 1
logger.info(
f"Building a tower of size {tower_size} with object indices {obj_indices} ..."
)
target_i = obj_indices[0]
for height, i in enumerate(obj_indices[1:]):
current_placement[i, 0] = current_placement[target_i, 0]
current_placement[i, 1] = current_placement[target_i, 1]
current_placement[i, 2] += object_size * (height + 1) * 2
return current_placement
