<commit_before>from __future__ import absolute_import
from __future__ import division
import numpy as np
import os
import visdom
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from utils.helpers import loggerConfig
from utils.sharedAdam import SharedAdam

CONFIGS = [
# agent_type, env_type,    game,                       model_type, memory_type
[ "empty",    "gym",       "CartPole-v0",              "mlp",      "none"        ],  # 0
[ "dqn",      "gym",       "CartPole-v0",              "mlp",      "sequential"  ],  # 1
[ "dqn",      "atari-ram", "Pong-ram-v0",              "mlp",      "sequential"  ],  # 2
[ "dqn",      "atari",     "PongDeterministic-v3",     "cnn",      "sequential"  ],  # 3
[ "dqn",      "atari",     "BreakoutDeterministic-v3", "cnn",      "sequential"  ],  # 4
[ "a3c",      "atari",     "PongDeterministic-v3",     "a3c-cnn",  "none"        ],  # 5
[ "a3c",      "gym",       "InvertedPendulum-v1",      "a3c-mjc",  "none"        ],  # 6
[ "a3c",      "minisim",   "minisim",                  "a3c-mlp",  "none"        ],  # 7
[ "a3c",      "minisim",   "minisim",                  "a3c-mlp-minisim",  "none"],  # 8
[ "a3c",      "minisim",   "minisim",                  "a3c-mlp-minisim-narrowing",  "none"],  # 9
[ "dqn",      "minisim",   "minisim",                  "mlp",      "sequential"  ]   # 10
]

minisim_num_robots = 1

class Params(object):   # NOTE: shared across all modules
    def __init__(self):
        self.verbose     = 2            # 0(warning) | 1(info) | 2(debug)

        # training signature
        self.machine     = "pearl10"     # "machine_id"
        self.timestamp   = "17080600"   # "yymmdd##"
        # training configuration
        self.mode        = 1            # 1(train) | 2(test model_file)
        self.config      = 8

        self.seed        = 123
        self.render      = False        # whether render the window from the original envs or not
        self.visualize   = True         # whether do online plotting and stuff or not
        self.save_best   = False        # save model w/ highest reward if True, otherwise always save the latest model

        self.agent_type, self.env_type, self.game, self.model_type, self.memory_type = CONFIGS[self.config]

        if self.agent_type == "dqn":
            self.enable_double_dqn  = False
            self.enable_dueling     = False
            self.dueling_type       = 'avg' # avg | max | naive

            if self.env_type == "gym":
                self.hist_len       = 1
                self.hidden_dim     = 16
            elif self.env_type == "minisim":
                self.hist_len = 4
                self.hidden_dim = 24
            else:
                self.hist_len       = 4
                self.hidden_dim     = 256

            self.use_cuda           = torch.cuda.is_available()
            self.dtype              = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
        elif self.agent_type == "a3c":
            self.enable_lstm        = True
            if self.model_type == "a3c-mjc":    # NOTE: should be set to True when training Mujoco envs
                self.enable_continuous  = True
            else:
                self.enable_continuous  = False
            self.num_processes      = 8

            self.hist_len           = 1
            self.hidden_dim         = 128

            if self.env_type == "minisim":
                self.hist_len = 4
                self.hidden_dim = 160
                self.num_processes = 6

            self.use_cuda           = False
            self.dtype              = torch.FloatTensor
        else:
            self.hist_len           = 1
            self.hidden_dim         = 256

            self.use_cuda           = torch.cuda.is_available()
            self.dtype              = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor

        # prefix for model/log/visdom
        self.refs        = self.machine + "_" + self.timestamp # NOTE: using this as env for visdom
        self.root_dir    = os.getcwd()

        # model files
        # NOTE: will save the current model to model_name
        self.model_name  = self.root_dir + "/models/" + self.refs + ".pth"
        # NOTE: will load pretrained model_file if not None
        self.model_file  = None#self.root_dir + "/models/{TODO:FILL_IN_PRETAINED_MODEL_FILE}.pth"
        if self.mode == 2:
            self.model_file  = self.model_name  # NOTE: so only need to change self.mode to 2 to test the current training
            assert self.model_file is not None, "Pre-Trained model is None, Testing aborted!!!"
            self.refs = self.refs + "_test"     # NOTE: using this as env for visdom for testing, to avoid accidentally redraw on the training plots

        # logging configs
        self.log_name    = self.root_dir + "/logs/" + self.refs + ".log"
        self.logger      = loggerConfig(self.log_name, self.verbose)
        self.logger.warning("<===================================>")

        if self.visualize:
            self.vis = visdom.Visdom()
            self.logger.warning("bash$: python -m visdom.server")           # activate visdom server on bash
            self.logger.warning("http://localhost:8097/env/" + self.refs)   # open this address on browser

class EnvParams(Params):    # settings for simulation environment
    def __init__(self):
        super(EnvParams, self).__init__()

        if self.env_type == "gym":
            pass
        elif self.env_type == "atari-ram":
            pass
        elif self.env_type == "atari":
            self.hei_state = 42
            self.wid_state = 42
            self.preprocess_mode = 3    # 0(nothing) | 1(rgb2gray) | 2(rgb2y) | 3(crop&resize)
        elif self.env_type == "lab":
            pass
        elif self.env_type == "gazebo":
            self.hei_state = 60
            self.wid_state = 80
            self.preprocess_mode = 3  # 0(nothing) | 1(rgb2gray) | 2(rgb2y) | 3(crop&resize depth)
            self.img_encoding_type = "passthrough"
        elif self.env_type == "minisim":
            self.num_robots = minisim_num_robots
            self.curriculum = True
            pass
        else:
            assert False, "env_type must be: gym | atari-ram | atari | lab | minisim"

class ModelParams(Params):  # settings for network architecture
    def __init__(self):
        super(ModelParams, self).__init__()

        self.state_shape = None # NOTE: set in fit_model of inherited Agents
        self.action_dim  = None # NOTE: set in fit_model of inherited Agents

        if self.env_type == "minisim":
            self.num_robots = minisim_num_robots

class MemoryParams(Params):     # settings for replay memory
    def __init__(self):
        super(MemoryParams, self).__init__()

        if self.env_type == "gym":
            self.memory_size = 50000
        elif self.env_type == "atari-ram":
            self.memory_size = 1000000
        elif self.env_type == "atari":
            self.memory_size = 1000000
        elif self.env_type == "lab":
            self.memory_size = 1000000
        elif self.env_type == "minisim":
            self.memory_size = 200000
        else:
            self.memory_size = 1000000

class AgentParams(Params):  # hyperparameters for drl agents
    def __init__(self):
        super(AgentParams, self).__init__()

        # criteria and optimizer
        if self.agent_type == "dqn":
            self.value_criteria = F.smooth_l1_loss
            self.optim          = optim.Adam
            # self.optim          = optim.RMSprop
        elif self.agent_type == "a3c":
            self.value_criteria = nn.MSELoss()
            self.optim          = SharedAdam    # share momentum across learners
        else:
            self.value_criteria = F.smooth_l1_loss
            self.optim          = optim.Adam
        # hyperparameters
        if self.agent_type == "dqn" and self.env_type == "gym":
            self.steps               = 100000   # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 1.#np.inf
            self.lr                  = 0.0001
            self.eval_freq           = 2500     # NOTE: here means every this many steps
            self.eval_steps          = 1000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 1

            self.learn_start         = 500      # start update params after this many steps
            self.batch_size          = 32
            self.valid_size          = 250
            self.eps_start           = 1
            self.eps_end             = 0.3
            self.eps_eval            = 0.#0.05
            self.eps_decay           = 50000
            self.target_model_update = 1000#0.0001
            self.action_repetition   = 1
            self.memory_interval     = 1
            self.train_interval      = 1
        elif self.agent_type == "dqn" and self.env_type == "atari-ram" or \
             self.agent_type == "dqn" and self.env_type == "atari":
            self.steps               = 50000000 # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 40.#np.inf
            self.lr                  = 0.00025
            self.eval_freq           = 250000#12500    # NOTE: here means every this many steps
            self.eval_steps          = 125000#2500
            self.prog_freq           = 10000#self.eval_freq
            self.test_nepisodes      = 1

            self.learn_start         = 50000    # start update params after this many steps
            self.batch_size          = 32
            self.valid_size          = 500
            self.eps_start           = 1
            self.eps_end             = 0.1
            self.eps_eval            = 0.#0.05
            self.eps_decay           = 1000000
            self.target_model_update = 10000
            self.action_repetition   = 4
            self.memory_interval     = 1
            self.train_interval      = 4
        elif self.agent_type == "dqn" and self.env_type == "minisim":
            self.steps               = 1000000   # max #iterations
            self.early_stop          = 15000     # max #steps per episode  # should be the average it takes to get to the goal
            self.gamma               = 0.99
            self.clip_grad           = np.inf
            self.lr                  = 0.0001
            self.eval_freq           = 150000    # NOTE: here means every this many steps
            self.eval_steps          = 15000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 10

            self.learn_start         = 500      # start update params after this many steps
            self.batch_size          = 64
            self.valid_size          = 250
            self.eps_start           = 1
            self.eps_end             = 0.2
            self.eps_eval            = 0.#0.05
            self.eps_decay           = 700000
            self.target_model_update = 2000#0.0001
            self.action_repetition   = 1
            self.memory_interval     = 1
            self.train_interval      = 32
        elif self.agent_type == "a3c" and self.env_type == "atari-ram" or \
             self.agent_type == "a3c" and (self.env_type == "atari" or self.env_type == "gym"):
            self.steps               = 20000000 # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 40.
            self.lr                  = 0.0001
            self.eval_freq           = 60       # NOTE: here means every this many seconds
            self.eval_steps          = 3000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 10

            self.rollout_steps       = 20       # max look-ahead steps in a single rollout
            self.tau                 = 1.
        elif self.agent_type == "a3c" and self.env_type == "minisim":
            self.steps               = 2000000   # max #iterations
            self.early_stop          = 26500     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 40.
            self.lr                  = 0.0001
            self.eval_freq           = 60       # NOTE: here means every this many seconds
            self.eval_steps          = 60000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 10

            self.rollout_steps       = 50       # max look-ahead steps in a single rollout
            self.tau                 = 1.
            self.entropy_weight      = 0.02
        else:
            self.steps               = 1000000  # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 1.#np.inf
            self.lr                  = 0.001
            self.eval_freq           = 2500     # NOTE: here means every this many steps
            self.eval_steps          = 1000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 10

            self.learn_start         = 25000    # start update params after this many steps
            self.batch_size          = 32
            self.valid_size          = 500
            self.eps_start           = 1
            self.eps_end             = 0.1
            self.eps_eval            = 0.#0.05
            self.eps_decay           = 50000
            self.target_model_update = 1000
            self.action_repetition   = 1
            self.memory_interval     = 1
            self.train_interval      = 4

            self.rollout_steps       = 20       # max look-ahead steps in a single rollout
            self.tau                 = 1.

        self.env_params    = EnvParams()
        self.model_params  = ModelParams()
        self.memory_params = MemoryParams()

class Options(Params):
    agent_params  = AgentParams()
<commit_msg>minor<commit_after>from __future__ import absolute_import
from __future__ import division
import numpy as np
import os
import visdom
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from utils.helpers import loggerConfig
from utils.sharedAdam import SharedAdam

CONFIGS = [
# agent_type, env_type,    game,                       model_type, memory_type
[ "empty",    "gym",       "CartPole-v0",              "mlp",      "none"        ],  # 0
[ "dqn",      "gym",       "CartPole-v0",              "mlp",      "sequential"  ],  # 1
[ "dqn",      "atari-ram", "Pong-ram-v0",              "mlp",      "sequential"  ],  # 2
[ "dqn",      "atari",     "PongDeterministic-v3",     "cnn",      "sequential"  ],  # 3
[ "dqn",      "atari",     "BreakoutDeterministic-v3", "cnn",      "sequential"  ],  # 4
[ "a3c",      "atari",     "PongDeterministic-v3",     "a3c-cnn",  "none"        ],  # 5
[ "a3c",      "gym",       "InvertedPendulum-v1",      "a3c-mjc",  "none"        ],  # 6
[ "a3c",      "minisim",   "minisim",                  "a3c-mlp",  "none"        ],  # 7
[ "a3c",      "minisim",   "minisim",                  "a3c-mlp-minisim",  "none"],  # 8
[ "a3c",      "minisim",   "minisim",                  "a3c-mlp-minisim-narrowing",  "none"],  # 9
[ "dqn",      "minisim",   "minisim",                  "mlp",      "sequential"  ]   # 10
]

minisim_num_robots = 1

class Params(object):   # NOTE: shared across all modules
    def __init__(self):
        self.verbose     = 2            # 0(warning) | 1(info) | 2(debug)

        # training signature
        self.machine     = "pearl3"     # "machine_id"
        self.timestamp   = "17080600"   # "yymmdd##"
        # training configuration
        self.mode        = 1            # 1(train) | 2(test model_file)
        self.config      = 9

        self.seed        = 123
        self.render      = False        # whether render the window from the original envs or not
        self.visualize   = True         # whether do online plotting and stuff or not
        self.save_best   = False        # save model w/ highest reward if True, otherwise always save the latest model

        self.agent_type, self.env_type, self.game, self.model_type, self.memory_type = CONFIGS[self.config]

        if self.agent_type == "dqn":
            self.enable_double_dqn  = False
            self.enable_dueling     = False
            self.dueling_type       = 'avg' # avg | max | naive

            if self.env_type == "gym":
                self.hist_len       = 1
                self.hidden_dim     = 16
            elif self.env_type == "minisim":
                self.hist_len = 4
                self.hidden_dim = 24
            else:
                self.hist_len       = 4
                self.hidden_dim     = 256

            self.use_cuda           = torch.cuda.is_available()
            self.dtype              = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
        elif self.agent_type == "a3c":
            self.enable_lstm        = True
            if self.model_type == "a3c-mjc":    # NOTE: should be set to True when training Mujoco envs
                self.enable_continuous  = True
            else:
                self.enable_continuous  = False
            self.num_processes      = 8

            self.hist_len           = 1
            self.hidden_dim         = 128

            if self.env_type == "minisim":
                self.hist_len = 4
                self.hidden_dim = 128
                self.num_processes = 6

            self.use_cuda           = False
            self.dtype              = torch.FloatTensor
        else:
            self.hist_len           = 1
            self.hidden_dim         = 256

            self.use_cuda           = torch.cuda.is_available()
            self.dtype              = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor

        # prefix for model/log/visdom
        self.refs        = self.machine + "_" + self.timestamp # NOTE: using this as env for visdom
        self.root_dir    = os.getcwd()

        # model files
        # NOTE: will save the current model to model_name
        self.model_name  = self.root_dir + "/models/" + self.refs + ".pth"
        # NOTE: will load pretrained model_file if not None
        self.model_file  = None#self.root_dir + "/models/{TODO:FILL_IN_PRETAINED_MODEL_FILE}.pth"
        if self.mode == 2:
            self.model_file  = self.model_name  # NOTE: so only need to change self.mode to 2 to test the current training
            assert self.model_file is not None, "Pre-Trained model is None, Testing aborted!!!"
            self.refs = self.refs + "_test"     # NOTE: using this as env for visdom for testing, to avoid accidentally redraw on the training plots

        # logging configs
        self.log_name    = self.root_dir + "/logs/" + self.refs + ".log"
        self.logger      = loggerConfig(self.log_name, self.verbose)
        self.logger.warning("<===================================>")

        if self.visualize:
            self.vis = visdom.Visdom()
            self.logger.warning("bash$: python -m visdom.server")           # activate visdom server on bash
            self.logger.warning("http://localhost:8097/env/" + self.refs)   # open this address on browser

class EnvParams(Params):    # settings for simulation environment
    def __init__(self):
        super(EnvParams, self).__init__()

        if self.env_type == "gym":
            pass
        elif self.env_type == "atari-ram":
            pass
        elif self.env_type == "atari":
            self.hei_state = 42
            self.wid_state = 42
            self.preprocess_mode = 3    # 0(nothing) | 1(rgb2gray) | 2(rgb2y) | 3(crop&resize)
        elif self.env_type == "lab":
            pass
        elif self.env_type == "gazebo":
            self.hei_state = 60
            self.wid_state = 80
            self.preprocess_mode = 3  # 0(nothing) | 1(rgb2gray) | 2(rgb2y) | 3(crop&resize depth)
            self.img_encoding_type = "passthrough"
        elif self.env_type == "minisim":
            self.num_robots = minisim_num_robots
            self.curriculum = True
            pass
        else:
            assert False, "env_type must be: gym | atari-ram | atari | lab | minisim"

class ModelParams(Params):  # settings for network architecture
    def __init__(self):
        super(ModelParams, self).__init__()

        self.state_shape = None # NOTE: set in fit_model of inherited Agents
        self.action_dim  = None # NOTE: set in fit_model of inherited Agents

        if self.env_type == "minisim":
            self.num_robots = minisim_num_robots

class MemoryParams(Params):     # settings for replay memory
    def __init__(self):
        super(MemoryParams, self).__init__()

        if self.env_type == "gym":
            self.memory_size = 50000
        elif self.env_type == "atari-ram":
            self.memory_size = 1000000
        elif self.env_type == "atari":
            self.memory_size = 1000000
        elif self.env_type == "lab":
            self.memory_size = 1000000
        elif self.env_type == "minisim":
            self.memory_size = 200000
        else:
            self.memory_size = 1000000

class AgentParams(Params):  # hyperparameters for drl agents
    def __init__(self):
        super(AgentParams, self).__init__()

        # criteria and optimizer
        if self.agent_type == "dqn":
            self.value_criteria = F.smooth_l1_loss
            self.optim          = optim.Adam
            # self.optim          = optim.RMSprop
        elif self.agent_type == "a3c":
            self.value_criteria = nn.MSELoss()
            self.optim          = SharedAdam    # share momentum across learners
        else:
            self.value_criteria = F.smooth_l1_loss
            self.optim          = optim.Adam
        # hyperparameters
        if self.agent_type == "dqn" and self.env_type == "gym":
            self.steps               = 100000   # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 1.#np.inf
            self.lr                  = 0.0001
            self.eval_freq           = 2500     # NOTE: here means every this many steps
            self.eval_steps          = 1000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 1

            self.learn_start         = 500      # start update params after this many steps
            self.batch_size          = 32
            self.valid_size          = 250
            self.eps_start           = 1
            self.eps_end             = 0.3
            self.eps_eval            = 0.#0.05
            self.eps_decay           = 50000
            self.target_model_update = 1000#0.0001
            self.action_repetition   = 1
            self.memory_interval     = 1
            self.train_interval      = 1
        elif self.agent_type == "dqn" and self.env_type == "atari-ram" or \
             self.agent_type == "dqn" and self.env_type == "atari":
            self.steps               = 50000000 # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 40.#np.inf
            self.lr                  = 0.00025
            self.eval_freq           = 250000#12500    # NOTE: here means every this many steps
            self.eval_steps          = 125000#2500
            self.prog_freq           = 10000#self.eval_freq
            self.test_nepisodes      = 1

            self.learn_start         = 50000    # start update params after this many steps
            self.batch_size          = 32
            self.valid_size          = 500
            self.eps_start           = 1
            self.eps_end             = 0.1
            self.eps_eval            = 0.#0.05
            self.eps_decay           = 1000000
            self.target_model_update = 10000
            self.action_repetition   = 4
            self.memory_interval     = 1
            self.train_interval      = 4
        elif self.agent_type == "dqn" and self.env_type == "minisim":
            self.steps               = 1000000   # max #iterations
            self.early_stop          = 15000     # max #steps per episode  # should be the average it takes to get to the goal
            self.gamma               = 0.99
            self.clip_grad           = np.inf
            self.lr                  = 0.0001
            self.eval_freq           = 150000    # NOTE: here means every this many steps
            self.eval_steps          = 15000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 10

            self.learn_start         = 500      # start update params after this many steps
            self.batch_size          = 64
            self.valid_size          = 250
            self.eps_start           = 1
            self.eps_end             = 0.2
            self.eps_eval            = 0.#0.05
            self.eps_decay           = 700000
            self.target_model_update = 2000#0.0001
            self.action_repetition   = 1
            self.memory_interval     = 1
            self.train_interval      = 32
        elif self.agent_type == "a3c" and self.env_type == "atari-ram" or \
             self.agent_type == "a3c" and (self.env_type == "atari" or self.env_type == "gym"):
            self.steps               = 20000000 # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 40.
            self.lr                  = 0.0001
            self.eval_freq           = 60       # NOTE: here means every this many seconds
            self.eval_steps          = 3000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 10

            self.rollout_steps       = 20       # max look-ahead steps in a single rollout
            self.tau                 = 1.
        elif self.agent_type == "a3c" and self.env_type == "minisim":
            self.steps               = 2000000   # max #iterations
            self.early_stop          = 26500     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 40.
            self.lr                  = 0.0001
            self.eval_freq           = 60       # NOTE: here means every this many seconds
            self.eval_steps          = 60000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 10

            self.rollout_steps       = 50       # max look-ahead steps in a single rollout
            self.tau                 = 1.
            self.entropy_weight      = 0.02
        else:
            self.steps               = 1000000  # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 1.#np.inf
            self.lr                  = 0.001
            self.eval_freq           = 2500     # NOTE: here means every this many steps
            self.eval_steps          = 1000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 10

            self.learn_start         = 25000    # start update params after this many steps
            self.batch_size          = 32
            self.valid_size          = 500
            self.eps_start           = 1
            self.eps_end             = 0.1
            self.eps_eval            = 0.#0.05
            self.eps_decay           = 50000
            self.target_model_update = 1000
            self.action_repetition   = 1
            self.memory_interval     = 1
            self.train_interval      = 4

            self.rollout_steps       = 20       # max look-ahead steps in a single rollout
            self.tau                 = 1.

        self.env_params    = EnvParams()
        self.model_params  = ModelParams()
        self.memory_params = MemoryParams()

class Options(Params):
    agent_params  = AgentParams()
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import os
import visdom
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from utils.helpers import loggerConfig
from optims.sharedAdam import SharedAdam
from optims.sharedRMSprop import SharedRMSprop

CONFIGS = [
# agent_type, env_type,    game,                       model_type,     memory_type
[ "empty",    "gym",       "CartPole-v0",              "empty",        "none"      ],  # 0
[ "dqn",      "gym",       "CartPole-v0",              "dqn-mlp",      "sequential"],  # 1
[ "dqn",      "atari-ram", "Pong-ram-v0",              "dqn-mlp",      "sequential"],  # 2
[ "dqn",      "atari",     "PongDeterministic-v4",     "dqn-cnn",      "sequential"],  # 3
[ "dqn",      "atari",     "BreakoutDeterministic-v4", "dqn-cnn",      "sequential"],  # 4
[ "a3c",      "atari",     "PongDeterministic-v4",     "a3c-cnn-dis",  "none"      ],  # 5
[ "a3c",      "gym",       "InvertedPendulum-v1",      "a3c-mlp-con",  "none"      ],  # 6
[ "acer",     "gym",       "CartPole-v1",              "acer-mlp-dis", "none"      ]   # 7  # NOTE: acer still under development, dont use this config
]

class Params(object):   # NOTE: shared across all modules
    def __init__(self):
        self.verbose     = 0            # 0(warning) | 1(info) | 2(debug)

        # training signature
        self.machine     = "daim"    # "machine_id"
        self.timestamp   = "17080900"   # "yymmdd##"
        # training configuration
        self.mode        = 1            # 1(train) | 2(test model_file)
        self.config      = 5

        self.seed        = 123
        self.render      = False        # whether render the window from the original envs or not
        self.visualize   = True         # whether do online plotting and stuff or not
        self.save_best   = False        # save model w/ highest reward if True, otherwise always save the latest model

        self.agent_type, self.env_type, self.game, self.model_type, self.memory_type = CONFIGS[self.config]

        if self.agent_type == "dqn":
            self.enable_double_dqn  = False
            self.enable_dueling     = False
            self.dueling_type       = 'avg' # avg | max | naive

            if self.env_type == "gym":
                self.hist_len       = 1
                self.hidden_dim     = 16
            else:
                self.hist_len       = 4
                self.hidden_dim     = 256

            self.use_cuda           = torch.cuda.is_available()
            self.dtype              = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
        elif self.agent_type == "a3c":
            self.enable_lstm        = True
            if "-con" in self.model_type:
                self.enable_continuous  = True
            else:
                self.enable_continuous  = False
            self.num_processes      = 8

            self.hist_len           = 1
            self.hidden_dim         = 128

            self.use_cuda           = False
            self.dtype              = torch.FloatTensor
        elif self.agent_type == "acer":
            self.enable_lstm        = True
            if "-con" in self.model_type:
                self.enable_continuous  = True
            else:
                self.enable_continuous  = False
            self.num_processes      = 8

            self.hist_len           = 1
            self.hidden_dim         = 32

            self.use_cuda           = False
            self.dtype              = torch.FloatTensor
        else:
            self.hist_len           = 1
            self.hidden_dim         = 256

            self.use_cuda           = torch.cuda.is_available()
            self.dtype              = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor

        # prefix for model/log/visdom
        self.refs        = self.machine + "_" + self.timestamp # NOTE: using this as env for visdom
        self.root_dir    = os.getcwd()

        # model files
        # NOTE: will save the current model to model_name
        self.model_name  = self.root_dir + "/models/" + self.refs + ".pth"
        # NOTE: will load pretrained model_file if not None
        self.model_file  = None#self.root_dir + "/models/{TODO:FILL_IN_PRETAINED_MODEL_FILE}.pth"
        if self.mode == 2:
            self.model_file  = self.model_name  # NOTE: so only need to change self.mode to 2 to test the current training
            assert self.model_file is not None, "Pre-Trained model is None, Testing aborted!!!"
            self.refs = self.refs + "_test"     # NOTE: using this as env for visdom for testing, to avoid accidentally redraw on the training plots

        # logging configs
        self.log_name    = self.root_dir + "/logs/" + self.refs + ".log"
        self.logger      = loggerConfig(self.log_name, self.verbose)
        self.logger.warning("<===================================>")

        if self.visualize:
            self.vis = visdom.Visdom()
            self.logger.warning("bash$: python -m visdom.server")           # activate visdom server on bash
            self.logger.warning("http://localhost:8097/env/" + self.refs)   # open this address on browser

class EnvParams(Params):    # settings for simulation environment
    def __init__(self):
        super(EnvParams, self).__init__()

        if self.env_type == "gym":
            pass
        elif self.env_type == "atari-ram":
            pass
        elif self.env_type == "atari":
            self.hei_state = 42
            self.wid_state = 42
            self.preprocess_mode = 3    # 0(nothing) | 1(rgb2gray) | 2(rgb2y) | 3(crop&resize)
        elif self.env_type == "lab":
            pass
        elif self.env_type == "gazebo":
            self.hei_state = 60
            self.wid_state = 80
            self.preprocess_mode = 3  # 0(nothing) | 1(rgb2gray) | 2(rgb2y) | 3(crop&resize depth)
            self.img_encoding_type = "passthrough"
        else:
            assert False, "env_type must be: gym | atari-ram | atari | lab"

class ModelParams(Params):  # settings for network architecture
    def __init__(self):
        super(ModelParams, self).__init__()

        self.state_shape = None # NOTE: set in fit_model of inherited Agents
        self.action_dim  = None # NOTE: set in fit_model of inherited Agents

class MemoryParams(Params):     # settings for replay memory
    def __init__(self):
        super(MemoryParams, self).__init__()

        if self.env_type == "gym":
            self.memory_size = 50000
        elif self.env_type == "atari-ram":
            self.memory_size = 1000000
        elif self.env_type == "atari":
            self.memory_size = 1000000
        elif self.env_type == "lab":
            self.memory_size = 1000000
        else:
            self.memory_size = 1000000

class AgentParams(Params):  # hyperparameters for drl agents
    def __init__(self):
        super(AgentParams, self).__init__()

        # criteria and optimizer
        if self.agent_type == "dqn":
            self.value_criteria = F.smooth_l1_loss
            self.optim          = optim.Adam
            # self.optim          = optim.RMSprop
        elif self.agent_type == "a3c":
            self.value_criteria = nn.MSELoss()
            self.optim          = SharedAdam    # share momentum across learners
        else:
            self.value_criteria = F.smooth_l1_loss
            self.optim          = optim.Adam
        # hyperparameters
        if self.agent_type == "dqn" and self.env_type == "gym":
            self.steps               = 100000   # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 1.#np.inf
            self.lr                  = 0.0001
            self.weight_decay        = 0.
            self.eval_freq           = 2500     # NOTE: here means every this many steps
            self.eval_steps          = 1000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 1

            self.learn_start         = 500      # start update params after this many steps
            self.batch_size          = 32
            self.valid_size          = 250
            self.eps_start           = 1
            self.eps_end             = 0.3
            self.eps_eval            = 0.#0.05
            self.eps_decay           = 50000
            self.target_model_update = 1000#0.0001
            self.action_repetition   = 1
            self.memory_interval     = 1
            self.train_interval      = 1
        elif self.agent_type == "dqn" and self.env_type == "atari-ram" or \
             self.agent_type == "dqn" and self.env_type == "atari":
            self.steps               = 50000000 # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 40.#np.inf
            self.lr                  = 0.00025
            self.weight_decay        = 0.
            self.eval_freq           = 250000#12500    # NOTE: here means every this many steps
            self.eval_steps          = 125000#2500
            self.prog_freq           = 10000#self.eval_freq
            self.test_nepisodes      = 1

            self.learn_start         = 50000    # start update params after this many steps
            self.batch_size          = 32
            self.valid_size          = 500
            self.eps_start           = 1
            self.eps_end             = 0.1
            self.eps_eval            = 0.#0.05
            self.eps_decay           = 1000000
            self.target_model_update = 10000
            self.action_repetition   = 4
            self.memory_interval     = 1
            self.train_interval      = 4
        elif self.agent_type == "a3c":
            self.steps               = 20000000 # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 40.
            self.lr                  = 0.0001
            self.weight_decay        = 1e-4 if self.enable_continuous else 0.
            self.eval_freq           = 60       # NOTE: here means every this many seconds
            self.eval_steps          = 3000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 10

            self.rollout_steps       = 20       # max look-ahead steps in a single rollout
            self.tau                 = 1.
            self.beta                = 0.01     # coefficient for entropy penalty
        elif self.agent_type == "acer":
            self.steps               = 20000000 # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 40.
            self.lr                  = 0.0001
            self.weight_decay        = 0.
            self.eval_freq           = 60       # NOTE: here means every this many seconds
            self.eval_steps          = 3000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 10

            self.rollout_steps       = 20       # max look-ahead steps in a single rollout
            self.tau                 = 1.
        else:
            self.steps               = 1000000  # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 1.#np.inf
            self.lr                  = 0.001
            self.weight_decay        = 0.
            self.eval_freq           = 2500     # NOTE: here means every this many steps
            self.eval_steps          = 1000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 10

            self.learn_start         = 25000    # start update params after this many steps
            self.batch_size          = 32
            self.valid_size          = 500
            self.eps_start           = 1
            self.eps_end             = 0.1
            self.eps_eval            = 0.#0.05
            self.eps_decay           = 50000
            self.target_model_update = 1000
            self.action_repetition   = 1
            self.memory_interval     = 1
            self.train_interval      = 4

            self.rollout_steps       = 20       # max look-ahead steps in a single rollout
            self.tau                 = 1.

        self.env_params    = EnvParams()
        self.model_params  = ModelParams()
        self.memory_params = MemoryParams()

class Options(Params):
    agent_params  = AgentParams()
<commit_msg>* small mods<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import os
import visdom
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from utils.helpers import loggerConfig
from optims.sharedAdam import SharedAdam
from optims.sharedRMSprop import SharedRMSprop

CONFIGS = [
# agent_type, env_type,    game,                       model_type,     memory_type
[ "empty",    "gym",       "CartPole-v0",              "empty",        "none"      ],  # 0
[ "dqn",      "gym",       "CartPole-v0",              "dqn-mlp",      "sequential"],  # 1
[ "dqn",      "atari-ram", "Pong-ram-v0",              "dqn-mlp",      "sequential"],  # 2
[ "dqn",      "atari",     "PongDeterministic-v4",     "dqn-cnn",      "sequential"],  # 3
[ "dqn",      "atari",     "BreakoutDeterministic-v4", "dqn-cnn",      "sequential"],  # 4
[ "a3c",      "atari",     "PongDeterministic-v4",     "a3c-cnn-dis",  "none"      ],  # 5
[ "a3c",      "gym",       "InvertedPendulum-v1",      "a3c-mlp-con",  "none"      ],  # 6
[ "acer",     "gym",       "CartPole-v1",              "acer-mlp-dis", "none"      ]   # 7  # NOTE: acer still under development, dont use this config
]

class Params(object):   # NOTE: shared across all modules
    def __init__(self):
        self.verbose     = 0            # 0(warning) | 1(info) | 2(debug)

        # training signature
        self.machine     = "daim"    # "machine_id"
        self.timestamp   = "17080900"   # "yymmdd##"
        # training configuration
        self.mode        = 1            # 1(train) | 2(test model_file)
        self.config      = 5

        self.seed        = 123
        self.render      = False        # whether render the window from the original envs or not
        self.visualize   = True         # whether do online plotting and stuff or not
        self.save_best   = False        # save model w/ highest reward if True, otherwise always save the latest model

        self.agent_type, self.env_type, self.game, self.model_type, self.memory_type = CONFIGS[self.config]

        if self.agent_type == "dqn":
            self.enable_double_dqn  = False
            self.enable_dueling     = False
            self.dueling_type       = 'avg' # avg | max | naive

            if self.env_type == "gym":
                self.hist_len       = 1
                self.hidden_dim     = 16
            else:
                self.hist_len       = 4
                self.hidden_dim     = 256

            self.use_cuda           = torch.cuda.is_available()
            self.dtype              = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
        elif self.agent_type == "a3c":
            self.enable_lstm        = True
            if "-con" in self.model_type:
                self.enable_continuous  = True
            else:
                self.enable_continuous  = False
            self.num_processes      = 16

            self.hist_len           = 1
            self.hidden_dim         = 128

            self.use_cuda           = False
            self.dtype              = torch.FloatTensor
        elif self.agent_type == "acer":
            self.enable_lstm        = True
            if "-con" in self.model_type:
                self.enable_continuous  = True
            else:
                self.enable_continuous  = False
            self.num_processes      = 16

            self.hist_len           = 1
            self.hidden_dim         = 32

            self.use_cuda           = False
            self.dtype              = torch.FloatTensor
        else:
            self.hist_len           = 1
            self.hidden_dim         = 256

            self.use_cuda           = torch.cuda.is_available()
            self.dtype              = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor

        # prefix for model/log/visdom
        self.refs        = self.machine + "_" + self.timestamp # NOTE: using this as env for visdom
        self.root_dir    = os.getcwd()

        # model files
        # NOTE: will save the current model to model_name
        self.model_name  = self.root_dir + "/models/" + self.refs + ".pth"
        # NOTE: will load pretrained model_file if not None
        self.model_file  = None#self.root_dir + "/models/{TODO:FILL_IN_PRETAINED_MODEL_FILE}.pth"
        if self.mode == 2:
            self.model_file  = self.model_name  # NOTE: so only need to change self.mode to 2 to test the current training
            assert self.model_file is not None, "Pre-Trained model is None, Testing aborted!!!"
            self.refs = self.refs + "_test"     # NOTE: using this as env for visdom for testing, to avoid accidentally redraw on the training plots

        # logging configs
        self.log_name    = self.root_dir + "/logs/" + self.refs + ".log"
        self.logger      = loggerConfig(self.log_name, self.verbose)
        self.logger.warning("<===================================>")

        if self.visualize:
            self.vis = visdom.Visdom()
            self.logger.warning("bash$: python -m visdom.server")           # activate visdom server on bash
            self.logger.warning("http://localhost:8097/env/" + self.refs)   # open this address on browser

class EnvParams(Params):    # settings for simulation environment
    def __init__(self):
        super(EnvParams, self).__init__()

        if self.env_type == "gym":
            pass
        elif self.env_type == "atari-ram":
            pass
        elif self.env_type == "atari":
            self.hei_state = 42
            self.wid_state = 42
            self.preprocess_mode = 3    # 0(nothing) | 1(rgb2gray) | 2(rgb2y) | 3(crop&resize)
        elif self.env_type == "lab":
            pass
        elif self.env_type == "gazebo":
            self.hei_state = 60
            self.wid_state = 80
            self.preprocess_mode = 3  # 0(nothing) | 1(rgb2gray) | 2(rgb2y) | 3(crop&resize depth)
            self.img_encoding_type = "passthrough"
        else:
            assert False, "env_type must be: gym | atari-ram | atari | lab"

class ModelParams(Params):  # settings for network architecture
    def __init__(self):
        super(ModelParams, self).__init__()

        self.state_shape = None # NOTE: set in fit_model of inherited Agents
        self.action_dim  = None # NOTE: set in fit_model of inherited Agents

class MemoryParams(Params):     # settings for replay memory
    def __init__(self):
        super(MemoryParams, self).__init__()

        if self.env_type == "gym":
            self.memory_size = 50000
        elif self.env_type == "atari-ram":
            self.memory_size = 1000000
        elif self.env_type == "atari":
            self.memory_size = 1000000
        elif self.env_type == "lab":
            self.memory_size = 1000000
        else:
            self.memory_size = 1000000

class AgentParams(Params):  # hyperparameters for drl agents
    def __init__(self):
        super(AgentParams, self).__init__()

        # criteria and optimizer
        if self.agent_type == "dqn":
            self.value_criteria = F.smooth_l1_loss
            self.optim          = optim.Adam
            # self.optim          = optim.RMSprop
        elif self.agent_type == "a3c":
            self.value_criteria = nn.MSELoss()
            self.optim          = SharedAdam    # share momentum across learners
        else:
            self.value_criteria = F.smooth_l1_loss
            self.optim          = optim.Adam
        # hyperparameters
        if self.agent_type == "dqn" and self.env_type == "gym":
            self.steps               = 100000   # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 1.#np.inf
            self.lr                  = 0.0001
            self.weight_decay        = 0.
            self.eval_freq           = 2500     # NOTE: here means every this many steps
            self.eval_steps          = 1000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 1

            self.learn_start         = 500      # start update params after this many steps
            self.batch_size          = 32
            self.valid_size          = 250
            self.eps_start           = 1
            self.eps_end             = 0.3
            self.eps_eval            = 0.#0.05
            self.eps_decay           = 50000
            self.target_model_update = 1000#0.0001
            self.action_repetition   = 1
            self.memory_interval     = 1
            self.train_interval      = 1
        elif self.agent_type == "dqn" and self.env_type == "atari-ram" or \
             self.agent_type == "dqn" and self.env_type == "atari":
            self.steps               = 50000000 # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 40.#np.inf
            self.lr                  = 0.00025
            self.weight_decay        = 0.
            self.eval_freq           = 250000#12500    # NOTE: here means every this many steps
            self.eval_steps          = 125000#2500
            self.prog_freq           = 10000#self.eval_freq
            self.test_nepisodes      = 1

            self.learn_start         = 50000    # start update params after this many steps
            self.batch_size          = 32
            self.valid_size          = 500
            self.eps_start           = 1
            self.eps_end             = 0.1
            self.eps_eval            = 0.#0.05
            self.eps_decay           = 1000000
            self.target_model_update = 10000
            self.action_repetition   = 4
            self.memory_interval     = 1
            self.train_interval      = 4
        elif self.agent_type == "a3c":
            self.steps               = 20000000 # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 40.
            self.lr                  = 0.0001
            self.weight_decay        = 1e-4 if self.enable_continuous else 0.
            self.eval_freq           = 60       # NOTE: here means every this many seconds
            self.eval_steps          = 3000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 10

            self.rollout_steps       = 20       # max look-ahead steps in a single rollout
            self.tau                 = 1.
            self.beta                = 0.01     # coefficient for entropy penalty
        elif self.agent_type == "acer":
            self.steps               = 20000000 # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 40.
            self.lr                  = 0.0001
            self.weight_decay        = 0.
            self.eval_freq           = 60       # NOTE: here means every this many seconds
            self.eval_steps          = 3000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 10

            self.rollout_steps       = 20       # max look-ahead steps in a single rollout
            self.tau                 = 1.
        else:
            self.steps               = 1000000  # max #iterations
            self.early_stop          = None     # max #steps per episode
            self.gamma               = 0.99
            self.clip_grad           = 1.#np.inf
            self.lr                  = 0.001
            self.weight_decay        = 0.
            self.eval_freq           = 2500     # NOTE: here means every this many steps
            self.eval_steps          = 1000
            self.prog_freq           = self.eval_freq
            self.test_nepisodes      = 10

            self.learn_start         = 25000    # start update params after this many steps
            self.batch_size          = 32
            self.valid_size          = 500
            self.eps_start           = 1
            self.eps_end             = 0.1
            self.eps_eval            = 0.#0.05
            self.eps_decay           = 50000
            self.target_model_update = 1000
            self.action_repetition   = 1
            self.memory_interval     = 1
            self.train_interval      = 4

            self.rollout_steps       = 20       # max look-ahead steps in a single rollout
            self.tau                 = 1.

        self.env_params    = EnvParams()
        self.model_params  = ModelParams()
        self.memory_params = MemoryParams()

class Options(Params):
    agent_params  = AgentParams()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Copyright (C) 2010-2015 Cuckoo Foundation.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os
import sys
import time
import logging
import argparse
import multiprocessing

logging.basicConfig(level=logging.INFO)
log = logging.getLogger()

sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), ".."))

from lib.cuckoo.common.config import Config
from lib.cuckoo.common.constants import CUCKOO_ROOT
from lib.cuckoo.core.database import Database, TASK_REPORTED, TASK_COMPLETED
from lib.cuckoo.core.database import TASK_FAILED_PROCESSING
from lib.cuckoo.core.plugins import RunProcessing, RunSignatures, RunReporting
from lib.cuckoo.core.startup import init_modules

def process(task_id, target=None, copy_path=None, report=False, auto=False):
    results = RunProcessing(task_id=task_id).run()
    RunSignatures(results=results).run()

    if report:
        RunReporting(task_id=task_id, results=results).run()
        Database().set_status(task_id)

        if auto:
            if cfg.cuckoo.delete_original and os.path.exists(target):
                os.unlink(target)

            if cfg.cuckoo.delete_bin_copy and os.path.exists(copy_path):
                os.unlink(copy_path)

def autoprocess(parallel=1):
    maxcount = cfg.cuckoo.max_analysis_count
    count = 0
    db = Database()
    pool = multiprocessing.Pool(parallel)
    pending_results = []

    # CAUTION - big ugly loop ahead.
    while count < maxcount or not maxcount:

        # Pending_results maintenance.
        for ar, tid, target, copy_path in list(pending_results):
            if ar.ready():
                if ar.successful():
                    log.info("Task #%d: reports generation completed", tid)
                else:
                    try:
                        ar.get()
                    except:
                        log.exception("Exception when processing task ID %u.", tid)
                        db.set_status(tid, TASK_FAILED_PROCESSING)

                pending_results.remove((ar, tid, target, copy_path))

        # If still full, don't add more (necessary despite pool).
        if len(pending_results) >= parallel:
            time.sleep(1)
            continue

        # If we're here, getting parallel tasks should at least
        # have one we don't know.
        tasks = db.list_tasks(status=TASK_COMPLETED, limit=parallel,
                              order_by="completed_on asc")

        # For loop to add only one, nice.
        for task in tasks:
            # Not-so-efficient lock.
            if task.id in [tid for ar, tid, target, copy_path
                           in pending_results]:
                continue

            log.info("Processing analysis data for Task #%d", task.id)

            sample = db.view_sample(task.sample_id)

            copy_path = os.path.join(CUCKOO_ROOT, "storage",
                                     "binaries", sample.sha256)

            args = task.id, task.target, copy_path
            kwargs = dict(report=True, auto=True)
            result = pool.apply_async(process, args, kwargs)

            pending_results.append((result, task.id, task.target, copy_path))

            count += 1
            break

        # If there wasn't anything to add, sleep tight.
        if not tasks:
            time.sleep(5)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("id", type=str, help="ID of the analysis to process (auto for continuous processing of unprocessed tasks).")
    parser.add_argument("-d", "--debug", help="Display debug messages", action="store_true", required=False)
    parser.add_argument("-r", "--report", help="Re-generate report", action="store_true", required=False)
    parser.add_argument("-p", "--parallel", help="Number of parallel threads to use (auto mode only).", type=int, required=False, default=1)
    args = parser.parse_args()

    if args.debug:
        log.setLevel(logging.DEBUG)

    init_modules()

    if args.id == "auto":
        autoprocess(parallel=args.parallel)
    else:
        process(args.id, report=args.report)

if __name__ == "__main__":
    cfg = Config()

    try:
        main()
    except KeyboardInterrupt:
        pass
<commit_msg>Fixed a couple of bugs blocking url analysis (Close #463)<commit_after>#!/usr/bin/env python
# Copyright (C) 2010-2015 Cuckoo Foundation.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os
import sys
import time
import logging
import argparse
import multiprocessing

logging.basicConfig(level=logging.INFO)
log = logging.getLogger()

sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), ".."))

from lib.cuckoo.common.config import Config
from lib.cuckoo.common.constants import CUCKOO_ROOT
from lib.cuckoo.core.database import Database, TASK_REPORTED, TASK_COMPLETED
from lib.cuckoo.core.database import TASK_FAILED_PROCESSING
from lib.cuckoo.core.plugins import RunProcessing, RunSignatures, RunReporting
from lib.cuckoo.core.startup import init_modules

def process(task_id, target=None, copy_path=None, report=False, auto=False):
    assert isinstance(task_id, int)

    results = RunProcessing(task_id=task_id).run()
    RunSignatures(results=results).run()

    if report:
        RunReporting(task_id=task_id, results=results).run()
        Database().set_status(task_id, TASK_REPORTED)

        if auto:
            if cfg.cuckoo.delete_original and os.path.exists(target):
                os.unlink(target)

            if cfg.cuckoo.delete_bin_copy and os.path.exists(copy_path):
                os.unlink(copy_path)

def autoprocess(parallel=1):
    maxcount = cfg.cuckoo.max_analysis_count
    count = 0
    db = Database()
    pool = multiprocessing.Pool(parallel)
    pending_results = []

    # CAUTION - big ugly loop ahead.
    while count < maxcount or not maxcount:

        # Pending_results maintenance.
        for ar, tid, target, copy_path in list(pending_results):
            if ar.ready():
                if ar.successful():
                    log.info("Task #%d: reports generation completed", tid)
                else:
                    try:
                        ar.get()
                    except:
                        log.exception("Exception when processing task ID %u.", tid)
                        db.set_status(tid, TASK_FAILED_PROCESSING)

                pending_results.remove((ar, tid, target, copy_path))

        # If still full, don't add more (necessary despite pool).
        if len(pending_results) >= parallel:
            time.sleep(1)
            continue

        # If we're here, getting parallel tasks should at least
        # have one we don't know.
        tasks = db.list_tasks(status=TASK_COMPLETED, limit=parallel,
                              order_by="completed_on asc")

        # For loop to add only one, nice.
        for task in tasks:
            # Not-so-efficient lock.
            if task.id in [tid for ar, tid, target, copy_path
                           in pending_results]:
                continue

            log.info("Processing analysis data for Task #%d", task.id)

            if task.category == "file":
                sample = db.view_sample(task.sample_id)

                copy_path = os.path.join(CUCKOO_ROOT, "storage",
                                         "binaries", sample.sha256)
            else:
                copy_path = None

            args = task.id, task.target, copy_path
            kwargs = dict(report=True, auto=True)
            result = pool.apply_async(process, args, kwargs)

            pending_results.append((result, task.id, task.target, copy_path))

            count += 1
            break

        # If there wasn't anything to add, sleep tight.
        if not tasks:
            time.sleep(5)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("id", type=str, help="ID of the analysis to process (auto for continuous processing of unprocessed tasks).")
    parser.add_argument("-d", "--debug", help="Display debug messages", action="store_true", required=False)
    parser.add_argument("-r", "--report", help="Re-generate report", action="store_true", required=False)
    parser.add_argument("-p", "--parallel", help="Number of parallel threads to use (auto mode only).", type=int, required=False, default=1)
    args = parser.parse_args()

    if args.debug:
        log.setLevel(logging.DEBUG)

    init_modules()

    if args.id == "auto":
        autoprocess(parallel=args.parallel)
    else:
        process(args.id, report=args.report)

if __name__ == "__main__":
    cfg = Config()

    try:
        main()
    except KeyboardInterrupt:
        pass
<|endoftext|>


<commit_before>#!/usr/bin/python
"""Underdark Web Framework -- uWeb"""

__author__ = 'Elmer de Looff <elmer@underdark.nl>'
__version__ = '0.14'

# Standard modules
import httplib
import os
import re
import sys

# Underdark package module
#
# If the underdark package is not directly available, add the 'ud_lib'
# directory to the module path.
try:
  import underdark
except ImportError:
  sys.path.append(
      os.path.abspath(os.path.join(os.path.dirname(__file__), 'ext_lib')))
# Underdark modules
from underdark.libs import app
from underdark.libs.app import logging

# Package modules
from . import pagemaker
from . import request

# Package classes
from .response import Response
from .response import Redirect
from .pagemaker import PageMaker
from .pagemaker import DebuggingPageMaker

class Error(Exception):
  """Superclass used for inheritance and external excepion handling."""

class ImmediateResponse(Exception):
  """Used to trigger an immediate repsonse, foregoing the regular returns."""

class NoRouteError(Error):
  """The server does not know how to route this request"""

class NewWeb(object):
  """Returns a configured closure for handling page requests.

  This closure is configured with a precomputed set of routes and handlers using
  the Router function. After this, incoming requests are processed and delegated
  to the correct PageMaker handler.

  The url in the received `req` object is taken and matches against the
  `router`` (refer to Router() for more documentation on this).

  Takes:
    @ page_class: PageMaker
      Class that holds request handling methods as defined in the `routes`
    @ router: request router
      The result of the Router() function.
    @ config: dict
      Configuration for the PageMaker. Typically contains entries for database
      connections, default search paths etc.

  Returns:
    RequestHandler: Configured closure that is ready to process requests.
  """
  def __init__(self, page_class, router, config):
    self.page_class = page_class
    self.router = router
    self.config = config

  def __call__(self, env, start_response):
    """WSGI request handler.

    Accpepts the WSGI `environment` dictionary and a function to start the
    response and returns a response iterator.
    """
    req = request.Request(env)
    pages = self.page_class(req, config=self.config)
    try:
      # We're specifically calling _PostInit here as promised in documentation.
      # pylint: disable=W0212
      pages._PostInit()
      # pylint: enable=W0212
      method, args = self.router(req.env['PATH_INFO'])
      response = getattr(pages, method)(*args)
    except pagemaker.ReloadModules, message:
      reload_message = reload(sys.modules[self.page_class.__module__])
      response = Response(
          content='%s\n%s' % (message, reload_message))
    except ImmediateResponse, response:
      response = response[0]
    except (NoRouteError, Exception):
      response = pages.InternalServerError(*sys.exc_info())

    if not isinstance(response, Response):
      response = Response(content=response)
    status = response.httpcode
    if isinstance(status, int):
      status = '%d %s' % (status, httplib.responses[status])
    headers = []
    for key, value in response.headers.items():
      headers.append((key, value.encode('utf8')))
    start_response(status, headers)
    return [response.content]

def Router(routes, prefix=None):
  """Returns the first request handler that matches the request URL.

  The `routes` argument is an iterable of 2-tuples, each of which contain a
  pattern (regex) and the name of the handler to use for matching requests.

  Before returning the closure, all regexen are compiled, and handler methods
  are retrieved from the provided `page_class`.

  Arguments:
    @ routes: iterable of 2-tuples.
      Each tuple is a pair of `pattern` and `handler`, both are strings.

  Returns:
    RequestRouter: Configured closure that processes urls.
  """
  if prefix is None:
    prefix = ''
  req_routes = []
  for pattern, method in routes:
    req_routes.append((re.compile(prefix + pattern + '$', re.UNICODE), method))

  def RequestRouter(url):
    """Returns the appropriate handler and arguments for the given `url`.

    The`url` is matched against the compiled patterns in the `req_routes`
    provided by the outer scope. Upon finding a pattern that matches, the
    match groups from the regex and the unbound handler method are returned.

    N.B. The rules are such that the first matching route will be used. There
    is no further concept of specificity. Routes should be written with this in
    mind.

    Arguments:
      @ url: str
        The URL requested by the client.

    Raises:
      NoRouteError: None of the patterns match the requested `url`.

    Returns:
      2-tuple: handler method (unbound), and tuple of pattern matches.
    """
    for pattern, handler in req_routes:
      match = pattern.match(url)
      if match:
        return handler, match.groups()
    raise NoRouteError(url +' cannot be handled')
  return RequestRouter

def ServerSetup(apache_logging=True):
  """Sets up and starts serving a WSGI application based on wsgiref

  If provided through the CONFIG constant, the configuration file will be read
  and parsed. This configuration will be used for the `StandAloneServer` for
  host and port configurations, and the PageMaker will use it for all sorts of
  configuration, for example database connection into and default search paths.

  Logging:
    For both `standalone` and `apache` mode, the PACKAGE constant will set the
    directory under which log files should be accumulated.
    * For `apache` this will create a log database 'apache.sqlite' only, and if
      the PACKAGE constant is not available, this will default to 'uweb_project'
    * For `standalone` mode, there will be '.sqlite' log files for each router,
      and the base-name will be the same as that of the router. Additionally
      there will be access and error logs, again sharing the base name with the
      router itself. The default directory name to bundle these files under will
      be the name of the directory one up from where the router runs.
  """

  # We need _getframe here, there's not really a neater way to do this.
  # pylint: disable=W0212
  entrypoint = sys._getframe(1)
  # pylint: enable=W0212
  router_file = entrypoint.f_code.co_filename

  # Configuration based on constants provided
  pages_class = entrypoint.f_globals['PAGE_CLASS']
  routes = entrypoint.f_globals['ROUTES']
  config_file = entrypoint.f_globals.get('CONFIG')
  if config_file:
    router_config = app.ParseConfig(os.path.join(
        os.path.dirname(router_file), config_file))
  else:
    router_config = {}
  req_router = Router(routes, prefix=entrypoint.f_globals.get('ROUTE_PREFIX'))
  application = NewWeb(pages_class, req_router, router_config)

  from wsgiref.simple_server import make_server
  wsgi_server = make_server('localhost', 8001, application)
  wsgi_server.serve_forever()
<commit_msg>Moved configuration parsing to the main uWeb module, removing the need for the app module.<commit_after>#!/usr/bin/python
"""Underdark Web Framework -- uWeb"""

__author__ = 'Elmer de Looff <elmer@underdark.nl>'
__version__ = '0.14'

# Standard modules
import ConfigParser
import httplib
import os
import re
import sys

# Underdark package module
#
# If the underdark package is not directly available, add the 'ud_lib'
# directory to the module path.
try:
  import underdark
except ImportError:
  sys.path.append(
      os.path.abspath(os.path.join(os.path.dirname(__file__), 'ext_lib')))
# Underdark modules
from underdark.libs.app import logging

# Package modules
from . import pagemaker
from . import request

# Package classes
from .response import Response
from .response import Redirect
from .pagemaker import PageMaker
from .pagemaker import DebuggingPageMaker

class Error(Exception):
  """Superclass used for inheritance and external excepion handling."""

class ImmediateResponse(Exception):
  """Used to trigger an immediate repsonse, foregoing the regular returns."""

class NoRouteError(Error):
  """The server does not know how to route this request"""

class NewWeb(object):
  """Returns a configured closure for handling page requests.

  This closure is configured with a precomputed set of routes and handlers using
  the Router function. After this, incoming requests are processed and delegated
  to the correct PageMaker handler.

  The url in the received `req` object is taken and matches against the
  `router`` (refer to Router() for more documentation on this).

  Takes:
    @ page_class: PageMaker
      Class that holds request handling methods as defined in the `routes`
    @ router: request router
      The result of the Router() function.
    @ config: dict
      Configuration for the PageMaker. Typically contains entries for database
      connections, default search paths etc.

  Returns:
    RequestHandler: Configured closure that is ready to process requests.
  """
  def __init__(self, page_class, router, config):
    self.page_class = page_class
    self.router = router
    self.config = config

  def __call__(self, env, start_response):
    """WSGI request handler.

    Accpepts the WSGI `environment` dictionary and a function to start the
    response and returns a response iterator.
    """
    req = request.Request(env)
    pages = self.page_class(req, config=self.config)
    try:
      # We're specifically calling _PostInit here as promised in documentation.
      # pylint: disable=W0212
      pages._PostInit()
      # pylint: enable=W0212
      method, args = self.router(req.env['PATH_INFO'])
      response = getattr(pages, method)(*args)
    except pagemaker.ReloadModules, message:
      reload_message = reload(sys.modules[self.page_class.__module__])
      response = Response(
          content='%s\n%s' % (message, reload_message))
    except ImmediateResponse, response:
      response = response[0]
    except (NoRouteError, Exception):
      response = pages.InternalServerError(*sys.exc_info())

    if not isinstance(response, Response):
      response = Response(content=response)
    status = response.httpcode
    if isinstance(status, int):
      status = '%d %s' % (status, httplib.responses[status])
    headers = []
    for key, value in response.headers.items():
      headers.append((key, value.encode('utf8')))
    start_response(status, headers)
    return [response.content]

def ParseConfig(config_file):
  """Parses the given `config_file` and returns it as a nested dictionary."""
  parser = ConfigParser.SafeConfigParser()
  config_file = os.path.join(os.getcwd(), config_file)
  try:
    parser.read(config_file)
  except ConfigParser.ParsingError:
    raise ValueError('Not a valid config file: %r.' % config_file)
  return dict((section, dict(parser.items(section)))
              for section in parser.sections())

def Router(routes, prefix=None):
  """Returns the first request handler that matches the request URL.

  The `routes` argument is an iterable of 2-tuples, each of which contain a
  pattern (regex) and the name of the handler to use for matching requests.

  Before returning the closure, all regexen are compiled, and handler methods
  are retrieved from the provided `page_class`.

  Arguments:
    @ routes: iterable of 2-tuples.
      Each tuple is a pair of `pattern` and `handler`, both are strings.

  Returns:
    RequestRouter: Configured closure that processes urls.
  """
  if prefix is None:
    prefix = ''
  req_routes = []
  for pattern, method in routes:
    req_routes.append((re.compile(prefix + pattern + '$', re.UNICODE), method))

  def RequestRouter(url):
    """Returns the appropriate handler and arguments for the given `url`.

    The`url` is matched against the compiled patterns in the `req_routes`
    provided by the outer scope. Upon finding a pattern that matches, the
    match groups from the regex and the unbound handler method are returned.

    N.B. The rules are such that the first matching route will be used. There
    is no further concept of specificity. Routes should be written with this in
    mind.

    Arguments:
      @ url: str
        The URL requested by the client.

    Raises:
      NoRouteError: None of the patterns match the requested `url`.

    Returns:
      2-tuple: handler method (unbound), and tuple of pattern matches.
    """
    for pattern, handler in req_routes:
      match = pattern.match(url)
      if match:
        return handler, match.groups()
    raise NoRouteError(url +' cannot be handled')
  return RequestRouter

def ServerSetup(apache_logging=True):
  """Sets up and starts serving a WSGI application based on wsgiref

  If provided through the CONFIG constant, the configuration file will be read
  and parsed. This configuration will be used for the `StandAloneServer` for
  host and port configurations, and the PageMaker will use it for all sorts of
  configuration, for example database connection into and default search paths.

  Logging:
    For both `standalone` and `apache` mode, the PACKAGE constant will set the
    directory under which log files should be accumulated.
    * For `apache` this will create a log database 'apache.sqlite' only, and if
      the PACKAGE constant is not available, this will default to 'uweb_project'
    * For `standalone` mode, there will be '.sqlite' log files for each router,
      and the base-name will be the same as that of the router. Additionally
      there will be access and error logs, again sharing the base name with the
      router itself. The default directory name to bundle these files under will
      be the name of the directory one up from where the router runs.
  """

  # We need _getframe here, there's not really a neater way to do this.
  # pylint: disable=W0212
  entrypoint = sys._getframe(1)
  # pylint: enable=W0212
  router_file = entrypoint.f_code.co_filename

  # Configuration based on constants provided
  pages_class = entrypoint.f_globals['PAGE_CLASS']
  routes = entrypoint.f_globals['ROUTES']
  config_file = entrypoint.f_globals.get('CONFIG')
  if config_file:
    router_config = ParseConfig(os.path.join(
        os.path.dirname(router_file), config_file))
  else:
    router_config = {}
  req_router = Router(routes, prefix=entrypoint.f_globals.get('ROUTE_PREFIX'))
  application = NewWeb(pages_class, req_router, router_config)

  from wsgiref.simple_server import make_server
  wsgi_server = make_server('localhost', 8001, application)
  wsgi_server.serve_forever()
<|endoftext|>


<commit_before>import os
from conans.errors import ConanException
import logging
from conans.util.env_reader import get_env
from conans.util.files import save, load
from six.moves.configparser import ConfigParser, NoSectionError
from conans.model.values import Values
import urllib
from conans.paths import conan_expand_user

MIN_SERVER_COMPATIBLE_VERSION = '0.6.0'

default_settings_yml = """
os: [Windows, Linux, Macos, Android, iOS]
arch: [x86, x86_64, armv6, armv7, armv7hf, armv8]
compiler:
    gcc:
        version: ["4.4", "4.5", "4.6", "4.7", "4.8", "4.9", "5.1", "5.2", "5.3", "5.4", "6.1", "6.2"]
        libcxx: [libstdc++, libstdc++11]
        threads: [None, posix, win32] #  Windows MinGW
        exception: [None, dwarf2, sjlj, seh] # Windows MinGW
    Visual Studio:
        runtime: [MD, MT, MTd, MDd]
        version: ["8", "9", "10", "11", "12", "14"]
    clang:
        version: ["3.3", "3.4", "3.5", "3.6", "3.7", "3.8"]
        libcxx: [libstdc++, libstdc++11, libc++]
    apple-clang:
        version: ["5.0", "5.1", "6.0", "6.1", "7.0", "7.3", "8.0"]
        libcxx: [libstdc++, libc++]

build_type: [None, Debug, Release]
"""

default_client_conf = '''
[storage]
# This is the default path, but you can write your own
path: ~/.conan/data

[proxies]
# Empty section will try to use system proxies.
# If don't want proxy at all, remove section [proxies]
# As documented in http://docs.python-requests.org/en/latest/user/advanced/#proxies
# http: http://user:pass@10.10.1.10:3128/
# http: http://10.10.1.10:3128
# https: http://10.10.1.10:1080

[settings_defaults]
'''

class ConanClientConfigParser(ConfigParser):

    def __init__(self, filename):
        ConfigParser.__init__(self)
        self.read(filename)

    def get_conf(self, varname):
        """Gets the section from config file or raises an exception"""
        try:
            return self.items(varname)
        except NoSectionError:
            raise ConanException("Invalid configuration, missing %s" % varname)

    @property
    def storage(self):
        return dict(self.get_conf("storage"))

    @property
    def storage_path(self):
        try:
            conan_user_home = os.getenv("CONAN_USER_HOME")
            if conan_user_home:
                storage = self.storage["path"]
                if storage[:2] == "~/":
                    storage = storage[2:]
                result = os.path.join(conan_user_home, storage)
            else:
                result = conan_expand_user(self.storage["path"])
        except KeyError:
            result = None
        result = get_env('CONAN_STORAGE_PATH', result)
        return result

    @property
    def deprecated_remotes(self):
        return self.get_conf("remotes")

    @property
    def proxies(self):
        """ optional field, might not exist
        """
        try:
            proxies = self.get_conf("proxies")
            # If there is proxies section, but empty, it will try to use system proxy
            if not proxies:
                return urllib.getproxies()
            return dict(proxies)
        except:
            return None

    @property
    def settings_defaults(self):
        default_settings = self.get_conf("settings_defaults")
        values = Values.from_list(default_settings)
        return values
<commit_msg>bumped min_server_compatible_version<commit_after>import os
from conans.errors import ConanException
import logging
from conans.util.env_reader import get_env
from conans.util.files import save, load
from six.moves.configparser import ConfigParser, NoSectionError
from conans.model.values import Values
import urllib
from conans.paths import conan_expand_user

MIN_SERVER_COMPATIBLE_VERSION = '0.12.0-dev'

default_settings_yml = """
os: [Windows, Linux, Macos, Android, iOS]
arch: [x86, x86_64, armv6, armv7, armv7hf, armv8]
compiler:
    gcc:
        version: ["4.4", "4.5", "4.6", "4.7", "4.8", "4.9", "5.1", "5.2", "5.3", "5.4", "6.1", "6.2"]
        libcxx: [libstdc++, libstdc++11]
        threads: [None, posix, win32] #  Windows MinGW
        exception: [None, dwarf2, sjlj, seh] # Windows MinGW
    Visual Studio:
        runtime: [MD, MT, MTd, MDd]
        version: ["8", "9", "10", "11", "12", "14"]
    clang:
        version: ["3.3", "3.4", "3.5", "3.6", "3.7", "3.8"]
        libcxx: [libstdc++, libstdc++11, libc++]
    apple-clang:
        version: ["5.0", "5.1", "6.0", "6.1", "7.0", "7.3", "8.0"]
        libcxx: [libstdc++, libc++]

build_type: [None, Debug, Release]
"""

default_client_conf = '''
[storage]
# This is the default path, but you can write your own
path: ~/.conan/data

[proxies]
# Empty section will try to use system proxies.
# If don't want proxy at all, remove section [proxies]
# As documented in http://docs.python-requests.org/en/latest/user/advanced/#proxies
# http: http://user:pass@10.10.1.10:3128/
# http: http://10.10.1.10:3128
# https: http://10.10.1.10:1080

[settings_defaults]
'''

class ConanClientConfigParser(ConfigParser):

    def __init__(self, filename):
        ConfigParser.__init__(self)
        self.read(filename)

    def get_conf(self, varname):
        """Gets the section from config file or raises an exception"""
        try:
            return self.items(varname)
        except NoSectionError:
            raise ConanException("Invalid configuration, missing %s" % varname)

    @property
    def storage(self):
        return dict(self.get_conf("storage"))

    @property
    def storage_path(self):
        try:
            conan_user_home = os.getenv("CONAN_USER_HOME")
            if conan_user_home:
                storage = self.storage["path"]
                if storage[:2] == "~/":
                    storage = storage[2:]
                result = os.path.join(conan_user_home, storage)
            else:
                result = conan_expand_user(self.storage["path"])
        except KeyError:
            result = None
        result = get_env('CONAN_STORAGE_PATH', result)
        return result

    @property
    def deprecated_remotes(self):
        return self.get_conf("remotes")

    @property
    def proxies(self):
        """ optional field, might not exist
        """
        try:
            proxies = self.get_conf("proxies")
            # If there is proxies section, but empty, it will try to use system proxy
            if not proxies:
                return urllib.getproxies()
            return dict(proxies)
        except:
            return None

    @property
    def settings_defaults(self):
        default_settings = self.get_conf("settings_defaults")
        values = Values.from_list(default_settings)
        return values
<|endoftext|>