File size: 5,014 Bytes
e479c46 | 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 | import argparse
import numpy as np
from sapien.core import Pose
from transforms3d.euler import euler2quat
from simpler_env.utils.io import DictAction
def parse_range_tuple(t):
return np.linspace(t[0], t[1], int(t[2]))
def get_args():
# parse command-line arguments
parser = argparse.ArgumentParser()
parser.add_argument(
"--policy-model",
type=str,
default="rt1",
help="Policy model type; e.g., 'rt1', 'octo-base', 'octo-small'",
)
parser.add_argument(
"--policy-setup",
type=str,
default="google_robot",
help="Policy model setup; e.g., 'google_robot', 'widowx_bridge'",
)
parser.add_argument("--ckpt-path", type=str, default=None)
parser.add_argument("--env-name", type=str, required=True)
parser.add_argument(
"--additional-env-save-tags",
type=str,
default=None,
help="Additional tags to save the environment eval results",
)
parser.add_argument("--scene-name", type=str, default="google_pick_coke_can_1_v4")
parser.add_argument("--enable-raytracing", action="store_true")
parser.add_argument("--robot", type=str, default="google_robot_static")
parser.add_argument(
"--obs-camera-name",
type=str,
default=None,
help="Obtain image observation from this camera for policy input. None = default",
)
parser.add_argument("--action-scale", type=float, default=1.0)
parser.add_argument("--control-freq", type=int, default=3)
parser.add_argument("--sim-freq", type=int, default=513)
parser.add_argument("--max-episode-steps", type=int, default=80)
parser.add_argument("--rgb-overlay-path", type=str, default=None)
parser.add_argument(
"--robot-init-x-range",
type=float,
nargs=3,
default=[0.35, 0.35, 1],
help="[xmin, xmax, num]",
)
parser.add_argument(
"--robot-init-y-range",
type=float,
nargs=3,
default=[0.20, 0.20, 1],
help="[ymin, ymax, num]",
)
parser.add_argument(
"--robot-init-rot-quat-center",
type=float,
nargs=4,
default=[1, 0, 0, 0],
help="[x, y, z, w]",
)
parser.add_argument(
"--robot-init-rot-rpy-range",
type=float,
nargs=9,
default=[0, 0, 1, 0, 0, 1, 0, 0, 1],
help="[rmin, rmax, rnum, pmin, pmax, pnum, ymin, ymax, ynum]",
)
parser.add_argument(
"--obj-variation-mode",
type=str,
default="xy",
choices=["xy", "episode"],
help="Whether to vary the xy position of a single object, or to vary predetermined episodes",
)
parser.add_argument("--obj-episode-range", type=int, nargs=2, default=[0, 60], help="[start, end]")
parser.add_argument(
"--obj-init-x-range",
type=float,
nargs=3,
default=[-0.35, -0.12, 5],
help="[xmin, xmax, num]",
)
parser.add_argument(
"--obj-init-y-range",
type=float,
nargs=3,
default=[-0.02, 0.42, 5],
help="[ymin, ymax, num]",
)
parser.add_argument(
"--additional-env-build-kwargs",
nargs="+",
action=DictAction,
help="Additional env build kwargs in xxx=yyy format. If the value "
'is a list, it should be like key="[a,b]" or key=a,b '
'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
"Note that the quotation marks are necessary and that no white space "
"is allowed.",
)
parser.add_argument("--logging-dir", type=str, default="./results")
parser.add_argument("--tf-memory-limit", type=int, default=3072, help="Tensorflow memory limit")
parser.add_argument("--octo-init-rng", type=int, default=0, help="Octo init rng seed")
args = parser.parse_args()
# env args: robot pose
args.robot_init_xs = parse_range_tuple(args.robot_init_x_range)
args.robot_init_ys = parse_range_tuple(args.robot_init_y_range)
args.robot_init_quats = []
for r in parse_range_tuple(args.robot_init_rot_rpy_range[:3]):
for p in parse_range_tuple(args.robot_init_rot_rpy_range[3:6]):
for y in parse_range_tuple(args.robot_init_rot_rpy_range[6:]):
args.robot_init_quats.append((Pose(q=euler2quat(r, p, y)) * Pose(q=args.robot_init_rot_quat_center)).q)
# env args: object position
if args.obj_variation_mode == "xy":
args.obj_init_xs = parse_range_tuple(args.obj_init_x_range)
args.obj_init_ys = parse_range_tuple(args.obj_init_y_range)
# update logging info (args.additional_env_save_tags) if using a different camera from default
if args.obs_camera_name is not None:
if args.additional_env_save_tags is None:
args.additional_env_save_tags = f"obs_camera_{args.obs_camera_name}"
else:
args.additional_env_save_tags = args.additional_env_save_tags + f"_obs_camera_{args.obs_camera_name}"
return args
|