in grolp/gen/scripts/generate_trajectories.py [0:0]
def sample_task_params(succ_traj, full_traj, fail_traj,
goal_candidates, pickup_candidates, movable_candidates, receptacle_candidates, scene_candidates,
inject_noise=10):
# Get the current conditional distributions of all variables (goal/pickup/receptacle/scene).
goal_weight = [(1 / (1 + np.random.randint(0, inject_noise + 1) + succ_traj.loc[
(succ_traj['pickup'].isin(pickup_candidates) if 'pickup' in goal_to_required_variables[c] else True) &
(succ_traj['movable'].isin(movable_candidates) if 'movable' in goal_to_required_variables[c] else True) &
(succ_traj['receptacle'].isin(receptacle_candidates) if 'receptacle' in goal_to_required_variables[c] else True)
& (succ_traj['scene'].isin(scene_candidates) if 'scene' in goal_to_required_variables[c] else True)]
['goal'].tolist().count(c))) # Conditional.
* (1 / (1 + succ_traj['goal'].tolist().count(c))) # Prior.
for c in goal_candidates]
goal_probs = [w / sum(goal_weight) for w in goal_weight]
pickup_weight = [(1 / (1 + np.random.randint(0, inject_noise + 1) +
sum([succ_traj.loc[
succ_traj['goal'].isin([g]) &
(succ_traj['movable'].isin(movable_candidates)
if 'movable' in goal_to_required_variables[g] else True) &
(succ_traj['receptacle'].isin(receptacle_candidates)
if 'receptacle' in goal_to_required_variables[g] else True) &
(succ_traj['scene'].isin(scene_candidates)
if 'scene' in goal_to_required_variables[g] else True)]
['pickup'].tolist().count(c) for g in goal_candidates])))
* (1 / (1 + succ_traj['pickup'].tolist().count(c)))
for c in pickup_candidates]
pickup_probs = [w / sum(pickup_weight) for w in pickup_weight]
movable_weight = [(1 / (1 + np.random.randint(0, inject_noise + 1) +
sum([succ_traj.loc[
succ_traj['goal'].isin([g]) &
(succ_traj['pickup'].isin(pickup_candidates)
if 'pickup' in goal_to_required_variables[g] else True) &
(succ_traj['receptacle'].isin(receptacle_candidates)
if 'receptacle' in goal_to_required_variables[g] else True) &
(succ_traj['scene'].isin(scene_candidates)
if 'scene' in goal_to_required_variables[g] else True)]
['movable'].tolist().count(c) for g in goal_candidates])))
* (1 / (1 + succ_traj['movable'].tolist().count(c)))
for c in movable_candidates]
movable_probs = [w / sum(movable_weight) for w in movable_weight]
receptacle_weight = [(1 / (1 + np.random.randint(0, inject_noise + 1) +
sum([succ_traj.loc[
succ_traj['goal'].isin([g]) &
(succ_traj['pickup'].isin(pickup_candidates)
if 'pickup' in goal_to_required_variables[g] else True) &
(succ_traj['movable'].isin(movable_candidates)
if 'movable' in goal_to_required_variables[g] else True) &
(succ_traj['scene'].isin(scene_candidates)
if 'scene' in goal_to_required_variables[g] else True)]
['receptacle'].tolist().count(c) for g in goal_candidates])))
* (1 / (1 + succ_traj['receptacle'].tolist().count(c)))
for c in receptacle_candidates]
receptacle_probs = [w / sum(receptacle_weight) for w in receptacle_weight]
scene_weight = [(1 / (1 + np.random.randint(0, inject_noise + 1) +
sum([succ_traj.loc[
succ_traj['goal'].isin([g]) &
(succ_traj['pickup'].isin(pickup_candidates)
if 'pickup' in goal_to_required_variables[g] else True) &
(succ_traj['movable'].isin(movable_candidates)
if 'movable' in goal_to_required_variables[g] else True) &
(succ_traj['receptacle'].isin(receptacle_candidates)
if 'receptacle' in goal_to_required_variables[g] else True)]
['scene'].tolist().count(c) for g in goal_candidates])))
* (1 / (1 + succ_traj['scene'].tolist().count(c)))
for c in scene_candidates]
scene_probs = [w / sum(scene_weight) for w in scene_weight]
# Calculate the probability difference between each value and the maximum so we can iterate over them to find a
# next-best candidate to sample subject to the constraints of knowing which will fail.
diffs = [("goal", goal_candidates[idx], goal_probs[idx] - min(goal_probs))
for idx in range(len(goal_candidates)) if len(goal_candidates) > 1]
diffs.extend([("pickup", pickup_candidates[idx], pickup_probs[idx] - min(pickup_probs))
for idx in range(len(pickup_candidates)) if len(pickup_candidates) > 1])
diffs.extend([("movable", movable_candidates[idx], movable_probs[idx] - min(movable_probs))
for idx in range(len(movable_candidates)) if len(movable_candidates) > 1])
diffs.extend([("receptacle", receptacle_candidates[idx], receptacle_probs[idx] - min(receptacle_probs))
for idx in range(len(receptacle_candidates)) if len(receptacle_candidates) > 1])
diffs.extend([("scene", scene_candidates[idx], scene_probs[idx] - min(scene_probs))
for idx in range(len(scene_candidates)) if len(scene_candidates) > 1])
# Iteratively pop the next biggest difference until we find a combination that is valid (e.g., not already
# flagged as impossible by the simulator).
variable_value_by_diff = {}
diffs_as_keys = [] # list of diffs; index into list will be used as key values.
for _, _, diff in diffs:
already_keyed = False
for existing_diff in diffs_as_keys:
if np.isclose(existing_diff, diff):
already_keyed = True
break
if not already_keyed:
diffs_as_keys.append(diff)
for variable, value, diff in diffs:
key = None
for kidx in range(len(diffs_as_keys)):
if np.isclose(diffs_as_keys[kidx], diff):
key = kidx
if key not in variable_value_by_diff:
variable_value_by_diff[key] = []
variable_value_by_diff[key].append((variable, value))
for key, diff in sorted(enumerate(diffs_as_keys), key=lambda x: x[1], reverse=True):
variable_value = variable_value_by_diff[key]
random.shuffle(variable_value)
for variable, value in variable_value:
# Select a goal.
if variable == "goal":
gtype = value
# print("sampled goal '%s' with prob %.4f" % (gtype, goal_probs[goal_candidates.index(gtype)]))
_goal_candidates = [gtype]
_pickup_candidates = pickup_candidates[:]
_movable_candidates = movable_candidates[:]
_receptacle_candidates = receptacle_candidates[:]
_scene_candidates = scene_candidates[:]
# Select a pickup object.
elif variable == "pickup":
pickup_obj = value
# print("sampled pickup object '%s' with prob %.4f" %
# (pickup_obj, pickup_probs[pickup_candidates.index(pickup_obj)]))
_pickup_candidates = [pickup_obj]
_goal_candidates = goal_candidates[:]
_movable_candidates = movable_candidates[:]
_receptacle_candidates = receptacle_candidates[:]
_scene_candidates = scene_candidates[:]
# Select a movable object.
elif variable == "movable":
movable_obj = value
# print("sampled movable object '%s' with prob %.4f" %
# (movable_obj, movable_probs[movable_candidates.index(movable_obj)]))
_movable_candidates = [movable_obj]
_goal_candidates = [g for g in goal_candidates if g == 'pick_and_place_with_movable_recep']
_pickup_candidates = pickup_candidates[:]
_receptacle_candidates = receptacle_candidates[:]
_scene_candidates = scene_candidates[:]
# Select a receptacle.
elif variable == "receptacle":
receptacle_obj = value
# print("sampled receptacle object '%s' with prob %.4f" %
# (receptacle_obj, receptacle_probs[receptacle_candidates.index(receptacle_obj)]))
_receptacle_candidates = [receptacle_obj]
_goal_candidates = goal_candidates[:]
_pickup_candidates = pickup_candidates[:]
_movable_candidates = movable_candidates[:]
_scene_candidates = scene_candidates[:]
# Select a scene.
else:
sampled_scene = value
# print("sampled scene %s with prob %.4f" %
# (sampled_scene, scene_probs[scene_candidates.index(sampled_scene)]))
_scene_candidates = [sampled_scene]
_goal_candidates = goal_candidates[:]
_pickup_candidates = pickup_candidates[:]
_movable_candidates = movable_candidates[:]
_receptacle_candidates = receptacle_candidates[:]
# Perform constraint propagation to determine whether this is a valid assignment.
propagation_finished = False
while not propagation_finished:
assignment_lens = (len(_goal_candidates), len(_pickup_candidates), len(_movable_candidates),
len(_receptacle_candidates), len(_scene_candidates))
# Constraints on goal.
_goal_candidates = [g for g in _goal_candidates if
(g not in goal_to_pickup_type or
len(set(_pickup_candidates).intersection( # Pickup constraint.
constants.VAL_ACTION_OBJECTS[goal_to_pickup_type[g]])) > 0)
and (g not in goal_to_receptacle_type or
np.any([r in constants.VAL_ACTION_OBJECTS[goal_to_receptacle_type[g]]
for r in _receptacle_candidates])) # Valid by goal receptacle const.
and (g not in goal_to_invalid_receptacle or
len(set(_receptacle_candidates).difference(
goal_to_invalid_receptacle[g])) > 0) # Invalid by goal receptacle const.
and len(set(_scene_candidates).intersection(
scenes_for_goal[g])) > 0 # Scene constraint
]
# Define whether to consider constraints for each role based on current set of candidate goals.
pickup_constrained = np.any(["pickup" in goal_to_required_variables[g] for g in _goal_candidates])
movable_constrained = np.any(["movable" in goal_to_required_variables[g] for g in _goal_candidates])
receptacle_constrained = np.any(["receptacle" in goal_to_required_variables[g]
for g in _goal_candidates])
scene_constrained = np.any(["scene" in goal_to_required_variables[g] for g in _goal_candidates])
# Constraints on pickup obj.
_pickup_candidates = [p for p in _pickup_candidates if
np.any([g not in goal_to_pickup_type or
p in constants.VAL_ACTION_OBJECTS[goal_to_pickup_type[g]]
for g in _goal_candidates]) # Goal constraint.
and (not movable_constrained or
np.any([p in constants.VAL_RECEPTACLE_OBJECTS[m]
for m in _movable_candidates])) # Movable constraint.
and (not receptacle_constrained or
np.any([r in constants.VAL_ACTION_OBJECTS["Toggleable"] or
p in constants.VAL_RECEPTACLE_OBJECTS[r]
for r in _receptacle_candidates])) # Receptacle constraint.
and (not scene_constrained or
np.any([s in obj_to_scene_ids[constants.OBJ_PARENTS[p]]
for s in _scene_candidates])) # Scene constraint
]
# Constraints on movable obj.
_movable_candidates = [m for m in _movable_candidates if
'pick_and_place_with_movable_recep' in _goal_candidates # Goal constraint
and (not pickup_constrained or
np.any([p in constants.VAL_RECEPTACLE_OBJECTS[m]
for p in _pickup_candidates])) # Pickup constraint.
and (not receptacle_constrained or
np.any([r in constants.VAL_RECEPTACLE_OBJECTS and
m in constants.VAL_RECEPTACLE_OBJECTS[r]
for r in _receptacle_candidates])) # Receptacle constraint.
and (not scene_constrained or
np.any([s in obj_to_scene_ids[constants.OBJ_PARENTS[m]]
for s in _scene_candidates])) # Scene constraint
]
# Constraints on receptacle obj.
_receptacle_candidates = [r for r in _receptacle_candidates if
np.any([(g not in goal_to_receptacle_type or
r in constants.VAL_ACTION_OBJECTS[goal_to_receptacle_type[g]]) and
(g not in goal_to_invalid_receptacle or
r not in goal_to_invalid_receptacle[g])
for g in _goal_candidates]) # Goal constraint.
and (not receptacle_constrained or
r in constants.VAL_ACTION_OBJECTS["Toggleable"] or
np.any([p in constants.VAL_RECEPTACLE_OBJECTS[r]
for p in _pickup_candidates])) # Pickup constraint.
and (not movable_constrained or
r in constants.VAL_ACTION_OBJECTS["Toggleable"] or
np.any([m in constants.VAL_RECEPTACLE_OBJECTS[r]
for m in _movable_candidates])) # Movable constraint.
and (not scene_constrained or
np.any([s in obj_to_scene_ids[constants.OBJ_PARENTS[r]]
for s in _scene_candidates])) # Scene constraint
]
# Constraints on scene.
_scene_candidates = [s for s in _scene_candidates if
np.any([s in scenes_for_goal[g]
for g in _goal_candidates]) # Goal constraint.
and (not pickup_constrained or
np.any([obj_to_scene_ids[constants.OBJ_PARENTS[p]]
for p in _pickup_candidates])) # Pickup constraint.
and (not movable_constrained or
np.any([obj_to_scene_ids[constants.OBJ_PARENTS[m]]
for m in _movable_candidates])) # Movable constraint.
and (not receptacle_constrained or
np.any([obj_to_scene_ids[constants.OBJ_PARENTS[r]]
for r in _receptacle_candidates])) # Receptacle constraint.
]
if assignment_lens == (len(_goal_candidates), len(_pickup_candidates), len(_movable_candidates),
len(_receptacle_candidates), len(_scene_candidates)):
propagation_finished = True
candidate_lens = {"goal": len(_goal_candidates), "pickup": len(_pickup_candidates),
"movable": len(_movable_candidates), "receptacle": len(_receptacle_candidates),
"scene": len(_scene_candidates)}
if candidate_lens["goal"] == 0:
# print("Goal over-constrained; skipping")
continue
if np.all([0 in [candidate_lens[v] for v in goal_to_required_variables[g]] for g in _goal_candidates]):
continue
# Ensure some combination of the remaining constraints is not in failures and is not already populated
# by the target number of repeats.
failure_ensured = True
full_ensured = True
for g in _goal_candidates:
pickup_iter = _pickup_candidates if "pickup" in goal_to_required_variables[g] else ["None"]
for p in pickup_iter:
movable_iter = _movable_candidates if "movable" in goal_to_required_variables[g] else ["None"]
for m in movable_iter:
receptacle_iter = _receptacle_candidates if "receptacle" in goal_to_required_variables[g] \
else ["None"]
for r in receptacle_iter:
scene_iter = _scene_candidates if "scene" in goal_to_required_variables[g] else ["None"]
for s in scene_iter:
if (g, p, m, r, s) not in fail_traj:
failure_ensured = False
if (g, p, m, r, s) not in full_traj:
full_ensured = False
if not failure_ensured and not full_ensured:
break
if not failure_ensured and not full_ensured:
break
if not failure_ensured and not full_ensured:
break
if not failure_ensured and not full_ensured:
break
if not failure_ensured and not full_ensured:
break
if failure_ensured:
continue
if full_ensured:
continue
if candidate_lens["goal"] > 1 or np.any([np.any([candidate_lens[v] > 1
for v in goal_to_required_variables[g]])
for g in _goal_candidates]):
task_sampler = sample_task_params(succ_traj, full_traj, fail_traj,
_goal_candidates, _pickup_candidates, _movable_candidates,
_receptacle_candidates, _scene_candidates)
sampled_task = next(task_sampler)
if sampled_task is None:
continue
else:
g = _goal_candidates[0]
p = _pickup_candidates[0] if "pickup" in goal_to_required_variables[g] else "None"
m = _movable_candidates[0] if "movable" in goal_to_required_variables[g] else "None"
r = _receptacle_candidates[0] if "receptacle" in goal_to_required_variables[g] else "None"
s = _scene_candidates[0] if "scene" in goal_to_required_variables[g] else "None"
sampled_task = (g, p, m, r, int(s))
yield sampled_task
yield None # Discovered that there are no valid assignments remaining.