def __init__(self, cost_fn, data_dir='vismpc/sv2p_data_cloth', model_dir='vismpc/sv2p_model_cloth', horizon=5, batch=False, viz=None): """ cost_fn: higher-order function """ params = DotMap() params.name = 'cloth' params.model_dir = model_dir params.data_dir = data_dir params.popsize = 2000 # must match _run_cem's popsize! params.nparts = 1 params.plan_hor = horizon params.adim = 4 params.stochastic_model = True sys.argv = sys.argv[:1] self.model = SV2P(params) self.cost_fn = cost_fn # TUNE CEM VARIANCE: # -0.4/0.4 work better for smoothing, -0.7/0.7 better for folding self.ac_lb = np.array([-1., -1., -0.7, -0.7]) self.ac_ub = np.array([1., 1., 0.7, 0.7]) self.act_dim = 4 self.plan_hor = horizon self.batch = batch self.prev_sol = np.tile((self.ac_lb + self.ac_ub) / 2, [self.plan_hor]) # / 16 works better for smoothing, /8 better for folding self.init_var = np.tile(np.square(self.ac_ub - self.ac_lb) / 8, [self.plan_hor]) self.viz = viz
pp.add_argument("--model_dir", type=str, default="sv2p_model_cloth") pp.add_argument("--horizon", type=int, default=5) pp.add_argument("--input_img", type=str, help="filepath of input image as .pkl file") pp.add_argument("--batch", action="store_true", default=False) # input_img is pkl of images and actions args = pp.parse_args() params = DotMap() params.name = 'cloth' params.model_dir = args.model_dir params.data_dir = args.data_dir params.popsize = 1 params.nparts = 1 params.plan_hor = args.horizon params.adim = 4 params.stochastic_model = True sys.argv = sys.argv[:1] sv2p = SV2P(params) if args.batch: # pkl is a list of length num_episodes with entries of the form {obs: (steps+1, 56, 56, 3), acts: (steps, 4)} # output is of the same form with entries {preds: (steps+1-horizon, horizon, 56, 56, 3), acts: (steps+1-horizon, horizon, 4)} pkl = pickle.load(open(args.input_img, 'rb')) output = list() for episode in pkl: num_steps = len(episode['act']) all_acts = np.array(episode['act']) preds = [] acts = [] for i in range(num_steps + 1 - args.horizon): currim = episode['obs'][i]