Spaces:
Sleeping
Sleeping
File size: 8,963 Bytes
96da58e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 |
"""
Convenience script to tune a camera view in a mujoco environment.
Allows keyboard presses to move a camera around in the viewer, and
then prints the final position and quaternion you should set
for your camera in the mujoco XML file.
"""
import argparse
import time
import xml.etree.ElementTree as ET
import numpy as np
from pynput.keyboard import Controller, Key, Listener
import robosuite
import robosuite.utils.transform_utils as T
from robosuite.utils.camera_utils import CameraMover
from robosuite.utils.mjcf_utils import find_elements, find_parent
# some settings
DELTA_POS_KEY_PRESS = 0.05 # delta camera position per key press
DELTA_ROT_KEY_PRESS = 1 # delta camera angle per key press
class KeyboardHandler:
def __init__(self, camera_mover):
"""
Store internal state here.
Args:
camera_mover (CameraMover): Playback camera class
cam_body_id (int): id corresponding to parent body of camera element
"""
self.camera_mover = camera_mover
# make a thread to listen to keyboard and register our callback functions
self.listener = Listener(on_press=self.on_press, on_release=self.on_release)
# start listening
self.listener.start()
def on_press(self, key):
"""
Key handler for key presses.
Args:
key (int): keycode corresponding to the key that was pressed
"""
try:
# controls for moving rotation
if key == Key.up:
# rotate up
self.camera_mover.rotate_camera(point=None, axis=[1.0, 0.0, 0.0], angle=DELTA_ROT_KEY_PRESS)
elif key == Key.down:
# rotate down
self.camera_mover.rotate_camera(point=None, axis=[-1.0, 0.0, 0.0], angle=DELTA_ROT_KEY_PRESS)
elif key == Key.left:
# rotate left
self.camera_mover.rotate_camera(point=None, axis=[0.0, 1.0, 0.0], angle=DELTA_ROT_KEY_PRESS)
elif key == Key.right:
# rotate right
self.camera_mover.rotate_camera(point=None, axis=[0.0, -1.0, 0.0], angle=DELTA_ROT_KEY_PRESS)
# controls for moving position
elif key.char == "w":
# move forward
self.camera_mover.move_camera(direction=[0.0, 0.0, -1.0], scale=DELTA_POS_KEY_PRESS)
elif key.char == "s":
# move backward
self.camera_mover.move_camera(direction=[0.0, 0.0, 1.0], scale=DELTA_POS_KEY_PRESS)
elif key.char == "a":
# move left
self.camera_mover.move_camera(direction=[-1.0, 0.0, 0.0], scale=DELTA_POS_KEY_PRESS)
elif key.char == "d":
# move right
self.camera_mover.move_camera(direction=[1.0, 0.0, 0.0], scale=DELTA_POS_KEY_PRESS)
elif key.char == "r":
# move up
self.camera_mover.move_camera(direction=[0.0, 1.0, 0.0], scale=DELTA_POS_KEY_PRESS)
elif key.char == "f":
# move down
self.camera_mover.move_camera(direction=[0.0, -1.0, 0.0], scale=DELTA_POS_KEY_PRESS)
elif key.char == ".":
# rotate counterclockwise
self.camera_mover.rotate_camera(point=None, axis=[0.0, 0.0, 1.0], angle=DELTA_ROT_KEY_PRESS)
elif key.char == "/":
# rotate clockwise
self.camera_mover.rotate_camera(point=None, axis=[0.0, 0.0, -1.0], angle=DELTA_ROT_KEY_PRESS)
except AttributeError as e:
pass
def on_release(self, key):
"""
Key handler for key releases.
Args:
key: [NOT USED]
"""
pass
def print_command(char, info):
"""
Prints out the command + relevant info entered by user
Args:
char (str): Command entered
info (str): Any additional info to print
"""
char += " " * (10 - len(char))
print("{}\t{}".format(char, info))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--env", type=str, default="Lift")
parser.add_argument("--robots", nargs="+", type=str, default="Sawyer", help="Which robot(s) to use in the env")
args = parser.parse_args()
print("\nWelcome to the camera tuning script! You will be able to tune a camera view")
print("by moving it around using your keyboard. The controls are printed below.")
print("")
print_command("Keys", "Command")
print_command("w-s", "zoom the camera in/out")
print_command("a-d", "pan the camera left/right")
print_command("r-f", "pan the camera up/down")
print_command("arrow keys", "rotate the camera to change view direction")
print_command(".-/", "rotate the camera view without changing view direction")
print("")
# read camera XML tag from user input
inp = input(
"\nPlease paste a camera name below \n"
"OR xml tag below (e.g. <camera ... />) \n"
"OR leave blank for an example:\n"
)
if len(inp) == 0:
if args.env != "Lift":
raise Exception("ERROR: env must be Lift to run default example.")
print("\nUsing an example tag corresponding to the frontview camera.")
print("This xml tag was copied from robosuite/models/assets/arenas/table_arena.xml")
inp = '<camera mode="fixed" name="frontview" pos="1.6 0 1.45" quat="0.56 0.43 0.43 0.56"/>'
# remember the tag and infer some properties
from_tag = "<" in inp
notify_str = (
"NOTE: using the following xml tag:\n"
if from_tag
else "NOTE: using the following camera (initialized at default sim location)\n"
)
print(notify_str)
print("{}\n".format(inp))
cam_tree = ET.fromstring(inp) if from_tag else ET.Element("camera", attrib={"name": inp})
CAMERA_NAME = cam_tree.get("name")
# make the environment
env = robosuite.make(
args.env,
robots=args.robots,
has_renderer=True,
has_offscreen_renderer=False,
ignore_done=True,
use_camera_obs=False,
control_freq=100,
)
env.reset()
# Create the camera mover
camera_mover = CameraMover(
env=env,
camera=CAMERA_NAME,
)
# Make sure we're using the camera that we're modifying
camera_id = env.sim.model.camera_name2id(CAMERA_NAME)
env.viewer.set_camera(camera_id=camera_id)
# Infer initial camera pose
if from_tag:
initial_file_camera_pos = np.array(cam_tree.get("pos").split(" ")).astype(float)
initial_file_camera_quat = T.convert_quat(np.array(cam_tree.get("quat").split(" ")).astype(float), to="xyzw")
# Set these values as well
camera_mover.set_camera_pose(pos=initial_file_camera_pos, quat=initial_file_camera_quat)
# Optionally set fov if specified
cam_fov = cam_tree.get("fovy", None)
if cam_fov is not None:
env.sim.model.cam_fovy[camera_id] = float(cam_fov)
else:
initial_file_camera_pos, initial_file_camera_quat = camera_mover.get_camera_pose()
# Define initial file camera pose
initial_file_camera_pose = T.make_pose(initial_file_camera_pos, T.quat2mat(initial_file_camera_quat))
# remember difference between camera pose in initial tag and absolute camera pose in world
initial_world_camera_pos, initial_world_camera_quat = camera_mover.get_camera_pose()
initial_world_camera_pose = T.make_pose(initial_world_camera_pos, T.quat2mat(initial_world_camera_quat))
world_in_file = initial_file_camera_pose.dot(T.pose_inv(initial_world_camera_pose))
# register callbacks to handle key presses in the viewer
key_handler = KeyboardHandler(camera_mover=camera_mover)
# just spin to let user interact with window
spin_count = 0
while True:
action = np.zeros(env.action_dim)
obs, reward, done, _ = env.step(action)
env.render()
spin_count += 1
if spin_count % 500 == 0:
# convert from world coordinates to file coordinates (xml subtree)
camera_pos, camera_quat = camera_mover.get_camera_pose()
world_camera_pose = T.make_pose(camera_pos, T.quat2mat(camera_quat))
file_camera_pose = world_in_file.dot(world_camera_pose)
# TODO: Figure out why numba causes black screen of death (specifically, during mat2pose --> mat2quat call below)
camera_pos, camera_quat = T.mat2pose(file_camera_pose)
camera_quat = T.convert_quat(camera_quat, to="wxyz")
print("\n\ncurrent camera tag you should copy")
cam_tree.set("pos", "{} {} {}".format(camera_pos[0], camera_pos[1], camera_pos[2]))
cam_tree.set("quat", "{} {} {} {}".format(camera_quat[0], camera_quat[1], camera_quat[2], camera_quat[3]))
print(ET.tostring(cam_tree, encoding="utf8").decode("utf8"))
|