<reponame>Asap7772/rail-rl-franka-eval<gh_stars>0
"""
Compare a gradient-based MPC controller with a stochastic-optimization based
MPC controller.
"""
import argparse
import uuid
import numpy as np

import joblib

from railrl.core import logger
from railrl.torch.mpc.controller import GradientBasedMPCController

filename = str(uuid.uuid4())


def rollout(env, agent, orig_policy, max_path_length=np.inf, animated=False):
    """
    The following value for the following keys will be a 2D array, with the
    first dimension corresponding to the time dimension.
     - observations
     - actions
     - rewards
     - next_observations
     - terminals

    The next two elements will be lists of dictionaries, with the index into
    the list being the index into the time
     - agent_infos
     - env_infos

    :param env:
    :param agent:
    :param max_path_length:
    :param animated:
    :return:
    """
    observations = []
    actions = []
    rewards = []
    terminals = []
    agent_infos = []
    env_infos = []
    o = env.reset()
    next_o = None
    path_length = 0
    if animated:
        env.render()
    while path_length < max_path_length:
        a, agent_info = agent.get_action(o)
        next_o, r, d, env_info = env.step(a)
        # print("gradient", a)
        # print("orig", orig_policy.get_action(o)[0])
        observations.append(o)
        rewards.append(r)
        terminals.append(d)
        actions.append(a)
        agent_infos.append(agent_info)
        env_infos.append(env_info)
        path_length += 1
        if d:
            break
        o = next_o
        if animated:
            env.render()

    actions = np.array(actions)
    if len(actions.shape) == 1:
        actions = np.expand_dims(actions, 1)
    observations = np.array(observations)
    if len(observations.shape) == 1:
        observations = np.expand_dims(observations, 1)
        next_o = np.array([next_o])
    next_observations = np.vstack(
        (
            observations[1:, :],
            np.expand_dims(next_o, 0)
        )
    )
    return dict(
        observations=observations,
        actions=actions,
        rewards=np.array(rewards).reshape(-1, 1),
        next_observations=next_observations,
        terminals=np.array(terminals).reshape(-1, 1),
        agent_infos=agent_infos,
        env_infos=env_infos,
    )


def simulate_policy(args):
    data = joblib.load(args.file)
    model = data['model']
    env = data['env']
    orig_policy = data['mpc_controller']
    print("Policy loaded")
    if args.pause:
        import ipdb; ipdb.set_trace()
    policy = GradientBasedMPCController(
        env,
        model,
        mpc_horizon=1,
        num_grad_steps=10,
        learning_rate=1e-1,
        warm_start=False,
    )
    while True:
        path = rollout(
            env,
            policy,
            orig_policy,
            max_path_length=args.H,
            animated=True,
        )
        if hasattr(env, "log_diagnostics"):
            env.log_diagnostics([path])
        logger.dump_tabular()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('file', type=str,
                        help='path to the snapshot file')
    parser.add_argument('--H', type=int, default=300,
                        help='Max length of rollout')
    parser.add_argument('--speedup', type=float, default=10,
                        help='Speedup')
    parser.add_argument('--gpu', action='store_true')
    parser.add_argument('--pause', action='store_true')
    args = parser.parse_args()

    simulate_policy(args)

"""
Compare a gradient-based MPC controller with a stochastic-optimization based
MPC controller.
"""
import argparse
import uuid
import numpy as np

import joblib

from railrl.core import logger
from railrl.torch.mpc.controller import GradientBasedMPCController

filename = str(uuid.uuid4())


def rollout(env, agent, orig_policy, max_path_length=np.inf, animated=False):
    """
    The following value for the following keys will be a 2D array, with the
    first dimension corresponding to the time dimension.
     - observations
     - actions
     - rewards
     - next_observations
     - terminals

    The next two elements will be lists of dictionaries, with the index into
    the list being the index into the time
     - agent_infos
     - env_infos

    :param env:
    :param agent:
    :param max_path_length:
    :param animated:
    :return:
    """
    observations = []
    actions = []
    rewards = []
    terminals = []
    agent_infos = []
    env_infos = []
    o = env.reset()
    next_o = None
    path_length = 0
    if animated:
        env.render()
    while path_length < max_path_length:
        a, agent_info = agent.get_action(o)
        next_o, r, d, env_info = env.step(a)
        # print("gradient", a)
        # print("orig", orig_policy.get_action(o)[0])
        observations.append(o)
        rewards.append(r)
        terminals.append(d)
        actions.append(a)
        agent_infos.append(agent_info)
        env_infos.append(env_info)
        path_length += 1
        if d:
            break
        o = next_o
        if animated:
            env.render()

    actions = np.array(actions)
    if len(actions.shape) == 1:
        actions = np.expand_dims(actions, 1)
    observations = np.array(observations)
    if len(observations.shape) == 1:
        observations = np.expand_dims(observations, 1)
        next_o = np.array([next_o])
    next_observations = np.vstack(
        (
            observations[1:, :],
            np.expand_dims(next_o, 0)
        )
    )
    return dict(
        observations=observations,
        actions=actions,
        rewards=np.array(rewards).reshape(-1, 1),
        next_observations=next_observations,
        terminals=np.array(terminals).reshape(-1, 1),
        agent_infos=agent_infos,
        env_infos=env_infos,
    )


def simulate_policy(args):
    data = joblib.load(args.file)
    model = data['model']
    env = data['env']
    orig_policy = data['mpc_controller']
    print("Policy loaded")
    if args.pause:
        import ipdb; ipdb.set_trace()
    policy = GradientBasedMPCController(
        env,
        model,
        mpc_horizon=1,
        num_grad_steps=10,
        learning_rate=1e-1,
        warm_start=False,
    )
    while True:
        path = rollout(
            env,
            policy,
            orig_policy,
            max_path_length=args.H,
            animated=True,
        )
        if hasattr(env, "log_diagnostics"):
            env.log_diagnostics([path])
        logger.dump_tabular()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('file', type=str,
                        help='path to the snapshot file')
    parser.add_argument('--H', type=int, default=300,
                        help='Max length of rollout')
    parser.add_argument('--speedup', type=float, default=10,
                        help='Speedup')
    parser.add_argument('--gpu', action='store_true')
    parser.add_argument('--pause', action='store_true')
    args = parser.parse_args()

    simulate_policy(args)
Compare a gradient-based MPC controller with a stochastic-optimization based MPC controller. The following value for the following keys will be a 2D array, with the     first dimension corresponding to the time dimension.      - observations      - actions      - rewards      - next_observations      - terminals      The next two elements will be lists of dictionaries, with the index into     the list being the index into the time      - agent_infos      - env_infos      :param env:     :param agent:     :param max_path_length:     :param animated:     :return: # print("gradient", a) # print("orig", orig_policy.get_action(o)[0])