def _choose_distractor_position(self) -> None:
self.distractor_position = None
theta = random.uniform(
*get_range(self.target_angle_range)) * random.choice([-1., 1.])
dx, dy, dz = [
self.get_record_dimensions(self.distractor)[i] *
self.distractor_scale[k] * 0.5 for i, k in enumerate(XYZ)
]
offset = max(dx, dy, dz)
try:
dpos = random.uniform(
*get_range(self.target_position_range)) + offset
except:
dpos = random.uniform(
*get_range(self.target_position_range['x'])) + offset
## if relation is null, make sure distractor is on opposite side
if self.left_or_right is None:
l_or_r = random.choice([-90, 90])
else:
l_or_r = -self.left_or_right
unit_v = self.rotate_vector_parallel_to_floor(
self.opposite_unit_vector, theta + l_or_r)
self.distractor_position = {
"x": unit_v["x"] * dpos,
"y": self.container_height,
"z": unit_v["z"] * dpos
}
## jitter position
for k in ["x", "z"]:
self.distractor_position[k] += random.uniform(
-self.target_position_jitter, self.target_position_jitter)
def _choose_target_position(self) -> None:
self.target_position = self.left_or_right = None
theta = random.uniform(
*get_range(self.target_angle_range)) * random.choice([-1., 1.])
tx, ty, tz = [
self.get_record_dimensions(self.target)[i] * self.target_scale[k] *
0.5 for i, k in enumerate(XYZ)
]
offset = max(tx, ty, tz)
try:
tpos = random.uniform(
*get_range(self.target_position_range)) + offset
except:
tpos = random.uniform(
*get_range(self.target_position_range['x'])) + offset
## if contain or support, place the target on the container;
if (self.relation == Relation.support) or (self.relation
== Relation.contain):
self.target_position = copy.deepcopy(self.container_position)
self.target_position["y"] = self.container_height
## elif occlude, put it mostly behind the container
elif self.relation == Relation.occlude:
unit_v = self.rotate_vector_parallel_to_floor(
self.opposite_unit_vector, theta)
self.target_position = {
"x": unit_v["x"] * tpos,
"y": self.container_height,
"z": unit_v["z"] * tpos
}
## elif null, put it to one side of the container
elif self.relation == Relation.null:
self.left_or_right = random.choice([-90, 90])
unit_v = self.rotate_vector_parallel_to_floor(
self.opposite_unit_vector, theta + self.left_or_right)
self.target_position = {
"x": unit_v["x"] * tpos,
"y": self.container_height,
"z": unit_v["z"] * tpos
}
else:
raise NotImplementedError(
"You need to institute a rule for this relation type")
## jitter position
for k in ["x", "z"]:
self.target_position[k] += random.uniform(
-self.target_position_jitter, self.target_position_jitter)
def __init__(self,
port=1071,
num_middle_objects=2,
middle_scale_range={
'x': 0.5,
'y': 1.0,
'z': 0.5
},
spacing_jitter=0.25,
lateral_jitter=0.0,
middle_friction=0.1,
**kwargs):
# middle config
Gravity.__init__(self,
port=port,
middle_friction=middle_friction,
**kwargs)
MultiDominoes.__init__(self,
launch_build=False,
middle_objects='cube',
num_middle_objects=num_middle_objects,
middle_scale_range=middle_scale_range,
spacing_jitter=spacing_jitter,
lateral_jitter=lateral_jitter,
**kwargs)
# raise the ramp
if hasattr(self.middle_scale_range, 'keys'):
self.pit_max_height = get_range(self.middle_scale_range['y'])[1]
elif hasattr(self.middle_scale_range, '__len__'):
self.pit_max_height = self.middle_scale_range[1]
else:
self.pit_max_height = self.middle_scale_range
self.ramp_base_height_range = [
r + self.pit_max_height
for r in get_range(self.ramp_base_height_range)
]
def _switch_target_link(self) -> List[dict]:
commands = []
if self.target_contained_range is None:
self.target_link_idx = random.choice(range(self.num_links))
elif hasattr(self.target_contained_range, '__len__'):
self.target_link_idx = int(
random.choice(range(*get_range(self.target_contained_range))))
self.target_link_idx = min(self.target_link_idx,
self.num_links - 1)
elif isinstance(self.target_contained_range, (int, float)):
self.target_link_idx = int(self.target_contained_range)
else:
return []
print("target is link idx %d" % self.target_link_idx)
if int(self.target_link_idx) not in range(self.num_links):
print("no target link")
return [] # no link is the target
record, data = self.blocks[self.target_link_idx]
o_id = data['id']
# update the data so that it's the target
if self.target_color is None:
self.target_color = self.random_color()
data['color'] = self.target_color
self._replace_target_with_object(record, data)
# add the commands to change the material and color
commands.extend(
self.get_object_material_commands(
record, o_id, self.get_material_name(self.target_material)))
# Scale the object and set its color.
rgb = data['color']
commands.extend([{
"$type": "set_color",
"color": {
"r": rgb[0],
"g": rgb[1],
"b": rgb[2],
"a": 1.
},
"id": o_id
}])
return commands
def _place_intermediate_object(self) -> List[dict]:
"""
Place a primitive object at the room center.
"""
# create a target object
# XXX TODO: Why is scaling part of random primitives
# but rotation and translation are not?
# Consider integrating!
record, data = self.random_primitive(self._target_types,
scale=self.target_scale_range,
color=self.probe_color)
o_id, scale, rgb = [data[k] for k in ["id", "scale", "color"]]
self.target_type = data["name"]
self.target = record
# self.target_id = o_id #for internal purposes, the other object is the target
# self.object_color = rgb
self.target_position = TDWUtils.VECTOR3_ZERO
# add the object
commands = []
if self.target_rotation is None:
self.target_rotation = self.get_rotation(self.target_rotation_range)
commands.extend(
self.add_physics_object(
record=record,
position=self.target_position,
rotation=self.target_rotation,
mass= random.uniform(*get_range(self.middle_mass_range)),
dynamic_friction=0.4, #increased friction
static_friction=0.4,
bounciness=0,
o_id=o_id))
# Set the object material
commands.extend(
self.get_object_material_commands(
record, o_id, self.get_material_name(self.target_material)))
# Scale the object and set its color.
commands.extend([
{"$type": "set_color",
"color": {"r": rgb[0], "g": rgb[1], "b": rgb[2], "a": 1.},
"id": o_id},
{"$type": "scale_object",
"scale_factor": scale,
"id": o_id}])
return commands
def _place_and_push_target_object(self) -> List[dict]:
"""
Place a probe object at the other end of the collision axis, then apply a force to push it.
Using probe mass and location here to allow for ramps. There is no dedicated probe in this controller.
"""
# create a target object
record, data = self.random_primitive(self._target_types,
scale=self.target_scale_range,
color=self.target_color,
add_data=False)
# add_data=(not self.remove_target)
o_id, scale, rgb = [data[k] for k in ["id", "scale", "color"]]
self.target = record
self.target_type = data["name"]
self.target_color = rgb
self.target_scale = self.middle_scale = scale
self.target_id = o_id
# Where to put the target
if self.target_rotation is None:
self.target_rotation = self.get_rotation(self.target_rotation_range)
# Add the object with random physics values
commands = []
### TODO: better sampling of random physics values
self.probe_mass = random.uniform(self.probe_mass_range[0], self.probe_mass_range[1])
self.probe_initial_position = {"x": -0.5*self.collision_axis_length, "y": self.target_lift, "z": 0.}
rot = self.get_rotation(self.target_rotation_range)
print(rot)
if self.use_ramp:
commands.extend(self._place_ramp_under_probe())
self.probe_initial_position['x'] += self.target_lift #HACK rotation might've led to the object falling of the back of the ramp, so we're moving it forward
# commands.extend(
# self.add_physics_object(
# record=record,
# position=self.probe_initial_position,
# rotation=rot,
# mass=self.probe_mass,
# # dynamic_friction=0.5,
# # static_friction=0.5,
# # bounciness=0.1,
# dynamic_friction=0.4,
# static_friction=0.4,
# bounciness=0,
# o_id=o_id))
commands.extend(
self.add_primitive(
record=record,
position=self.probe_initial_position,
rotation=rot,
mass=self.probe_mass,
scale_mass=False,
material=self.target_material,
color=rgb,
scale=scale,
# dynamic_friction=0.5,
# static_friction=0.5,
# bounciness=0.1,
dynamic_friction=0.4,
static_friction=0.4,
bounciness=0,
o_id=o_id,
add_data=True
))
# Set the target material
# commands.extend(
# self.get_object_material_commands(
# record, o_id, self.get_material_name(self.target_material)))
# the target is the probe
self.target_position = self.probe_initial_position
# Scale the object and set its color.
# commands.extend([
# {"$type": "set_color",
# "color": {"r": rgb[0], "g": rgb[1], "b": rgb[2], "a": 1.},
# "id": o_id},
# {"$type": "scale_object",
# "scale_factor": scale,
# "id": o_id}])
# Set its collision mode
commands.extend([
# {"$type": "set_object_collision_detection_mode",
# "mode": "continuous_speculative",
# "id": o_id},
{"$type": "set_object_drag",
"id": o_id,
"drag": 0, "angular_drag": 0}])
# Apply a force to the target object
self.push_force = self.get_push_force(
scale_range=self.probe_mass * np.array(self.force_scale_range),
angle_range=self.force_angle_range,
yforce=self.fupforce)
self.push_force = self.rotate_vector_parallel_to_floor(
self.push_force, -rot['y'], degrees=True)
self.push_position = self.probe_initial_position
if self.use_ramp:
self.push_cmd = {
"$type": "apply_force_to_object",
"force": self.push_force,
"id": int(o_id)
}
else:
self.push_position = {
k:v+self.force_offset[k]*self.rotate_vector_parallel_to_floor(
self.target_scale, rot['y'])[k]
for k,v in self.push_position.items()}
self.push_position = {
k:v+random.uniform(-self.force_offset_jitter, self.force_offset_jitter)
for k,v in self.push_position.items()}
self.push_cmd = {
"$type": "apply_force_at_position",
"force": self.push_force,
"position": self.push_position,
"id": int(o_id)
}
# decide when to apply the force
self.force_wait = int(random.uniform(*get_range(self.force_wait_range)))
print("force wait", self.force_wait)
if self.force_wait == 0:
commands.append(self.push_cmd)
return commands
# def _place_target_object(self) -> List[dict]:
"""
Place a primitive object at one end of the collision axis.
"""
# create a target object
record, data = self.random_primitive(self._target_types,
scale=self.target_scale_range,
color=self.target_color,
add_data=(not self.remove_target)
)
o_id, scale, rgb = [data[k] for k in ["id", "scale", "color"]]
self.target = record
self.target_type = data["name"]
self.target_color = rgb
self.target_scale = self.middle_scale = scale
self.target_id = o_id
if any((s <= 0 for s in scale.values())):
self.remove_target = True
# Where to put the target
if self.target_rotation is None:
self.target_rotation = self.get_rotation(self.target_rotation_range)
if self.target_position is None:
self.target_position = {
"x": 0.5 * self.collision_axis_length,
"y": 0. if not self.remove_target else 10.0,
"z": 0. if not self.remove_target else 10.0
}
# Commands for adding hte object
commands = []
commands.extend(
self.add_physics_object(
record=record,
position=self.target_position,
rotation=self.target_rotation,
mass=2.0,
dynamic_friction=0.5,
static_friction=0.5,
bounciness=0.0,
o_id=o_id,
add_data=(not self.remove_target)
))
# Set the object material
commands.extend(
self.get_object_material_commands(
record, o_id, self.get_material_name(self.target_material)))
# Scale the object and set its color.
commands.extend([
{"$type": "set_color",
"color": {"r": rgb[0], "g": rgb[1], "b": rgb[2], "a": 1.},
"id": o_id},
{"$type": "scale_object",
"scale_factor": scale if not self.remove_target else TDWUtils.VECTOR3_ZERO,
"id": o_id}])
# If this scene won't have a target
if self.remove_target:
commands.append(
{"$type": self._get_destroy_object_command_name(o_id),
"id": int(o_id)})
self.object_ids = self.object_ids[:-1]
return commands
cmds = []
# ramp params
self.ramp = random.choice(self.DEFAULT_RAMPS)
rgb = self.ramp_color or np.array([0.75,0.75,1.0])
ramp_pos = copy.deepcopy(self.probe_initial_position)
ramp_pos['y'] += self.zone_scale['y'] if not self.remove_zone else 0.0 # don't intersect w zone
ramp_rot = self.get_y_rotation([180,180])
ramp_id = self._get_next_object_id()
# figure out scale
r_len, r_height, r_dep = self.get_record_dimensions(self.ramp)
scale_x = (0.75 * self.collision_axis_length) / r_len
if self.ramp_scale is None:
self.ramp_scale = arr_to_xyz([scale_x, self.scale_to(r_height, 1.5), 0.75 * scale_x])
# optionally add base
self.ramp_base_height = random.uniform(*get_range(self.ramp_base_height_range))
if self.ramp_base_height > 0.01:
self.ramp_base = self.CUBE
self.ramp_base_scale = arr_to_xyz([
float(scale_x * r_len), float(self.ramp_base_height), float(0.75 * scale_x * r_dep)])
self.ramp_base_id = self._get_next_object_id()
cmds.extend(
self.add_physics_object(
record=self.ramp_base,
position=copy.deepcopy(ramp_pos),
rotation=TDWUtils.VECTOR3_ZERO,
mass=500,
dynamic_friction=0.01,
static_friction=0.01,
bounciness=0.0,
o_id=self.ramp_base_id,
add_data=True))
_,rb_height,_ = self.get_record_dimensions(self.ramp_base)
ramp_pos['y'] += self.ramp_base_scale['y']
# scale it, color it, fix it
cmds.extend(
self.get_object_material_commands(
self.ramp_base, self.ramp_base_id, self.get_material_name(self.zone_material)))
cmds.extend([
{"$type": "scale_object",
"scale_factor": self.ramp_base_scale,
"id": self.ramp_base_id},
{"$type": "set_color",
"color": {"r": rgb[0], "g": rgb[1], "b": rgb[2], "a": 1.},
"id": self.ramp_base_id},
{"$type": "set_object_collision_detection_mode",
"mode": "continuous_speculative",
"id": self.ramp_base_id},
{"$type": "set_kinematic_state",
"id": self.ramp_base_id,
"is_kinematic": True,
"use_gravity": True}])
cmds.extend(
self.add_ramp(
record = self.ramp,
position=ramp_pos,
rotation=ramp_rot,
scale=self.ramp_scale,
o_id=ramp_id,
add_data=True))
# give the ramp a texture and color
cmds.extend(
self.get_object_material_commands(
self.ramp, ramp_id, self.get_material_name(self.zone_material)))
cmds.append(
{"$type": "set_color",
"color": {"r": rgb[0], "g": rgb[1], "b": rgb[2], "a": 1.},
"id": ramp_id})
print("ramp commands")
print(cmds)
# need to adjust probe height as a result of ramp placement
self.probe_initial_position['x'] -= 0.5 * self.ramp_scale['x'] * r_len - 0.15
self.probe_initial_position['y'] = self.ramp_scale['y'] * r_height + self.ramp_base_height
return cmds
def _place_attachment(self) -> List[dict]:
commands = []
record, data = self.random_primitive(self._attachment_types,
scale=self.attachment_scale_range,
color=self.attachment_color,
exclude_color=self.target_color,
add_data=self.use_attachment)
o_id, scale, rgb = [data[k] for k in ["id", "scale", "color"]]
self.attachment = record
self.attachment_id = data['id']
self.attachment_type = data['name']
mass = random.uniform(*get_range(self.attachment_mass_range))
mass *= (np.prod(xyz_to_arr(scale)) /
np.prod(xyz_to_arr(self.STANDARD_BLOCK_SCALE)))
if self.attachment_type == 'cylinder':
mass *= (np.pi / 4.0)
elif self.attachment_type == 'cone':
mass *= (np.pi / 12.0)
print("attachment mass", mass)
commands.extend(
self.add_physics_object(record=record,
position={
"x": 0.,
"y": self.tower_height,
"z": 0.
},
rotation={
"x": 0.,
"y": 0.,
"z": 0.
},
mass=mass,
dynamic_friction=0.5,
static_friction=0.5,
bounciness=0,
o_id=o_id,
add_data=self.use_attachment))
commands.extend(
self.get_object_material_commands(
record, o_id,
self.get_material_name(self.attachment_material)))
# Scale the object and set its color.
commands.extend([{
"$type": "set_color",
"color": {
"r": rgb[0],
"g": rgb[1],
"b": rgb[2],
"a": 1.
},
"id": o_id
}, {
"$type": "scale_object",
"scale_factor": scale,
"id": o_id
}])
# for an attachment that is wider at base, better to place links a little higher
a_len, a_height, a_dep = self.get_record_dimensions(record)
if self.attachment_type == 'cone' and self.use_attachment:
self.tower_height += 0.25 * a_height * (
np.sqrt(scale['x']**2 + scale['z']**2) / scale['y'])
if not self.use_attachment:
commands.append({
"$type":
self._get_destroy_object_command_name(o_id),
"id":
int(o_id)
})
self.object_ids = self.object_ids[:-1]
# fix it to ground or block
if self.attachment_fixed_to_base and self.use_attachment:
# make it kinematic
if self.use_base:
commands.append({
"$type": "add_fixed_joint",
"parent_id": self.base_id,
"id": o_id
})
elif not self.use_base: # make kinematic
commands.extend([{
"$type": "set_object_collision_detection_mode",
"mode": "continuous_speculative",
"id": o_id
}, {
"$type": "set_kinematic_state",
"id": o_id,
"is_kinematic": True,
"use_gravity": True
}])
# add a cap
if self.use_cap and self.use_attachment:
commands.extend(
self._build_base(height=(self.tower_height +
a_height * scale['y']),
as_cap=True))
if self.attachment_fixed_to_base:
commands.append({
"$type": "add_fixed_joint",
"parent_id": o_id,
"id": self.cap_id
})
return commands
def _build_base(self, height, as_cap=False) -> List[dict]:
commands = []
record, data = self.random_primitive(self._base_types,
scale=self.base_scale_range,
color=self.base_color,
exclude_color=self.target_color,
add_data=self.use_base)
o_id, scale, rgb = [data[k] for k in ["id", "scale", "color"]]
if as_cap:
self.cap_id = data['id']
self.cap_type = data['name']
else:
self.base_id = data['id']
self.base_type = data['name']
mass = random.uniform(*get_range(self.base_mass_range))
mass *= (np.prod(xyz_to_arr(scale)) /
np.prod(xyz_to_arr(self.STANDARD_BLOCK_SCALE)))
print("base mass", mass)
commands.extend(
self.add_physics_object(record=record,
position={
"x": 0.,
"y": height,
"z": 0.
},
rotation={
"x": 0.,
"y": 0.,
"z": 0.
},
mass=mass,
dynamic_friction=0.5,
static_friction=0.5,
bounciness=0,
o_id=o_id,
add_data=self.use_base))
commands.extend(
self.get_object_material_commands(
record, o_id, self.get_material_name(self.base_material)))
# Scale the object and set its color.
commands.extend([{
"$type": "set_color",
"color": {
"r": rgb[0],
"g": rgb[1],
"b": rgb[2],
"a": 1.
},
"id": o_id
}, {
"$type": "scale_object",
"scale_factor": scale,
"id": o_id
}])
if self.use_base:
if self.base_type not in ['pipe', 'torus']:
b_len, b_height, b_dep = self.get_record_dimensions(record)
self.tower_height += b_height * scale['y']
else:
commands.append({
"$type":
self._get_destroy_object_command_name(o_id),
"id":
int(o_id)
})
self.object_ids = self.object_ids[:-1]
self.tower_height = 0.0
return commands
def _add_cap(self) -> List[dict]:
commands = []
# the cap becomes the target
record, data = self.random_primitive(self._cap_types,
scale=self.target_scale_range,
color=self.target_color,
add_data=self.use_cap)
o_id, scale, rgb = [data[k] for k in ["id", "scale", "color"]]
self.cap = record
self.cap_type = data["name"]
if self.use_cap:
self._replace_target_with_object(record, data)
mass = random.uniform(*get_range(self.middle_mass_range))
mass *= (np.prod(xyz_to_arr(scale)) /
np.prod(xyz_to_arr(self.STANDARD_BLOCK_SCALE)))
mass *= self.STANDARD_MASS_FACTOR
commands.extend(
self.add_physics_object(record=record,
position={
"x": 0.,
"y": self.tower_height,
"z": 0.
},
rotation={
"x": 0.,
"y": 0.,
"z": 0.
},
mass=mass,
dynamic_friction=0.5,
static_friction=0.5,
bounciness=0,
o_id=o_id,
add_data=self.use_cap))
# Set the cap object material
commands.extend(
self.get_object_material_commands(
record, o_id, self.get_material_name(self.target_material)))
# Scale the object and set its color.
commands.extend([{
"$type": "set_color",
"color": {
"r": rgb[0],
"g": rgb[1],
"b": rgb[2],
"a": 1.
},
"id": o_id
}, {
"$type": "scale_object",
"scale_factor": scale,
"id": o_id
}])
if not self.use_cap:
commands.append({
"$type":
self._get_destroy_object_command_name(o_id),
"id":
int(o_id)
})
self.object_ids = self.object_ids[:-1]
else:
self.tower_height += scale["y"]
return commands
def _build_stack(self) -> List[dict]:
commands = []
height = self.tower_height
height += self.zone_scale['y'] if not self.remove_zone else 0.0
# build the block scales
mid = self.num_blocks / 2.0
grad = self.middle_scale_gradient
self.block_scales = [
self._get_block_scale(offset=grad * (mid - i))
for i in range(self.num_blocks)
]
self.blocks = []
# place the blocks
for m in range(self.num_blocks):
record, data = self.random_primitive(
self._middle_types,
scale=self.block_scales[m],
color=self.middle_color,
exclude_color=self.target_color)
self.middle_type = data["name"]
o_id, scale, rgb = [data[k] for k in ["id", "scale", "color"]]
block_pos = self._get_block_position(scale, height)
block_rot = self.get_y_rotation(self.middle_rotation_range)
# scale the mass to give each block a uniform density
block_mass = random.uniform(*get_range(self.middle_mass_range))
block_mass *= (np.prod(xyz_to_arr(scale)) /
np.prod(xyz_to_arr(self.STANDARD_BLOCK_SCALE)))
block_mass *= self.STANDARD_MASS_FACTOR
commands.extend(
self.add_physics_object(record=record,
position=block_pos,
rotation=block_rot,
mass=block_mass,
dynamic_friction=random.uniform(
0, 0.9),
static_friction=random.uniform(0, 0.9),
bounciness=random.uniform(0, 1),
o_id=o_id))
# Set the block object material
commands.extend(
self.get_object_material_commands(
record, o_id,
self.get_material_name(self.middle_material)))
# Scale the object and set its color.
commands.extend([{
"$type": "set_color",
"color": {
"r": rgb[0],
"g": rgb[1],
"b": rgb[2],
"a": 1.
},
"id": o_id
}, {
"$type": "scale_object",
"scale_factor": scale,
"id": o_id
}])
print("placed middle object %s" % str(m + 1))
# update height
_y = record.bounds['top']['y'] if self.middle_type != 'bowl' else (
record.bounds['bottom']['y'] + 0.1)
height += scale["y"] * _y
data.update({
'position': block_pos,
'rotation': block_rot,
'mass': block_mass
})
print("middle object data", data)
self.blocks.append((record, data))
self.tower_height = height
return commands
def _place_and_push_probe_object(self) -> List[dict]:
"""
Place a probe object at the other end of the collision axis, then apply a force to push it.
"""
exclude = not (self.monochrome and self.match_probe_and_target_color)
record, data = self.random_primitive(
self._probe_types,
scale=self.probe_scale_range,
color=self.probe_color,
exclude_color=(self.target_color if exclude else None),
exclude_range=0.25)
o_id, scale, rgb = [data[k] for k in ["id", "scale", "color"]]
self.probe = record
self.probe_type = data["name"]
self.probe_scale = scale
self.probe_id = o_id
# Add the object with random physics values
commands = []
### TODO: better sampling of random physics values
self.probe_mass = random.uniform(self.probe_mass_range[0],
self.probe_mass_range[1])
self.probe_initial_position = {
"x": -0.5 * self.collision_axis_length,
"y": self.probe_lift,
"z": 0.
}
rot = self.get_rotation(self.probe_rotation_range)
if self.use_ramp:
commands.extend(self._place_ramp_under_probe())
commands.extend(
self.add_physics_object(
record=record,
position=self.probe_initial_position,
rotation=rot,
mass=self.probe_mass,
# dynamic_friction=0.5,
# static_friction=0.5,
# bounciness=0.1,
dynamic_friction=0.4,
static_friction=0.4,
bounciness=0,
o_id=o_id))
# Set the probe material
commands.extend(
self.get_object_material_commands(
record, o_id, self.get_material_name(self.probe_material)))
# Scale the object and set its color.
commands.extend([{
"$type": "set_color",
"color": {
"r": rgb[0],
"g": rgb[1],
"b": rgb[2],
"a": 1.
},
"id": o_id
}, {
"$type": "scale_object",
"scale_factor": scale,
"id": o_id
}])
# Set its collision mode
commands.extend([
# {"$type": "set_object_collision_detection_mode",
# "mode": "continuous_speculative",
# "id": o_id},
{
"$type": "set_object_drag",
"id": o_id,
"drag": 0,
"angular_drag": 0
}
])
# Apply a force to the probe object
self.push_force = self.get_push_force(
scale_range=self.probe_mass * np.array(self.force_scale_range),
angle_range=self.force_angle_range,
yforce=self.fupforce)
self.push_force = self.rotate_vector_parallel_to_floor(self.push_force,
0,
degrees=True)
self.push_position = self.probe_initial_position
if self.use_ramp:
self.push_cmd = {
"$type": "apply_force_to_object",
"force": self.push_force,
"id": int(o_id)
}
else:
self.push_position = {
k: v +
self.force_offset[k] * self.rotate_vector_parallel_to_floor(
self.probe_scale, rot['y'])[k]
for k, v in self.push_position.items()
}
self.push_position = {
k: v + random.uniform(-self.force_offset_jitter,
self.force_offset_jitter)
for k, v in self.push_position.items()
}
self.push_cmd = {
"$type": "apply_force_at_position",
"force": self.push_force,
"position": self.push_position,
"id": int(o_id)
}
# decide when to apply the force
self.force_wait = int(
random.uniform(*get_range(self.force_wait_range)))
if self.PRINT:
print("force wait", self.force_wait)
if self.force_wait == 0:
commands.append(self.push_cmd)
return commands
def _build_intermediate_structure(self) -> List[dict]:
commands = []
print("THIS IS A PIT!")
print(self.num_middle_objects)
# get the scale of the total pit object
scale = get_random_xyz_transform(self.middle_scale_range)
self.pit_mass = random.uniform(*get_range(self.middle_mass_range))
# get color and texture
self.pit_color = self.middle_color or self.random_color(
exclude=self.target_color)
self.pit_material = self.middle_material
# how wide are the pits?
self.pit_widths = [
random.uniform(0.0, self.spacing_jitter) * scale['x']
for _ in range(self.num_middle_objects - 1)
]
# make M cubes and scale in x accordingly
x_remaining = scale['x'] - self.pit_widths[0]
x_filled = 0.0
print("PIT WIDTHS", self.pit_widths)
for m in range(self.num_middle_objects):
print("x_filled, remaining", x_filled, x_remaining)
m_rec = random.choice(self._middle_types)
x_scale = random.uniform(0.0, x_remaining)
x_len, _, _ = self.get_record_dimensions(m_rec)
x_len *= x_scale
x_pos = self.ramp_end_x + x_filled + (0.5 * x_len)
z_pos = random.uniform(-self.lateral_jitter, self.lateral_jitter)
print(m)
print("ramp_end", self.ramp_end_x)
print("x_len", x_len)
print("x_scale", x_scale)
print("x_pos", x_pos)
print("z_pos", z_pos)
m_scale = arr_to_xyz([x_scale, scale['y'], scale['z']])
commands.extend(
self.add_primitive(record=m_rec,
position=arr_to_xyz([x_pos, 0., z_pos]),
rotation=TDWUtils.VECTOR3_ZERO,
scale=m_scale,
color=self.pit_color,
exclude_color=self.target_color,
material=self.pit_material,
mass=self.pit_mass,
dynamic_friction=self.middle_friction,
static_friction=self.middle_friction,
scale_mass=True,
make_kinematic=True,
add_data=True,
obj_list=self.middle_objects))
if m < len(self.pit_widths):
x_filled += self.pit_widths[m] + x_len
x_remaining -= (self.pit_widths[m] + x_len)
commands.extend(Gravity._build_intermediate_structure(self))
print("INTERMEDIATE")
print(commands)
return commands
