def _write_frame(self, frames_grp: h5py.Group, resp: List[bytes], frame_num: int) -> \
Tuple[h5py.Group, h5py.Group, dict, bool]:
frame, objs, tr, done = super()._write_frame(frames_grp=frames_grp, resp=resp, frame_num=frame_num)
num_objects = len(self.object_ids)
# Physics data.
velocities = np.empty(dtype=np.float32, shape=(num_objects, 3))
angular_velocities = np.empty(dtype=np.float32, shape=(num_objects, 3))
# Collision data.
collision_ids = np.empty(dtype=np.int32, shape=(0, 2))
collision_relative_velocities = np.empty(dtype=np.float32, shape=(0, 3))
collision_contacts = np.empty(dtype=np.float32, shape=(0, 2, 3))
# Environment Collision data.
env_collision_ids = np.empty(dtype=np.int32, shape=(0, 1))
env_collision_contacts = np.empty(dtype=np.float32, shape=(0, 2, 3))
sleeping = True
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
if r_id == "rigi":
ri = Rigidbodies(r)
ri_dict = dict()
for i in range(ri.get_num()):
ri_dict.update({ri.get_id(i): {"vel": ri.get_velocity(i),
"ang": ri.get_angular_velocity(i)}})
# Check if any objects are sleeping that aren't in the abyss.
if not ri.get_sleeping(i) and tr[ri.get_id(i)]["pos"][1] >= -1:
sleeping = False
# Add the Rigibodies data.
for o_id, i in zip(self.object_ids, range(num_objects)):
velocities[i] = ri_dict[o_id]["vel"]
angular_velocities[i] = ri_dict[o_id]["ang"]
elif r_id == "coll":
co = Collision(r)
collision_ids = np.append(collision_ids, [co.get_collider_id(), co.get_collidee_id()])
collision_relative_velocities = np.append(collision_relative_velocities, co.get_relative_velocity())
for i in range(co.get_num_contacts()):
collision_contacts = np.append(collision_contacts, (co.get_contact_normal(i),
co.get_contact_point(i)))
elif r_id == "enco":
en = EnvironmentCollision(r)
env_collision_ids = np.append(env_collision_ids, en.get_object_id())
for i in range(en.get_num_contacts()):
env_collision_contacts = np.append(env_collision_contacts, (en.get_contact_normal(i),
en.get_contact_point(i)))
objs.create_dataset("velocities", data=velocities.reshape(num_objects, 3), compression="gzip")
objs.create_dataset("angular_velocities", data=angular_velocities.reshape(num_objects, 3), compression="gzip")
collisions = frame.create_group("collisions")
collisions.create_dataset("object_ids", data=collision_ids.reshape((-1, 2)), compression="gzip")
collisions.create_dataset("relative_velocities", data=collision_relative_velocities.reshape((-1, 3)),
compression="gzip")
collisions.create_dataset("contacts", data=collision_contacts.reshape((-1, 2, 3)), compression="gzip")
env_collisions = frame.create_group("env_collisions")
env_collisions.create_dataset("object_ids", data=env_collision_ids, compression="gzip")
env_collisions.create_dataset("contacts", data=env_collision_contacts.reshape((-1, 2, 3)),
compression="gzip")
return frame, objs, tr, sleeping
def get_object_target_collision(self, obj_id: int, target_id: int, resp: List[bytes]):
contact_points = []
contact_normals = []
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
if r_id == "coll":
co = Collision(r)
coll_ids = [co.get_collider_id(), co.get_collidee_id()]
if [obj_id, target_id] == coll_ids or [target_id, obj_id] == coll_ids:
contact_points = [co.get_contact_point(i) for i in range(co.get_num_contacts())]
contact_normals = [co.get_contact_normal(i) for i in range(co.get_num_contacts())]
return (contact_points, contact_normals)
def on_frame(self, resp: List[bytes]) -> List[dict]:
"""
Update the avatar based on its current arm-bending goals and its state.
If the avatar has achieved a goal (for example, picking up an object), it will stop moving that arm.
Update the avatar's state as needed.
:param resp: The response from the build.
:return: A list of commands to pick up, stop moving, etc.
"""
def _get_mitten_position(a: Arm) -> np.array:
"""
:param a: The arm.
:return: The position of a mitten.
"""
if a == Arm.left:
return np.array(frame.get_mitten_center_left_position())
else:
return np.array(frame.get_mitten_center_right_position())
# Update dynamic data.
frame = self._get_frame(resp=resp)
# Update dynamic collision data.
self.collisions.clear()
self.env_collisions.clear()
# Get each collision.
for i in range(len(resp) - 1):
r_id = OutputData.get_data_type_id(resp[i])
if r_id == "coll":
coll = Collision(resp[i])
collider_id = coll.get_collider_id()
collidee_id = coll.get_collidee_id()
# Check if this was a mitten, if we're supposed to stop if there's a collision,
# and if the collision was not with the target.
for arm in self._ik_goals:
if self._ik_goals[arm] is not None:
if (collider_id == self.mitten_ids[arm] or collidee_id == self.mitten_ids[arm]) and \
(collider_id not in frame.get_held_right() and
collider_id not in frame.get_held_left() and
collidee_id not in frame.get_held_right() and
collidee_id not in frame.get_held_left()) and \
self._ik_goals[arm].stop_on_mitten_collision and \
(self._ik_goals[arm].target is None or
(self._ik_goals[arm].pick_up_id != collidee_id and
self._ik_goals[arm].pick_up_id != collider_id)) and \
(collidee_id not in self.body_parts_static or
collider_id not in self.body_parts_static):
self.status = TaskStatus.mitten_collision
self._ik_goals[arm] = None
if self._debug:
print(
"Stopping because the mitten collided with something."
)
return self._stop_arm(arm=arm)
# Check if the collision includes a body part.
if collider_id in self.body_parts_static and collidee_id not in self.body_parts_static:
if collider_id not in self.collisions:
self.collisions[collider_id] = []
self.collisions[collider_id].append(collidee_id)
elif collidee_id in self.body_parts_static and collider_id not in self.body_parts_static:
if collidee_id not in self.collisions:
self.collisions[collidee_id] = []
self.collisions[collidee_id].append(collider_id)
elif r_id == "enco":
coll = EnvironmentCollision(resp[i])
collider_id = coll.get_object_id()
if collider_id in self.body_parts_static:
self.env_collisions.append(collider_id)
# Check if IK goals are done.
temp_goals: Dict[Arm, Optional[_IKGoal]] = dict()
# Get commands for the next frame.
commands: List[dict] = []
for arm in self._ik_goals:
# No IK goal on this arm.
if self._ik_goals[arm] is None:
temp_goals[arm] = None
# This is a dummy IK goal. Let it run.
elif self._ik_goals[arm].target is None:
temp_goals[arm] = self._ik_goals[arm]
else:
# Is the arm at the target?
mitten_position = _get_mitten_position(arm)
# If we're not trying to pick something up, check if we are at the target position.
if self._ik_goals[arm].pick_up_id is None:
# If we're at the position, stop.
d = np.linalg.norm(mitten_position -
self._ik_goals[arm].target)
if d < self._ik_goals[arm].precision:
if self._debug:
print(
f"{arm.name} mitten is at target position {self._ik_goals[arm].target}. Stopping."
)
commands.extend(self._stop_arm(arm=arm))
temp_goals[arm] = None
self.status = TaskStatus.success
# Keep bending the arm.
else:
temp_goals[arm] = self._ik_goals[arm]
self._ik_goals[arm].previous_distance = d
# If we're trying to pick something, check if it was picked up on the previous frame.
else:
if self._ik_goals[arm].pick_up_id in frame.get_held_left() or self._ik_goals[arm]. \
pick_up_id in frame.get_held_right():
if self._debug:
print(
f"{arm.name} mitten picked up {self._ik_goals[arm].pick_up_id}. Stopping."
)
commands.extend(self._stop_arm(arm=arm))
temp_goals[arm] = None
self.status = TaskStatus.success
# Keep bending the arm and trying to pick up the object.
else:
commands.extend([{
"$type":
"pick_up_proximity",
"distance":
0.02,
"radius":
0.05,
"grip":
1000,
"is_left":
arm == Arm.left,
"avatar_id":
self.id,
"object_ids": [self._ik_goals[arm].pick_up_id]
}, {
"$type":
"pick_up",
"grip":
1000,
"is_left":
arm == Arm.left,
"object_ids": [self._ik_goals[arm].pick_up_id],
"avatar_id":
self.id
}])
temp_goals[arm] = self._ik_goals[arm]
self._ik_goals = temp_goals
# Check if the arms are still moving.
temp_goals: Dict[Arm, Optional[_IKGoal]] = dict()
for arm in self._ik_goals:
# No IK goal on this arm.
if self._ik_goals[arm] is None:
temp_goals[arm] = None
else:
# Get the past and present angles.
if arm == Arm.left:
angles_0 = self.frame.get_angles_left()
angles_1 = frame.get_angles_left()
else:
angles_0 = self.frame.get_angles_right()
angles_1 = frame.get_angles_right()
# Try to stop any moving joints.
if self._ik_goals[arm].rotations is not None and self._ik_goals[
arm].pick_up_id is not None:
joint_profile = self._get_default_sticky_mitten_profile()
for angle, joint_name in zip(angles_1, Avatar.ANGLE_ORDER):
target_angle = self._ik_goals[arm].rotations[
joint_name]
# Check if the joint stopped moving. Ignore if the joint already stopped.
if target_angle > 0.01 and np.abs(angle - target_angle) < 0.01 and \
joint_name in self._ik_goals[arm].moving_joints:
self._ik_goals[arm].moving_joints.remove(
joint_name)
j = joint_name.split("_")
j_name = f"{j[0]}_{arm.name}"
axis = j[1]
# Set the name of the elbow to the expected profile key.
if "elbow" in joint_name:
profile_key = "elbow"
else:
profile_key = joint_name
if self._debug:
print(
f"{joint_name} {arm.name} slowing down: {np.abs(angle - target_angle)}"
)
# Stop the joint from moving any more.
# Set the damper, force, and angular drag to "default" (non-moving) values.
commands.extend([{
"$type":
"set_joint_damper",
"joint":
j_name,
"axis":
axis,
"damper":
joint_profile[profile_key]["damper"],
"avatar_id":
self.id
}, {
"$type":
"set_joint_force",
"joint":
j_name,
"axis":
axis,
"force":
joint_profile[profile_key]["force"],
"avatar_id":
self.id
}, {
"$type":
"set_joint_angular_drag",
"joint":
j_name,
"axis":
axis,
"angular_drag":
joint_profile[profile_key]["angular_drag"],
"avatar_id":
self.id
}])
# Is any joint still moving?
moving = False
for a0, a1 in zip(angles_0, angles_1):
if np.abs(a0 - a1) > 0.03:
moving = True
break
# Keep moving.
if moving:
temp_goals[arm] = self._ik_goals[arm]
else:
if self._ik_goals[arm].rotations is not None:
# This is a reset arm action.
if self._ik_goals[arm].target is None:
mitten_position = _get_mitten_position(
arm) - frame.get_position()
d = np.linalg.norm(
self._initial_mitten_positions[arm] -
mitten_position)
# The reset arm action ended with the mitten very close to the initial position.
if d < self._ik_goals[arm].precision:
self.status = TaskStatus.success
else:
self.status = TaskStatus.no_longer_bending
# This is a regular action.
# It ended with the arm no longer moving but having never reached the target.
else:
if self._debug:
print(
f"{arm.name} is no longer bending. Cancelling."
)
self.status = TaskStatus.no_longer_bending
commands.extend(self._stop_arm(arm=arm))
temp_goals[arm] = None
self._ik_goals = temp_goals
self.frame = frame
return commands
def trial(self, d_f, s_f, b_f, d_c, s_c, b_c, collision_types):
"""
Collide a chair with a fridge.
:param d_f: The dynamic friction of the fridge.
:param s_f: The static friction of the fridge.
:param b_f: The bounciness of the fridge.
:param d_c: The dynamic friction of the chair.
:param s_c: The static friction of the chair.
:param b_c: The bounciness of the chair.
:param collision_types: The types of collisions to listen for.
"""
print("###\n\nNew Trial\n\n###\n")
fridge_id = 0
chair_id = 1
# Destroy all objects currently in the scene.
init_commands = [{"$type": "destroy_all_objects"}]
# Create the avatar.
init_commands.extend(
TDWUtils.create_avatar(position={
"x": 1,
"y": 2.5,
"z": 5
},
look_at=TDWUtils.VECTOR3_ZERO))
# Add the objects.
# Set the masses and physic materials.
# Apply a force to the chair.
# Receive collision data (note that by setting "stay" to True, you will receive a LOT of data;
# see "Performance Optimizations" documentation.)
init_commands.extend([
self.get_add_object("fridge_large", object_id=fridge_id),
self.get_add_object("chair_billiani_doll",
object_id=chair_id,
position={
"x": 4,
"y": 0,
"z": 0
}), {
"$type": "set_mass",
"id": fridge_id,
"mass": 40
}, {
"$type": "set_mass",
"id": chair_id,
"mass": 20
}, {
"$type": "set_physic_material",
"id": fridge_id,
"dynamic_friction": d_f,
"static_friction": s_f,
"bounciness": b_f
}, {
"$type": "set_physic_material",
"id": chair_id,
"dynamic_friction": d_c,
"static_friction": s_c,
"bounciness": b_c
}, {
"$type": "apply_force_to_object",
"force": {
"x": -200,
"y": 0,
"z": 0
},
"id": chair_id
}, {
"$type": "send_collisions",
"collision_types": collision_types,
"enter": True,
"exit": True,
"stay": True
}
])
self.communicate(init_commands)
# Iterate through 500 frames.
# Every frame, listen for collisions, and parse the output data.
for i in range(500):
resp = self.communicate([])
if len(resp) > 1:
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
# There was a collision between two objects.
if r_id == "coll":
collision = Collision(r)
print("Collision between two objects:")
print("\tEvent: " + collision.get_state())
print("\tCollider: " +
str(collision.get_collider_id()))
print("\tCollidee: " +
str(collision.get_collidee_id()))
print("\tRelative velocity: " +
str(collision.get_relative_velocity()))
print("\tContacts:")
for j in range(collision.get_num_contacts()):
print(
str(collision.get_contact_normal(j)) + "\t" +
str(collision.get_contact_point(j)))
# There was a collision between an object and the environment.
elif r_id == "enco":
collision = EnvironmentCollision(r)
print(
"Collision between an object and the environment:")
print("\tEvent: " + collision.get_state())
print("\tCollider: " + str(collision.get_object_id()))
print("\tContacts:")
for j in range(collision.get_num_contacts()):
print(
str(collision.get_contact_normal(j)) + "\t" +
str(collision.get_contact_point(j)))
else:
raise Exception(r_id)