PWM / PWM_code /policy_rl /act /constants.py
hchen657Naza's picture
Add files using upload-large-folder tool
23e1f91 verified
Raw History Blame Contribute Delete
9.95 kB
import pathlib
### Task parameters
DATA_DIR = '/data0/hongyi/act_data'
SIM_TASK_CONFIGS = {
'sim_open_drawer_real30_dream87_fh': { # 30 real + 9 (b1) + 45 (b2) + 33 (b3: 15/10/8) OSCAR fixed-horizon dreams, stateless
'dataset_dir': '/root/pwm_local/data0/act_data_mix/open_drawer_real30_dream87_fh',
'num_episodes': 117,
'episode_len': 200, # longest demo is 192 (dreams); pad to 200
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 25,
},
'sim_open_drawer_real30_dream54_fh': { # 30 real + 9 (batch1) + 45 (batch2) OSCAR fixed-horizon dreams, stateless
'dataset_dir': '/root/pwm_local/data0/act_data_mix/open_drawer_real30_dream54_fh',
'num_episodes': 84,
'episode_len': 200, # longest demo is 192 (dreams); pad to 200
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 25,
},
'sim_open_drawer_real30_dream9_fh': { # 30 real + 9 OSCAR fixed-horizon dreams (3/room, rooms 2/4/5), stateless
'dataset_dir': '/root/pwm_local/data0/act_data_mix/open_drawer_real30_dream9_fh',
'num_episodes': 39,
'episode_len': 200, # longest demo is 192 (dreams); pad to 200
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 25,
},
'sim_open_drawer_real30_dream8': { # 30 real + 8 OSCAR dreams (room2), stateless
'dataset_dir': '/root/pwm_local/data0/act_data_mix/open_drawer_real30_dream8',
'num_episodes': 38,
'episode_len': 200, # longest demo is 177; pad to 200
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 25,
},
'sim_open_drawer_100demos': {
'dataset_dir': DATA_DIR + '/open_drawer_100demos',
'num_episodes': 100,
'episode_len': 200, # longest demo is 177; pad to 200
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 50, # ~25% of episode_len, matches ALOHA's 100/400 ratio
},
'sim_push_box_100demos': {
'dataset_dir': DATA_DIR + '/push_box_100demos',
'num_episodes': 100,
'episode_len': 320, # longest demo is 307; pad to 320
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 25, # matches open_drawer_100demos_chunk25_30k recipe
},
'sim_push_box_50demos': {
# First 50 of the 100 push-box demos (num_episodes=N loads episode_0..N-1).
# Deliberately WEAKER BC base — to test whether WM RL has more headroom to
# improve when the BC starting point is lower than the 100-demo (80%) policy.
'dataset_dir': DATA_DIR + '/push_box_100demos',
'num_episodes': 50,
'episode_len': 320,
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 25,
},
'sim_push_box_100demos_table2': {
# Same demos (action/qpos identical) but RGB rendered at table_idx=2 — so the
# ACT is native to the table the Push-Box WM was built at (robot_new/Push-Box_halffail).
'dataset_dir': DATA_DIR + '/push_box_100demos_table2',
'num_episodes': 100,
'episode_len': 320,
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 25,
},
'sim_open_laptop_100demos': {
# Built directly from EgoVLA_SIM/Open-Laptop hdf5 (build_act_from_egovla.py).
# 38-d upper-body proprio/action (matches act/eval_act_in_sim.py).
'dataset_dir': DATA_DIR + '/open_laptop_100demos',
'num_episodes': 100,
'episode_len': 160, # longest demo 146; pad to 160
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 25,
},
'sim_pour_balls_100demos': {
'dataset_dir': DATA_DIR + '/pour_balls_100demos',
'num_episodes': 100,
'episode_len': 250, # longest demo 236; pad to 250
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 25,
},
'sim_open_drawer_30demos': {
# First 30 (sorted) EgoVLA_SIM/Open-Drawer demos via build_act_from_egovla.py.
'dataset_dir': DATA_DIR + '/open_drawer_30demos',
'num_episodes': 30,
'episode_len': 200, # longest demo 177; pad to 200
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 25,
},
'sim_open_laptop_30demos': {
'dataset_dir': DATA_DIR + '/open_laptop_30demos',
'num_episodes': 30,
'episode_len': 160, # longest demo 146; pad to 160
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 25,
},
'sim_pour_balls_30demos': {
'dataset_dir': DATA_DIR + '/pour_balls_30demos',
'num_episodes': 30,
'episode_len': 200, # longest demo 191; pad to 200
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 25,
},
'sim_push_box_30demos': {
'dataset_dir': DATA_DIR + '/push_box_30demos',
'num_episodes': 30,
'episode_len': 280, # longest demo 267; pad to 280
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 25,
},
'sim_multitask_30each': { # 4 tasks x 30 demos: open_drawer, open_laptop, pour_balls, push_box (120 eps)
'dataset_dir': DATA_DIR + '/multitask_30each',
'num_episodes': 120,
'episode_len': 280, # max len over the 4 (push_box 267); pad to 280
'camera_names': ['top'],
'state_dim': 38,
'action_dim': 38,
'chunk_size': 25,
},
'sim_transfer_cube_scripted':{
'dataset_dir': DATA_DIR + '/sim_transfer_cube_scripted',
'num_episodes': 50,
'episode_len': 400,
'camera_names': ['top']
},
'sim_transfer_cube_human':{
'dataset_dir': DATA_DIR + '/sim_transfer_cube_human',
'num_episodes': 50,
'episode_len': 400,
'camera_names': ['top']
},
'sim_insertion_scripted': {
'dataset_dir': DATA_DIR + '/sim_insertion_scripted',
'num_episodes': 50,
'episode_len': 400,
'camera_names': ['top']
},
'sim_insertion_human': {
'dataset_dir': DATA_DIR + '/sim_insertion_human',
'num_episodes': 50,
'episode_len': 500,
'camera_names': ['top']
},
}
### Simulation envs fixed constants
DT = 0.02
JOINT_NAMES = ["waist", "shoulder", "elbow", "forearm_roll", "wrist_angle", "wrist_rotate"]
START_ARM_POSE = [0, -0.96, 1.16, 0, -0.3, 0, 0.02239, -0.02239, 0, -0.96, 1.16, 0, -0.3, 0, 0.02239, -0.02239]
XML_DIR = str(pathlib.Path(__file__).parent.resolve()) + '/assets/' # note: absolute path
# Left finger position limits (qpos[7]), right_finger = -1 * left_finger
MASTER_GRIPPER_POSITION_OPEN = 0.02417
MASTER_GRIPPER_POSITION_CLOSE = 0.01244
PUPPET_GRIPPER_POSITION_OPEN = 0.05800
PUPPET_GRIPPER_POSITION_CLOSE = 0.01844
# Gripper joint limits (qpos[6])
MASTER_GRIPPER_JOINT_OPEN = 0.3083
MASTER_GRIPPER_JOINT_CLOSE = -0.6842
PUPPET_GRIPPER_JOINT_OPEN = 1.4910
PUPPET_GRIPPER_JOINT_CLOSE = -0.6213
############################ Helper functions ############################
MASTER_GRIPPER_POSITION_NORMALIZE_FN = lambda x: (x - MASTER_GRIPPER_POSITION_CLOSE) / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE)
PUPPET_GRIPPER_POSITION_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_POSITION_CLOSE) / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE)
MASTER_GRIPPER_POSITION_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) + MASTER_GRIPPER_POSITION_CLOSE
PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) + PUPPET_GRIPPER_POSITION_CLOSE
MASTER2PUPPET_POSITION_FN = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN(MASTER_GRIPPER_POSITION_NORMALIZE_FN(x))
MASTER_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE)
PUPPET_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE)
MASTER_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
MASTER2PUPPET_JOINT_FN = lambda x: PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN(MASTER_GRIPPER_JOINT_NORMALIZE_FN(x))
MASTER_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE)
PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE)
MASTER_POS2JOINT = lambda x: MASTER_GRIPPER_POSITION_NORMALIZE_FN(x) * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
MASTER_JOINT2POS = lambda x: MASTER_GRIPPER_POSITION_UNNORMALIZE_FN((x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE))
PUPPET_POS2JOINT = lambda x: PUPPET_GRIPPER_POSITION_NORMALIZE_FN(x) * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
PUPPET_JOINT2POS = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN((x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE))
MASTER_GRIPPER_JOINT_MID = (MASTER_GRIPPER_JOINT_OPEN + MASTER_GRIPPER_JOINT_CLOSE)/2