# Path: pytorchrl/misc/tensor_utils.py
# def running_average_tensor_list(first_list, second_list, rate):
#         """
#         Return the result of
#         first_list * (1 - rate) + second_list * rate
# 
#         Parameter
#         ---------
#         first_list (list) : A list of pytorch Tensors
#         second_list (list) : A list of pytorch Tensors, should have the same
#             format as first_list.
#         rate (float): a learning rate, in [0, 1]
# 
#         Returns
#         -------
#         results (list): A list of Tensors with computed results.
#         """
#         results = []
#         assert len(first_list) == len(second_list)
#         for first_t, second_t in zip(first_list, second_list):
#             assert first_t.shape == second_t.shape
#             result_tensor = first_t * (1 - rate) + second_t * rate
#             results.append(result_tensor)
# 
#         return results
# 
# Path: pytorchrl/sampler/parallel_sampler.py
# def _worker_init(G, id):
# def initialize(n_parallel):
# def _get_scoped_G(G, scope):
# def _worker_populate_task(G, env, policy, scope=None):
# def _worker_terminate_task(G, scope=None):
# def populate_task(env, policy, scope=None):
# def terminate_task(scope=None):
# def _worker_set_seed(_, seed):
# def set_seed(seed):
# def _worker_set_policy_params(G, params, scope=None):
# def _worker_set_env_params(G,params,scope=None):
# def _worker_collect_one_path(G, max_path_length, scope=None):
# def sample_paths(
#         policy_params,
#         max_samples,
#         max_path_length=np.inf,
#         env_params=None,
#         scope=None):
# def truncate_paths(paths, max_samples):
#     G = _get_scoped_G(G, scope)
#     G = _get_scoped_G(G, scope)
#         G = _get_scoped_G(singleton_pool.G, scope)
#     G = _get_scoped_G(G, scope)
#     G = _get_scoped_G(G, scope)
#     G = _get_scoped_G(G, scope)
# 
# Path: pytorchrl/sampler/utils.py
# def rollout(env, agent, max_path_length=np.inf, animated=False, speedup=1,
#             always_return_paths=False):
#     observations = []
#     actions = []
#     rewards = []
#     agent_infos = []
#     env_infos = []
#     o = env.reset()
#     agent.reset()
#     path_length = 0
#     if animated:
#         env.render()
#     while path_length < max_path_length:
#         a, agent_info = agent.get_action(o)
#         # print(a)
#         next_o, r, d, env_info = env.step(a)
#         observations.append(o)
#         rewards.append(r)
#         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()
#             timestep = 0.05
#             time.sleep(timestep / speedup)
#     if animated and not always_return_paths:
#         return
# 
#     return dict(
#         observations=tensor_utils.stack_tensor_list(observations),
#         actions=tensor_utils.stack_tensor_list(actions),
#         rewards=tensor_utils.stack_tensor_list(rewards),
#         agent_infos=tensor_utils.stack_tensor_dict_list(agent_infos),
#         env_infos=tensor_utils.stack_tensor_dict_list(env_infos),
#     )


import copy
import numpy as np
import pyprind
import rllab.misc.logger as logger
import torch
import torch.nn as nn
from rllab.algos.base import RLAlgorithm
from rllab.algos.ddpg import SimpleReplayPool
from rllab.misc import special
from rllab.misc import ext
from torch.autograd import Variable
from torch import optim
from pytorchrl.misc.tensor_utils import running_average_tensor_list
from pytorchrl.sampler import parallel_sampler
from pytorchrl.sampler.utils import rollout
        q_vals = self.qf(obs, actions)

        # Define loss function
        loss_fn = nn.MSELoss()
        loss = loss_fn(q_vals, expected_q)

        # Backpropagation and gradient descent
        self.qf_optimizer.zero_grad()
        loss.backward()
        self.qf_optimizer.step()

        return loss.data.numpy()

    def train_policy(self, obs_val):
        """
        Given the mini-batch, do gradient ascent on policy
        """
        obs = Variable(torch.from_numpy(obs_val)).type(torch.FloatTensor)

        # Do gradient descent, so need to add minus in front
        average_q = -self.qf(obs, self.policy(obs)).mean()

        self.policy_optimizer.zero_grad()
        average_q.backward()
        self.policy_optimizer.step()

        return average_q.data.numpy()

    def evaluate(self, epoch, pool):
        logger.log("Collecting samples for evaluation")

# Path: pytorchrl/misc/tensor_utils.py
# def running_average_tensor_list(first_list, second_list, rate):
#         """
#         Return the result of
#         first_list * (1 - rate) + second_list * rate
# 
#         Parameter
#         ---------
#         first_list (list) : A list of pytorch Tensors
#         second_list (list) : A list of pytorch Tensors, should have the same
#             format as first_list.
#         rate (float): a learning rate, in [0, 1]
# 
#         Returns
#         -------
#         results (list): A list of Tensors with computed results.
#         """
#         results = []
#         assert len(first_list) == len(second_list)
#         for first_t, second_t in zip(first_list, second_list):
#             assert first_t.shape == second_t.shape
#             result_tensor = first_t * (1 - rate) + second_t * rate
#             results.append(result_tensor)
# 
#         return results
# 
# Path: pytorchrl/sampler/parallel_sampler.py
# def _worker_init(G, id):
# def initialize(n_parallel):
# def _get_scoped_G(G, scope):
# def _worker_populate_task(G, env, policy, scope=None):
# def _worker_terminate_task(G, scope=None):
# def populate_task(env, policy, scope=None):
# def terminate_task(scope=None):
# def _worker_set_seed(_, seed):
# def set_seed(seed):
# def _worker_set_policy_params(G, params, scope=None):
# def _worker_set_env_params(G,params,scope=None):
# def _worker_collect_one_path(G, max_path_length, scope=None):
# def sample_paths(
#         policy_params,
#         max_samples,
#         max_path_length=np.inf,
#         env_params=None,
#         scope=None):
# def truncate_paths(paths, max_samples):
#     G = _get_scoped_G(G, scope)
#     G = _get_scoped_G(G, scope)
#         G = _get_scoped_G(singleton_pool.G, scope)
#     G = _get_scoped_G(G, scope)
#     G = _get_scoped_G(G, scope)
#     G = _get_scoped_G(G, scope)
# 
# Path: pytorchrl/sampler/utils.py
# def rollout(env, agent, max_path_length=np.inf, animated=False, speedup=1,
#             always_return_paths=False):
#     observations = []
#     actions = []
#     rewards = []
#     agent_infos = []
#     env_infos = []
#     o = env.reset()
#     agent.reset()
#     path_length = 0
#     if animated:
#         env.render()
#     while path_length < max_path_length:
#         a, agent_info = agent.get_action(o)
#         # print(a)
#         next_o, r, d, env_info = env.step(a)
#         observations.append(o)
#         rewards.append(r)
#         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()
#             timestep = 0.05
#             time.sleep(timestep / speedup)
#     if animated and not always_return_paths:
#         return
# 
#     return dict(
#         observations=tensor_utils.stack_tensor_list(observations),
#         actions=tensor_utils.stack_tensor_list(actions),
#         rewards=tensor_utils.stack_tensor_list(rewards),
#         agent_infos=tensor_utils.stack_tensor_dict_list(agent_infos),
#         env_infos=tensor_utils.stack_tensor_dict_list(env_infos),
#     )


# Path: pytorchrl/algos/ddpg.py
import copy
import numpy as np
import pyprind
import rllab.misc.logger as logger
import torch
import torch.nn as nn
from rllab.algos.base import RLAlgorithm
from rllab.algos.ddpg import SimpleReplayPool
from rllab.misc import special
from rllab.misc import ext
from torch.autograd import Variable
from torch import optim
from pytorchrl.misc.tensor_utils import running_average_tensor_list
from pytorchrl.sampler import parallel_sampler
from pytorchrl.sampler.utils import rollout

        q_vals = self.qf(obs, actions)

        # Define loss function
        loss_fn = nn.MSELoss()
        loss = loss_fn(q_vals, expected_q)

        # Backpropagation and gradient descent
        self.qf_optimizer.zero_grad()
        loss.backward()
        self.qf_optimizer.step()

        return loss.data.numpy()

    def train_policy(self, obs_val):
        """
        Given the mini-batch, do gradient ascent on policy
        """
        obs = Variable(torch.from_numpy(obs_val)).type(torch.FloatTensor)

        # Do gradient descent, so need to add minus in front
        average_q = -self.qf(obs, self.policy(obs)).mean()

        self.policy_optimizer.zero_grad()
        average_q.backward()
        self.policy_optimizer.step()

        return average_q.data.numpy()

    def evaluate(self, epoch, pool):
        logger.log("Collecting samples for evaluation")

        paths = parallel_sampler.sample_paths(