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]
