def test_agent_collect_info():
vec_env, custom_draws = make_block_push_env(two_d=True)
red_plan = red_skeleton()
green_plan = green_skeleton()
blue_plan = blue_skeleton()
#plans = [blue_plan, green_plan, red_plan]
plans = [red_plan]
agent = Agent("test_collect_info")
agent.collect_transition_data(vec_env, plans)
def test_one_push_planner():
vec_env, custom_draws = make_block_push_env(two_d=True)
start_state_str = vec_env.get_pillar_state()[0]
start_state = State.create_from_serialized_string(start_state_str)
goal_state = State.create_from_serialized_string(start_state_str)
goal_pose = np.array(start_state.get_values_as_vec([block_pos_fqn]))
goal_pose[0] -= 0.05
goal_state.set_values_from_vec([block_pos_fqn], goal_pose.tolist())
planner = Planner(vec_env.cfg)
one_push_plan = planner.plan(start_state.get_serialized_string(),
goal_state.get_serialized_string())
plans = [one_push_plan]
agent = Agent("test_planner")
agent.collect_transition_data(vec_env, plans)
