<commit_before>#!/usr/bin/python
# add-node-self.py -dnsName <dnsName> -siteName siteName -nickname <sshNickname>
# DNSName for the new node.  Should be a single name or, if fully-qualified, end in <domainName>
# add-node-self.py -dnsName washington -siteName washington -nickname ig-uwashington
# add-node-self.py -dnsName washington<domainName> -siteName washington -nickname ig-uwashington
# Both generate a node with  washington<domainName> which will appear as ig-uwashington in the ssh file
# siteName and/or nickname can be omitted; if so they will default to the unqualified dns name
# add-node-self.py -dnsName washington
# add-node-self.py -dnsName washington<domainName>
# Both generate a node with  washington<domainName> which will appear as washington in the ssh file
# just updates the database; a daemon will have to update the ssh/ansible file

import sys
import json
import argparse
import urllib2
import netaddr

domainName = '.planet-ignite.net'
#
# Get my ipaddress
#
def getIP():
    try:
        resultString = urllib2.urlopen('http://api.ipify.org?format=JSON').read()
        # return json.loads(resultString)['ip']
        return resultString
    except Exception as e:
        print 'Failed to get local IP address, exiting, error message is ' + e.message

#
# An exception that is tossed for an error
#

class ValidationError(Exception):
    pass

#
# make sure IP Address is a valid IP address
#

def checkIP(aString):
    fields = aString.split('.')
    if (len(fields) != 4):
        raise ValidationError('Bad IP Address ' + aString + '.  Must have four fields, not %d' % len(fields))
    try:
        numbers = [int(field) for field in fields]
    except ValueError:
        raise ValidationError('Non-numeric field in IP Address ' + aString)
    for field in numbers:
        if (field < 0 or field > 255):
            raise ValidationError('IP Address Component %d not permitted in iP Address %s' % (field, aString))
    return aString

#
# Make sure a name (aString) is a valid hostname for dnsName or an sshNickname.  Raises a ValidationError if invalid.
# fieldName is for error-reporting: it's the field we're checking (either dnsName or sshNickname)
#
def checkName(aString, fieldName):
    firstChar = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
    allChars = firstChar + '0123456789-_'
    #
    # maybe conservative, but what the heck
    #
    if (len(aString) < 1):
        raise ValidationError("%s must be a non-empty string" % fieldName)
    if firstChar.find(aString[0]) < 0:
        raise ValidationError('First character of field %s must be in %s, not %s' % (fieldName, firstChar, aString[0]))
    for char in aString[1:]:
        if allChars.find(char) < 0:
            raise ValidationError('Characters in %s must be from the set %s, not %s' % (fieldName, allChars, char))
    return aString
#
# Return the stripped version of an FQDN ending in domainName.  If it doesn't end in domainName
# just return the string itself
#
def stripDNSName(aName):
    if (aName.endswith(domainName)):
        return aName[:-len(domainName)]
    else:
        return aName

#
# makeDNSName(aName): make aName into aName<domainName> if aName<domainName> won't give dns heartburn.
# uses checkName for that, and throws a ValidationError otherwise
#
def makeDNSName(aName):
    aName = stripDNSName(aName)
    checkName(aName, 'DNSName')
    return (aName + domainName).lower()

requestHeader = 'http://www.gee-project.org:9999/rest/add_node'

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Add myself to the GEE')
    parser.add_argument('-dnsName', type=str, nargs=1, required=True, help='DNSName for the new node.  Should be a single name or, if fully-qualified, end in <domainName>')
    parser.add_argument('-nickname', type=str, nargs=1, required=False, help='SSH Nickname for the new node.  Should be shell-friendly as it will be used in command lines')
    parser.add_argument('-siteName', type=str, nargs=1, required=False, help='Site name for the new node.  For documentation purposes only')
    parser.add_argument('-ipAddress', type=str, nargs=1, required=False, help='ip Address for the new node')
    args = parser.parse_args()
    if args.ipAddress:
        if (netaddr.IPAddress(args.ipAddress[0]).is_private() or netaddr.IPAddress(args.ipAddress[0]).is_loopback()):
            ip = getIP()
        else: ip = args.ipAddress[0]
    else : ip = getIP()
    # ip = getIP()
    try:
        dnsName = makeDNSName(args.dnsName[0])
        siteName = nickname = stripDNSName(args.dnsName[0])
        if (args.nickname): nickname = args.nickname[0]
        if (args.siteName): siteName = args.siteName[0]
        checkName(nickname, 'sshNickname')
        if ip:
            argString = '?ipAddress=%s&sshNickname=%s&dnsName=%s&siteName=%s' % (ip, nickname, dnsName, siteName)
            requestURL  = requestHeader + argString
            resultString = urllib2.urlopen(requestURL).read()
            result = json.loads(resultString)
            if (result['success']):
                print 'Successfully added node at ip address %s with DNS Name %s and ssh nickname %s' % (ip, dnsName, nickname)
            else:
                print 'add node failed with message %s' % result['message']
        else:
            print 'Could not determine local IP address, aborting'
    except Exception as e:
        print 'Failed to add node, exception is %s' % str(e)
<commit_msg>changed to port 80<commit_after>#!/usr/bin/python
# add-node-self.py -dnsName <dnsName> -siteName siteName -nickname <sshNickname>
# DNSName for the new node.  Should be a single name or, if fully-qualified, end in <domainName>
# add-node-self.py -dnsName washington -siteName washington -nickname ig-uwashington
# add-node-self.py -dnsName washington<domainName> -siteName washington -nickname ig-uwashington
# Both generate a node with  washington<domainName> which will appear as ig-uwashington in the ssh file
# siteName and/or nickname can be omitted; if so they will default to the unqualified dns name
# add-node-self.py -dnsName washington
# add-node-self.py -dnsName washington<domainName>
# Both generate a node with  washington<domainName> which will appear as washington in the ssh file
# just updates the database; a daemon will have to update the ssh/ansible file

import sys
import json
import argparse
import urllib2
import netaddr

domainName = '.planet-ignite.net'
#
# Get my ipaddress
#
def getIP():
    try:
        resultString = urllib2.urlopen('http://api.ipify.org?format=JSON').read()
        # return json.loads(resultString)['ip']
        return resultString
    except Exception as e:
        print 'Failed to get local IP address, exiting, error message is ' + e.message

#
# An exception that is tossed for an error
#

class ValidationError(Exception):
    pass

#
# make sure IP Address is a valid IP address
#

def checkIP(aString):
    fields = aString.split('.')
    if (len(fields) != 4):
        raise ValidationError('Bad IP Address ' + aString + '.  Must have four fields, not %d' % len(fields))
    try:
        numbers = [int(field) for field in fields]
    except ValueError:
        raise ValidationError('Non-numeric field in IP Address ' + aString)
    for field in numbers:
        if (field < 0 or field > 255):
            raise ValidationError('IP Address Component %d not permitted in iP Address %s' % (field, aString))
    return aString

#
# Make sure a name (aString) is a valid hostname for dnsName or an sshNickname.  Raises a ValidationError if invalid.
# fieldName is for error-reporting: it's the field we're checking (either dnsName or sshNickname)
#
def checkName(aString, fieldName):
    firstChar = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
    allChars = firstChar + '0123456789-_'
    #
    # maybe conservative, but what the heck
    #
    if (len(aString) < 1):
        raise ValidationError("%s must be a non-empty string" % fieldName)
    if firstChar.find(aString[0]) < 0:
        raise ValidationError('First character of field %s must be in %s, not %s' % (fieldName, firstChar, aString[0]))
    for char in aString[1:]:
        if allChars.find(char) < 0:
            raise ValidationError('Characters in %s must be from the set %s, not %s' % (fieldName, allChars, char))
    return aString
#
# Return the stripped version of an FQDN ending in domainName.  If it doesn't end in domainName
# just return the string itself
#
def stripDNSName(aName):
    if (aName.endswith(domainName)):
        return aName[:-len(domainName)]
    else:
        return aName

#
# makeDNSName(aName): make aName into aName<domainName> if aName<domainName> won't give dns heartburn.
# uses checkName for that, and throws a ValidationError otherwise
#
def makeDNSName(aName):
    aName = stripDNSName(aName)
    checkName(aName, 'DNSName')
    return (aName + domainName).lower()

requestHeader = 'http://planet-ignite.org/rest/add_node'

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Add myself to PlanetIgnite')
    parser.add_argument('-dnsName', type=str, nargs=1, required=True, help='DNSName for the new node.  Should be a single name or, if fully-qualified, end in <domainName>')
    parser.add_argument('-nickname', type=str, nargs=1, required=False, help='SSH Nickname for the new node.  Should be shell-friendly as it will be used in command lines')
    parser.add_argument('-siteName', type=str, nargs=1, required=False, help='Site name for the new node.  For documentation purposes only')
    parser.add_argument('-ipAddress', type=str, nargs=1, required=False, help='ip Address for the new node')
    args = parser.parse_args()
    if args.ipAddress:
        if (netaddr.IPAddress(args.ipAddress[0]).is_private() or netaddr.IPAddress(args.ipAddress[0]).is_loopback()):
            ip = getIP()
        else: ip = args.ipAddress[0]
    else : ip = getIP()
    # ip = getIP()
    try:
        dnsName = makeDNSName(args.dnsName[0])
        siteName = nickname = stripDNSName(args.dnsName[0])
        if (args.nickname): nickname = args.nickname[0]
        if (args.siteName): siteName = args.siteName[0]
        checkName(nickname, 'sshNickname')
        if ip:
            argString = '?ipAddress=%s&sshNickname=%s&dnsName=%s&siteName=%s' % (ip, nickname, dnsName, siteName)
            requestURL  = requestHeader + argString
            resultString = urllib2.urlopen(requestURL).read()
            result = json.loads(resultString)
            if (result['success']):
                print 'Successfully added node at ip address %s with DNS Name %s and ssh nickname %s' % (ip, dnsName, nickname)
            else:
                print 'add node failed with message %s' % result['message']
        else:
            print 'Could not determine local IP address, aborting'
    except Exception as e:
        print 'Failed to add node, exception is %s' % str(e)
<|endoftext|>


<commit_before>from slm_lab.agent import net
from slm_lab.agent.algorithm import policy_util
from slm_lab.agent.algorithm.actor_critic import ActorCritic
from slm_lab.agent.net import net_util
from slm_lab.lib import logger, util
from slm_lab.lib.decorator import lab_api
import numpy as np
import torch

logger = logger.get_logger(__name__)

class SoftActorCritic(ActorCritic):
    '''
    Implementation of Soft Actor-Critic (SAC)
    Original paper: "Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor"
    https://arxiv.org/abs/1801.01290
    Improvement of SAC paper: "Soft Actor-Critic Algorithms and Applications"
    https://arxiv.org/abs/1812.05905

    e.g. algorithm_spec
    "algorithm": {
        "name": "SoftActorCritic",
        "action_pdtype": "default",
        "action_policy": "default",
        "gamma": 0.99,
        "training_frequency": 1,
    }
    '''
    @lab_api
    def init_algorithm_params(self):
        '''Initialize other algorithm parameters'''
        # set default
        util.set_attr(self, dict(
            action_pdtype='default',
            action_policy='default',
            training_iter=self.body.env.num_envs,
            training_start_step=self.body.memory.batch_size,
        ))
        util.set_attr(self, self.algorithm_spec, [
            'action_pdtype',
            'action_policy',
            'gamma',  # the discount factor
            'training_iter',
            'training_frequency',
        ])
        if self.body.is_discrete:
            assert self.action_pdtype == 'GumbelSoftmax'
        self.to_train = 0
        self.action_policy = getattr(policy_util, self.action_policy)

    @lab_api
    def init_nets(self, global_nets=None):
        '''
        Networks: net(actor/policy), q1_net, target_q1_net, q2_net, target_q2_net
        All networks are separate, and have the same hidden layer architectures and optim specs, so tuning is minimal
        '''
        self.shared = False  # SAC does not share networks
        NetClass = getattr(net, self.net_spec['type'])
        out_dim = net_util.get_out_dim(self.body)
        # main actor network
        self.net = NetClass(self.net_spec, self.body.state_dim, out_dim)
        self.net_names = ['net']
        # two critic Q-networks to mitigate positive bias in q_loss and speed up training
        q_in_dim = self.body.state_dim + self.body.action_dim  # NOTE concat s, a for now
        val_out_dim = 1
        self.q1_net = NetClass(self.net_spec, q_in_dim, val_out_dim)
        self.target_q1_net = NetClass(self.net_spec, q_in_dim, val_out_dim)
        self.q2_net = NetClass(self.net_spec, q_in_dim, val_out_dim)
        self.target_q2_net = NetClass(self.net_spec, q_in_dim, val_out_dim)
        self.net_names += ['q1_net', 'target_q1_net', 'q2_net', 'target_q2_net']
        net_util.copy(self.q1_net, self.target_q1_net)
        net_util.copy(self.q2_net, self.target_q2_net)
        # temperature variable to be learned, and its target entropy
        self.log_alpha = torch.zeros(1, requires_grad=True)
        self.alpha = self.log_alpha.detach().exp()
        self.target_entropy = - np.product(self.body.action_space.shape)

        # init net optimizer and its lr scheduler
        self.optim = net_util.get_optim(self.net, self.net.optim_spec)
        self.lr_scheduler = net_util.get_lr_scheduler(self.optim, self.net.lr_scheduler_spec)
        self.q1_optim = net_util.get_optim(self.q1_net, self.q1_net.optim_spec)
        self.q1_lr_scheduler = net_util.get_lr_scheduler(self.q1_optim, self.q1_net.lr_scheduler_spec)
        self.q2_optim = net_util.get_optim(self.q2_net, self.q2_net.optim_spec)
        self.q2_lr_scheduler = net_util.get_lr_scheduler(self.q2_optim, self.q2_net.lr_scheduler_spec)
        self.alpha_optim = net_util.get_optim(self.log_alpha, self.net.optim_spec)
        self.alpha_lr_scheduler = net_util.get_lr_scheduler(self.alpha_optim, self.net.lr_scheduler_spec)
        net_util.set_global_nets(self, global_nets)
        self.post_init_nets()

    @lab_api
    def act(self, state):
        if self.body.env.clock.frame < self.training_start_step:
            return policy_util.random(state, self, self.body).cpu().squeeze().numpy()
        else:
            action = self.action_policy(state, self, self.body)
            if self.body.is_discrete:
                # discrete output is RelaxedOneHotCategorical, need to sample to int
                action = torch.distributions.Categorical(probs=action).sample()
            else:
                action = self.scale_action(torch.tanh(action))  # continuous action bound
            return action.cpu().squeeze().numpy()

    def scale_action(self, action):
        '''Scale continuous actions from tanh range'''
        action_space = self.body.action_space
        low, high = torch.from_numpy(action_space.low), torch.from_numpy(action_space.high)
        return action * (high - low) / 2 + (low + high) / 2

    def guard_q_actions(self, actions):
        '''Guard to convert actions to one-hot for input to Q-network'''
        if self.body.is_discrete:
            # TODO support multi-discrete actions
            actions = F.one_hot(actions.long(), self.body.action_dim).float()
        return actions

    def calc_log_prob_action(self, action_pd, reparam=False):
        '''Calculate log_probs and actions with option to reparametrize from paper eq. 11'''
        samples = action_pd.rsample() if reparam else action_pd.sample()
        if self.body.is_discrete:  # this is straightforward using GumbelSoftmax
            actions = samples
            log_probs = action_pd.log_prob(actions)
        else:
            mus = samples
            actions = self.scale_action(torch.tanh(mus))
            if reparam:
                # paper Appendix C. Enforcing Action Bounds for continuous actions
                log_probs = (action_pd.log_prob(mus) - torch.log(1 - actions.pow(2) + 1e-6)).sum(1)
            else:
                log_probs = action_pd.log_prob(actions)
        return log_probs, actions

    def calc_q(self, state, action, net):
        '''Forward-pass to calculate the predicted state-action-value from q1_net.'''
        s_a = torch.cat((state, action), dim=-1)
        q_pred = net(s_a).view(-1)
        return q_pred

    def calc_q_targets(self, batch):
        '''Q_tar = r + gamma * (target_Q(s', a') - alpha * log pi(a'|s'))'''
        next_states = batch['next_states']
        with torch.no_grad():
            pdparams = self.calc_pdparam(next_states)
            action_pd = policy_util.init_action_pd(self.body.ActionPD, pdparams)
            next_log_probs, next_actions = self.calc_log_prob_action(action_pd)

            next_target_q1_preds = self.calc_q(next_states, next_actions, self.target_q1_net)
            next_target_q2_preds = self.calc_q(next_states, next_actions, self.target_q2_net)
            next_target_q_preds = torch.min(next_target_q1_preds, next_target_q2_preds)
            q_targets = batch['rewards'] + self.gamma * (1 - batch['dones']) * (next_target_q_preds - self.alpha * next_log_probs)
        return q_targets

    def calc_reg_loss(self, preds, targets):
        '''Calculate the regression loss for V and Q values, using the same loss function from net_spec'''
        assert preds.shape == targets.shape, f'{preds.shape} != {targets.shape}'
        reg_loss = self.net.loss_fn(preds, targets)
        return reg_loss

    def calc_policy_loss(self, batch, log_probs, reparam_actions):
        '''policy_loss = alpha * log pi(f(a)|s) - Q1(s, f(a)), where f(a) = reparametrized action'''
        states = batch['states']
        q1_preds = self.calc_q(states, reparam_actions, self.q1_net)
        q2_preds = self.calc_q(states, reparam_actions, self.q2_net)
        q_preds = torch.min(q1_preds, q2_preds)
        policy_loss = (self.alpha * log_probs - q_preds).mean()
        return policy_loss

    def calc_alpha_loss(self, log_probs):
        alpha_loss = - (self.log_alpha * (log_probs.detach() + self.target_entropy)).mean()
        return alpha_loss

    def try_update_per(self, q_preds, q_targets):
        if 'Prioritized' in util.get_class_name(self.body.memory):  # PER
            with torch.no_grad():
                errors = (q_preds - q_targets).abs().cpu().numpy()
            self.body.memory.update_priorities(errors)

    def train_alpha(self, alpha_loss):
        '''Custom method to train the alpha variable'''
        self.alpha_lr_scheduler.step(epoch=self.body.env.clock.frame)
        self.alpha_optim.zero_grad()
        alpha_loss.backward()
        self.alpha_optim.step()
        self.alpha = self.log_alpha.detach().exp()

    def train(self):
        '''Train actor critic by computing the loss in batch efficiently'''
        if util.in_eval_lab_modes():
            return np.nan
        clock = self.body.env.clock
        if self.to_train == 1:
            for _ in range(self.training_iter):
                batch = self.sample()
                clock.set_batch_size(len(batch))

                states = batch['states']
                actions = self.guard_q_actions(batch['actions'])
                q_targets = self.calc_q_targets(batch)
                # Q-value loss for both Q nets
                q1_preds = self.calc_q(states, actions, self.q1_net)
                q1_loss = self.calc_reg_loss(q1_preds, q_targets)
                self.q1_net.train_step(q1_loss, self.q1_optim, self.q1_lr_scheduler, clock=clock, global_net=self.global_q1_net)

                q2_preds = self.calc_q(states, actions, self.q2_net)
                q2_loss = self.calc_reg_loss(q2_preds, q_targets)
                self.q2_net.train_step(q2_loss, self.q2_optim, self.q2_lr_scheduler, clock=clock, global_net=self.global_q2_net)

                # policy loss
                action_pd = policy_util.init_action_pd(self.body.ActionPD, self.calc_pdparam(states))
                log_probs, reparam_actions = self.calc_log_prob_action(action_pd, reparam=True)
                policy_loss = self.calc_policy_loss(batch, log_probs, reparam_actions)
                self.net.train_step(policy_loss, self.optim, self.lr_scheduler, clock=clock, global_net=self.global_net)

                # alpha loss
                alpha_loss = self.calc_alpha_loss(log_probs)
                self.train_alpha(alpha_loss)

                loss = q1_loss + q2_loss + policy_loss + alpha_loss
                # update target networks
                self.update_nets()
                # update PER priorities if availalbe
                self.try_update_per(torch.min(q1_preds, q2_preds), q_targets)

            # reset
            self.to_train = 0
            logger.debug(f'Trained {self.name} at epi: {clock.epi}, frame: {clock.frame}, t: {clock.t}, total_reward so far: {self.body.env.total_reward}, loss: {loss:g}')
            return loss.item()
        else:
            return np.nan

    def update_nets(self):
        '''Update target networks'''
        if util.frame_mod(self.body.env.clock.frame, self.q1_net.update_frequency, self.body.env.num_envs):
            if self.q1_net.update_type == 'replace':
                net_util.copy(self.q1_net, self.target_q1_net)
                net_util.copy(self.q2_net, self.target_q2_net)
            elif self.q1_net.update_type == 'polyak':
                net_util.polyak_update(self.q1_net, self.target_q1_net, self.q1_net.polyak_coef)
                net_util.polyak_update(self.q2_net, self.target_q2_net, self.q2_net.polyak_coef)
            else:
                raise ValueError('Unknown q1_net.update_type. Should be "replace" or "polyak". Exiting.')

    @lab_api
    def update(self):
        '''Override parent method to do nothing'''
        return self.body.explore_var
<commit_msg>sum external log prob too<commit_after>from slm_lab.agent import net
from slm_lab.agent.algorithm import policy_util
from slm_lab.agent.algorithm.actor_critic import ActorCritic
from slm_lab.agent.net import net_util
from slm_lab.lib import logger, util
from slm_lab.lib.decorator import lab_api
import numpy as np
import torch

logger = logger.get_logger(__name__)

class SoftActorCritic(ActorCritic):
    '''
    Implementation of Soft Actor-Critic (SAC)
    Original paper: "Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor"
    https://arxiv.org/abs/1801.01290
    Improvement of SAC paper: "Soft Actor-Critic Algorithms and Applications"
    https://arxiv.org/abs/1812.05905

    e.g. algorithm_spec
    "algorithm": {
        "name": "SoftActorCritic",
        "action_pdtype": "default",
        "action_policy": "default",
        "gamma": 0.99,
        "training_frequency": 1,
    }
    '''
    @lab_api
    def init_algorithm_params(self):
        '''Initialize other algorithm parameters'''
        # set default
        util.set_attr(self, dict(
            action_pdtype='default',
            action_policy='default',
            training_iter=self.body.env.num_envs,
            training_start_step=self.body.memory.batch_size,
        ))
        util.set_attr(self, self.algorithm_spec, [
            'action_pdtype',
            'action_policy',
            'gamma',  # the discount factor
            'training_iter',
            'training_frequency',
        ])
        if self.body.is_discrete:
            assert self.action_pdtype == 'GumbelSoftmax'
        self.to_train = 0
        self.action_policy = getattr(policy_util, self.action_policy)

    @lab_api
    def init_nets(self, global_nets=None):
        '''
        Networks: net(actor/policy), q1_net, target_q1_net, q2_net, target_q2_net
        All networks are separate, and have the same hidden layer architectures and optim specs, so tuning is minimal
        '''
        self.shared = False  # SAC does not share networks
        NetClass = getattr(net, self.net_spec['type'])
        out_dim = net_util.get_out_dim(self.body)
        # main actor network
        self.net = NetClass(self.net_spec, self.body.state_dim, out_dim)
        self.net_names = ['net']
        # two critic Q-networks to mitigate positive bias in q_loss and speed up training
        q_in_dim = self.body.state_dim + self.body.action_dim  # NOTE concat s, a for now
        val_out_dim = 1
        self.q1_net = NetClass(self.net_spec, q_in_dim, val_out_dim)
        self.target_q1_net = NetClass(self.net_spec, q_in_dim, val_out_dim)
        self.q2_net = NetClass(self.net_spec, q_in_dim, val_out_dim)
        self.target_q2_net = NetClass(self.net_spec, q_in_dim, val_out_dim)
        self.net_names += ['q1_net', 'target_q1_net', 'q2_net', 'target_q2_net']
        net_util.copy(self.q1_net, self.target_q1_net)
        net_util.copy(self.q2_net, self.target_q2_net)
        # temperature variable to be learned, and its target entropy
        self.log_alpha = torch.zeros(1, requires_grad=True)
        self.alpha = self.log_alpha.detach().exp()
        self.target_entropy = - np.product(self.body.action_space.shape)

        # init net optimizer and its lr scheduler
        self.optim = net_util.get_optim(self.net, self.net.optim_spec)
        self.lr_scheduler = net_util.get_lr_scheduler(self.optim, self.net.lr_scheduler_spec)
        self.q1_optim = net_util.get_optim(self.q1_net, self.q1_net.optim_spec)
        self.q1_lr_scheduler = net_util.get_lr_scheduler(self.q1_optim, self.q1_net.lr_scheduler_spec)
        self.q2_optim = net_util.get_optim(self.q2_net, self.q2_net.optim_spec)
        self.q2_lr_scheduler = net_util.get_lr_scheduler(self.q2_optim, self.q2_net.lr_scheduler_spec)
        self.alpha_optim = net_util.get_optim(self.log_alpha, self.net.optim_spec)
        self.alpha_lr_scheduler = net_util.get_lr_scheduler(self.alpha_optim, self.net.lr_scheduler_spec)
        net_util.set_global_nets(self, global_nets)
        self.post_init_nets()

    @lab_api
    def act(self, state):
        if self.body.env.clock.frame < self.training_start_step:
            return policy_util.random(state, self, self.body).cpu().squeeze().numpy()
        else:
            action = self.action_policy(state, self, self.body)
            if self.body.is_discrete:
                # discrete output is RelaxedOneHotCategorical, need to sample to int
                action = torch.distributions.Categorical(probs=action).sample()
            else:
                action = self.scale_action(torch.tanh(action))  # continuous action bound
            return action.cpu().squeeze().numpy()

    def scale_action(self, action):
        '''Scale continuous actions from tanh range'''
        action_space = self.body.action_space
        low, high = torch.from_numpy(action_space.low), torch.from_numpy(action_space.high)
        return action * (high - low) / 2 + (low + high) / 2

    def guard_q_actions(self, actions):
        '''Guard to convert actions to one-hot for input to Q-network'''
        if self.body.is_discrete:
            # TODO support multi-discrete actions
            actions = F.one_hot(actions.long(), self.body.action_dim).float()
        return actions

    def calc_log_prob_action(self, action_pd, reparam=False):
        '''Calculate log_probs and actions with option to reparametrize from paper eq. 11'''
        samples = action_pd.rsample() if reparam else action_pd.sample()
        if self.body.is_discrete:  # this is straightforward using GumbelSoftmax
            actions = samples
            log_probs = action_pd.log_prob(actions)
        else:
            mus = samples
            actions = self.scale_action(torch.tanh(mus))
            if reparam:
                # paper Appendix C. Enforcing Action Bounds for continuous actions
                log_probs = (action_pd.log_prob(mus) - torch.log(1 - actions.pow(2) + 1e-6)).sum(1)
            else:
                log_probs = action_pd.log_prob(actions).sum(1)
        return log_probs, actions

    def calc_q(self, state, action, net):
        '''Forward-pass to calculate the predicted state-action-value from q1_net.'''
        s_a = torch.cat((state, action), dim=-1)
        q_pred = net(s_a).view(-1)
        return q_pred

    def calc_q_targets(self, batch):
        '''Q_tar = r + gamma * (target_Q(s', a') - alpha * log pi(a'|s'))'''
        next_states = batch['next_states']
        with torch.no_grad():
            pdparams = self.calc_pdparam(next_states)
            action_pd = policy_util.init_action_pd(self.body.ActionPD, pdparams)
            next_log_probs, next_actions = self.calc_log_prob_action(action_pd)

            next_target_q1_preds = self.calc_q(next_states, next_actions, self.target_q1_net)
            next_target_q2_preds = self.calc_q(next_states, next_actions, self.target_q2_net)
            next_target_q_preds = torch.min(next_target_q1_preds, next_target_q2_preds)
            q_targets = batch['rewards'] + self.gamma * (1 - batch['dones']) * (next_target_q_preds - self.alpha * next_log_probs)
        return q_targets

    def calc_reg_loss(self, preds, targets):
        '''Calculate the regression loss for V and Q values, using the same loss function from net_spec'''
        assert preds.shape == targets.shape, f'{preds.shape} != {targets.shape}'
        reg_loss = self.net.loss_fn(preds, targets)
        return reg_loss

    def calc_policy_loss(self, batch, log_probs, reparam_actions):
        '''policy_loss = alpha * log pi(f(a)|s) - Q1(s, f(a)), where f(a) = reparametrized action'''
        states = batch['states']
        q1_preds = self.calc_q(states, reparam_actions, self.q1_net)
        q2_preds = self.calc_q(states, reparam_actions, self.q2_net)
        q_preds = torch.min(q1_preds, q2_preds)
        policy_loss = (self.alpha * log_probs - q_preds).mean()
        return policy_loss

    def calc_alpha_loss(self, log_probs):
        alpha_loss = - (self.log_alpha * (log_probs.detach() + self.target_entropy)).mean()
        return alpha_loss

    def try_update_per(self, q_preds, q_targets):
        if 'Prioritized' in util.get_class_name(self.body.memory):  # PER
            with torch.no_grad():
                errors = (q_preds - q_targets).abs().cpu().numpy()
            self.body.memory.update_priorities(errors)

    def train_alpha(self, alpha_loss):
        '''Custom method to train the alpha variable'''
        self.alpha_lr_scheduler.step(epoch=self.body.env.clock.frame)
        self.alpha_optim.zero_grad()
        alpha_loss.backward()
        self.alpha_optim.step()
        self.alpha = self.log_alpha.detach().exp()

    def train(self):
        '''Train actor critic by computing the loss in batch efficiently'''
        if util.in_eval_lab_modes():
            return np.nan
        clock = self.body.env.clock
        if self.to_train == 1:
            for _ in range(self.training_iter):
                batch = self.sample()
                clock.set_batch_size(len(batch))

                states = batch['states']
                actions = self.guard_q_actions(batch['actions'])
                q_targets = self.calc_q_targets(batch)
                # Q-value loss for both Q nets
                q1_preds = self.calc_q(states, actions, self.q1_net)
                q1_loss = self.calc_reg_loss(q1_preds, q_targets)
                self.q1_net.train_step(q1_loss, self.q1_optim, self.q1_lr_scheduler, clock=clock, global_net=self.global_q1_net)

                q2_preds = self.calc_q(states, actions, self.q2_net)
                q2_loss = self.calc_reg_loss(q2_preds, q_targets)
                self.q2_net.train_step(q2_loss, self.q2_optim, self.q2_lr_scheduler, clock=clock, global_net=self.global_q2_net)

                # policy loss
                action_pd = policy_util.init_action_pd(self.body.ActionPD, self.calc_pdparam(states))
                log_probs, reparam_actions = self.calc_log_prob_action(action_pd, reparam=True)
                policy_loss = self.calc_policy_loss(batch, log_probs, reparam_actions)
                self.net.train_step(policy_loss, self.optim, self.lr_scheduler, clock=clock, global_net=self.global_net)

                # alpha loss
                alpha_loss = self.calc_alpha_loss(log_probs)
                self.train_alpha(alpha_loss)

                loss = q1_loss + q2_loss + policy_loss + alpha_loss
                # update target networks
                self.update_nets()
                # update PER priorities if availalbe
                self.try_update_per(torch.min(q1_preds, q2_preds), q_targets)

            # reset
            self.to_train = 0
            logger.debug(f'Trained {self.name} at epi: {clock.epi}, frame: {clock.frame}, t: {clock.t}, total_reward so far: {self.body.env.total_reward}, loss: {loss:g}')
            return loss.item()
        else:
            return np.nan

    def update_nets(self):
        '''Update target networks'''
        if util.frame_mod(self.body.env.clock.frame, self.q1_net.update_frequency, self.body.env.num_envs):
            if self.q1_net.update_type == 'replace':
                net_util.copy(self.q1_net, self.target_q1_net)
                net_util.copy(self.q2_net, self.target_q2_net)
            elif self.q1_net.update_type == 'polyak':
                net_util.polyak_update(self.q1_net, self.target_q1_net, self.q1_net.polyak_coef)
                net_util.polyak_update(self.q2_net, self.target_q2_net, self.q2_net.polyak_coef)
            else:
                raise ValueError('Unknown q1_net.update_type. Should be "replace" or "polyak". Exiting.')

    @lab_api
    def update(self):
        '''Override parent method to do nothing'''
        return self.body.explore_var
<|endoftext|>


<commit_before>from slm_lab.agent import net
from slm_lab.agent.algorithm import policy_util
from slm_lab.agent.algorithm.actor_critic import ActorCritic
from slm_lab.agent.net import net_util
from slm_lab.lib import logger, util
from slm_lab.lib.decorator import lab_api
import numpy as np
import torch
import torch.nn.functional as F

logger = logger.get_logger(__name__)

class SoftActorCritic(ActorCritic):
    '''
    Implementation of Soft Actor-Critic (SAC)
    Original paper: "Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor"
    https://arxiv.org/abs/1801.01290
    Improvement of SAC paper: "Soft Actor-Critic Algorithms and Applications"
    https://arxiv.org/abs/1812.05905

    e.g. algorithm_spec
    "algorithm": {
        "name": "SoftActorCritic",
        "action_pdtype": "default",
        "action_policy": "default",
        "gamma": 0.99,
        "training_frequency": 1,
    }
    '''
    @lab_api
    def init_algorithm_params(self):
        '''Initialize other algorithm parameters'''
        # set default
        util.set_attr(self, dict(
            action_pdtype='default',
            action_policy='default',
            training_iter=self.body.env.num_envs,
            training_start_step=self.body.memory.batch_size,
        ))
        util.set_attr(self, self.algorithm_spec, [
            'action_pdtype',
            'action_policy',
            'gamma',  # the discount factor
            'training_iter',
            'training_frequency',
            'training_start_step',
        ])
        if self.body.is_discrete:
            assert self.action_pdtype == 'GumbelSoftmax'
        self.to_train = 0
        self.action_policy = getattr(policy_util, self.action_policy)

    @lab_api
    def init_nets(self, global_nets=None):
        '''
        Networks: net(actor/policy), q1_net, target_q1_net, q2_net, target_q2_net
        All networks are separate, and have the same hidden layer architectures and optim specs, so tuning is minimal
        '''
        self.shared = False  # SAC does not share networks
        NetClass = getattr(net, self.net_spec['type'])
        # main actor network
        self.net = NetClass(self.net_spec, self.body.state_dim, net_util.get_out_dim(self.body))
        self.net_names = ['net']
        # two critic Q-networks to mitigate positive bias in q_loss and speed up training, uses q_net.py with prefix Q
        QNetClass = getattr(net, 'Q' + self.net_spec['type'])
        q_in_dim = [self.body.state_dim, self.body.action_dim]
        self.q1_net = QNetClass(self.net_spec, q_in_dim, 1)
        self.target_q1_net = QNetClass(self.net_spec, q_in_dim, 1)
        self.q2_net = QNetClass(self.net_spec, q_in_dim, 1)
        self.target_q2_net = QNetClass(self.net_spec, q_in_dim, 1)
        self.net_names += ['q1_net', 'target_q1_net', 'q2_net', 'target_q2_net']
        net_util.copy(self.q1_net, self.target_q1_net)
        net_util.copy(self.q2_net, self.target_q2_net)
        # temperature variable to be learned, and its target entropy
        self.log_alpha = torch.zeros(1, requires_grad=True, device=self.net.device)
        self.alpha = self.log_alpha.detach().exp()
        if self.body.is_discrete:
            self.target_entropy = - self.body.action_space.n
        else:
            self.target_entropy = - np.product(self.body.action_space.shape)

        # init net optimizer and its lr scheduler
        self.optim = net_util.get_optim(self.net, self.net.optim_spec)
        self.lr_scheduler = net_util.get_lr_scheduler(self.optim, self.net.lr_scheduler_spec)
        self.q1_optim = net_util.get_optim(self.q1_net, self.q1_net.optim_spec)
        self.q1_lr_scheduler = net_util.get_lr_scheduler(self.q1_optim, self.q1_net.lr_scheduler_spec)
        self.q2_optim = net_util.get_optim(self.q2_net, self.q2_net.optim_spec)
        self.q2_lr_scheduler = net_util.get_lr_scheduler(self.q2_optim, self.q2_net.lr_scheduler_spec)
        self.alpha_optim = net_util.get_optim(self.log_alpha, self.net.optim_spec)
        self.alpha_lr_scheduler = net_util.get_lr_scheduler(self.alpha_optim, self.net.lr_scheduler_spec)
        net_util.set_global_nets(self, global_nets)
        self.post_init_nets()

    @lab_api
    def act(self, state):
        if self.body.env.clock.frame < self.training_start_step:
            return policy_util.random(state, self, self.body).cpu().squeeze().numpy()
        else:
            action = self.action_policy(state, self, self.body)
            if not self.body.is_discrete:
                action = self.scale_action(torch.tanh(action))  # continuous action bound
            return action.cpu().squeeze().numpy()

    def scale_action(self, action):
        '''Scale continuous actions from tanh range'''
        action_space = self.body.action_space
        low, high = torch.from_numpy(action_space.low), torch.from_numpy(action_space.high)
        return action * (high - low) / 2 + (low + high) / 2

    def guard_q_actions(self, actions):
        '''Guard to convert actions to one-hot for input to Q-network'''
        if self.body.is_discrete:
            # TODO support multi-discrete actions
            actions = F.one_hot(actions.long(), self.body.action_dim).float()
        return actions

    def calc_log_prob_action(self, action_pd, reparam=False):
        '''Calculate log_probs and actions with option to reparametrize from paper eq. 11'''
        samples = action_pd.rsample() if reparam else action_pd.sample()
        if self.body.is_discrete:  # this is straightforward using GumbelSoftmax
            actions = samples
            log_probs = action_pd.log_prob(actions)
        else:
            mus = samples
            actions = self.scale_action(torch.tanh(mus))
            if actions.dim() == 1:  # handle shape consistency for single actions
                actions = actions.unsqueeze(dim=-1)
            # paper Appendix C. Enforcing Action Bounds for continuous actions
            log_probs = (action_pd.log_prob(mus) - torch.log(1 - actions.pow(2) + 1e-6).sum(1))
        return log_probs, actions

    def calc_q(self, state, action, net):
        '''Forward-pass to calculate the predicted state-action-value from q1_net.'''
        if not self.body.is_discrete and action.dim() == 1:  # handle shape consistency for single continuous action
            action = action.unsqueeze(dim=-1)
        q_pred = net(state, action).view(-1)
        return q_pred

    def calc_q_targets(self, batch):
        '''Q_tar = r + gamma * (target_Q(s', a') - alpha * log pi(a'|s'))'''
        next_states = batch['next_states']
        with torch.no_grad():
            pdparams = self.calc_pdparam(next_states)
            action_pd = policy_util.init_action_pd(self.body.ActionPD, pdparams)
            next_log_probs, next_actions = self.calc_log_prob_action(action_pd)
            next_actions = self.guard_q_actions(next_actions)  # non-reparam discrete actions need to be converted into one-hot

            next_target_q1_preds = self.calc_q(next_states, next_actions, self.target_q1_net)
            next_target_q2_preds = self.calc_q(next_states, next_actions, self.target_q2_net)
            next_target_q_preds = torch.min(next_target_q1_preds, next_target_q2_preds)
            q_targets = batch['rewards'] + self.gamma * (1 - batch['dones']) * (next_target_q_preds - self.alpha * next_log_probs)
        return q_targets

    def calc_reg_loss(self, preds, targets):
        '''Calculate the regression loss for V and Q values, using the same loss function from net_spec'''
        assert preds.shape == targets.shape, f'{preds.shape} != {targets.shape}'
        reg_loss = self.net.loss_fn(preds, targets)
        self.body.tb_tracker['reg_loss'] = reg_loss
        return reg_loss

    def calc_policy_loss(self, batch, log_probs, reparam_actions):
        '''policy_loss = alpha * log pi(f(a)|s) - Q1(s, f(a)), where f(a) = reparametrized action'''
        states = batch['states']
        q1_preds = self.calc_q(states, reparam_actions, self.q1_net)
        q2_preds = self.calc_q(states, reparam_actions, self.q2_net)
        q_preds = torch.min(q1_preds, q2_preds)
        policy_loss = (self.alpha * log_probs - q_preds).mean()
        self.body.tb_tracker['policy_loss'] = policy_loss
        return policy_loss

    def calc_alpha_loss(self, log_probs):
        # clamp for safety
        self.log_alpha = self.log_alpha.clamp(-10, 10)
        alpha_loss = - (self.log_alpha * (log_probs.detach() + self.target_entropy)).mean()
        self.body.tb_tracker['alpha_loss'] = alpha_loss
        return alpha_loss

    def try_update_per(self, q_preds, q_targets):
        if 'Prioritized' in util.get_class_name(self.body.memory):  # PER
            with torch.no_grad():
                errors = (q_preds - q_targets).abs().cpu().numpy()
            self.body.memory.update_priorities(errors)

    def train_alpha(self, alpha_loss):
        '''Custom method to train the alpha variable'''
        self.alpha_lr_scheduler.step(epoch=self.body.env.clock.frame)
        self.alpha_optim.zero_grad()
        alpha_loss.backward()
        self.alpha_optim.step()
        self.alpha = self.log_alpha.detach().exp()
        self.body.tb_tracker['alpha'] = self.alpha

    def train(self):
        '''Train actor critic by computing the loss in batch efficiently'''
        if util.in_eval_lab_modes():
            return np.nan
        clock = self.body.env.clock
        if self.to_train == 1:
            for _ in range(self.training_iter):
                batch = self.sample()
                clock.set_batch_size(len(batch))

                states = batch['states']
                actions = self.guard_q_actions(batch['actions'])
                q_targets = self.calc_q_targets(batch)
                # Q-value loss for both Q nets
                q1_preds = self.calc_q(states, actions, self.q1_net)
                q1_loss = self.calc_reg_loss(q1_preds, q_targets)
                self.q1_net.train_step(q1_loss, self.q1_optim, self.q1_lr_scheduler, clock=clock, global_net=self.global_q1_net)

                q2_preds = self.calc_q(states, actions, self.q2_net)
                q2_loss = self.calc_reg_loss(q2_preds, q_targets)
                self.q2_net.train_step(q2_loss, self.q2_optim, self.q2_lr_scheduler, clock=clock, global_net=self.global_q2_net)

                # policy loss
                action_pd = policy_util.init_action_pd(self.body.ActionPD, self.calc_pdparam(states))
                log_probs, reparam_actions = self.calc_log_prob_action(action_pd, reparam=True)
                policy_loss = self.calc_policy_loss(batch, log_probs, reparam_actions)
                self.net.train_step(policy_loss, self.optim, self.lr_scheduler, clock=clock, global_net=self.global_net)

                # alpha loss
                alpha_loss = self.calc_alpha_loss(log_probs)
                self.train_alpha(alpha_loss)

                loss = q1_loss + q2_loss + policy_loss + alpha_loss
                # update target networks
                self.update_nets()
                # update PER priorities if availalbe
                self.try_update_per(torch.min(q1_preds, q2_preds), q_targets)

            # reset
            self.to_train = 0
            logger.debug(f'Trained {self.name} at epi: {clock.epi}, frame: {clock.frame}, t: {clock.t}, total_reward so far: {self.body.env.total_reward}, loss: {loss:g}')
            return loss.item()
        else:
            return np.nan

    def update_nets(self):
        '''Update target networks'''
        if util.frame_mod(self.body.env.clock.frame, self.q1_net.update_frequency, self.body.env.num_envs):
            if self.q1_net.update_type == 'replace':
                net_util.copy(self.q1_net, self.target_q1_net)
                net_util.copy(self.q2_net, self.target_q2_net)
            elif self.q1_net.update_type == 'polyak':
                net_util.polyak_update(self.q1_net, self.target_q1_net, self.q1_net.polyak_coef)
                net_util.polyak_update(self.q2_net, self.target_q2_net, self.q2_net.polyak_coef)
            else:
                raise ValueError('Unknown q1_net.update_type. Should be "replace" or "polyak". Exiting.')

    @lab_api
    def update(self):
        '''Override parent method to do nothing'''
        return self.body.explore_var
<commit_msg>clamp after train<commit_after>from slm_lab.agent import net
from slm_lab.agent.algorithm import policy_util
from slm_lab.agent.algorithm.actor_critic import ActorCritic
from slm_lab.agent.net import net_util
from slm_lab.lib import logger, util
from slm_lab.lib.decorator import lab_api
import numpy as np
import torch
import torch.nn.functional as F

logger = logger.get_logger(__name__)

class SoftActorCritic(ActorCritic):
    '''
    Implementation of Soft Actor-Critic (SAC)
    Original paper: "Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor"
    https://arxiv.org/abs/1801.01290
    Improvement of SAC paper: "Soft Actor-Critic Algorithms and Applications"
    https://arxiv.org/abs/1812.05905

    e.g. algorithm_spec
    "algorithm": {
        "name": "SoftActorCritic",
        "action_pdtype": "default",
        "action_policy": "default",
        "gamma": 0.99,
        "training_frequency": 1,
    }
    '''
    @lab_api
    def init_algorithm_params(self):
        '''Initialize other algorithm parameters'''
        # set default
        util.set_attr(self, dict(
            action_pdtype='default',
            action_policy='default',
            training_iter=self.body.env.num_envs,
            training_start_step=self.body.memory.batch_size,
        ))
        util.set_attr(self, self.algorithm_spec, [
            'action_pdtype',
            'action_policy',
            'gamma',  # the discount factor
            'training_iter',
            'training_frequency',
            'training_start_step',
        ])
        if self.body.is_discrete:
            assert self.action_pdtype == 'GumbelSoftmax'
        self.to_train = 0
        self.action_policy = getattr(policy_util, self.action_policy)

    @lab_api
    def init_nets(self, global_nets=None):
        '''
        Networks: net(actor/policy), q1_net, target_q1_net, q2_net, target_q2_net
        All networks are separate, and have the same hidden layer architectures and optim specs, so tuning is minimal
        '''
        self.shared = False  # SAC does not share networks
        NetClass = getattr(net, self.net_spec['type'])
        # main actor network
        self.net = NetClass(self.net_spec, self.body.state_dim, net_util.get_out_dim(self.body))
        self.net_names = ['net']
        # two critic Q-networks to mitigate positive bias in q_loss and speed up training, uses q_net.py with prefix Q
        QNetClass = getattr(net, 'Q' + self.net_spec['type'])
        q_in_dim = [self.body.state_dim, self.body.action_dim]
        self.q1_net = QNetClass(self.net_spec, q_in_dim, 1)
        self.target_q1_net = QNetClass(self.net_spec, q_in_dim, 1)
        self.q2_net = QNetClass(self.net_spec, q_in_dim, 1)
        self.target_q2_net = QNetClass(self.net_spec, q_in_dim, 1)
        self.net_names += ['q1_net', 'target_q1_net', 'q2_net', 'target_q2_net']
        net_util.copy(self.q1_net, self.target_q1_net)
        net_util.copy(self.q2_net, self.target_q2_net)
        # temperature variable to be learned, and its target entropy
        self.log_alpha = torch.zeros(1, requires_grad=True, device=self.net.device)
        self.alpha = self.log_alpha.detach().exp()
        if self.body.is_discrete:
            self.target_entropy = - self.body.action_space.n
        else:
            self.target_entropy = - np.product(self.body.action_space.shape)

        # init net optimizer and its lr scheduler
        self.optim = net_util.get_optim(self.net, self.net.optim_spec)
        self.lr_scheduler = net_util.get_lr_scheduler(self.optim, self.net.lr_scheduler_spec)
        self.q1_optim = net_util.get_optim(self.q1_net, self.q1_net.optim_spec)
        self.q1_lr_scheduler = net_util.get_lr_scheduler(self.q1_optim, self.q1_net.lr_scheduler_spec)
        self.q2_optim = net_util.get_optim(self.q2_net, self.q2_net.optim_spec)
        self.q2_lr_scheduler = net_util.get_lr_scheduler(self.q2_optim, self.q2_net.lr_scheduler_spec)
        self.alpha_optim = net_util.get_optim(self.log_alpha, self.net.optim_spec)
        self.alpha_lr_scheduler = net_util.get_lr_scheduler(self.alpha_optim, self.net.lr_scheduler_spec)
        net_util.set_global_nets(self, global_nets)
        self.post_init_nets()

    @lab_api
    def act(self, state):
        if self.body.env.clock.frame < self.training_start_step:
            return policy_util.random(state, self, self.body).cpu().squeeze().numpy()
        else:
            action = self.action_policy(state, self, self.body)
            if not self.body.is_discrete:
                action = self.scale_action(torch.tanh(action))  # continuous action bound
            return action.cpu().squeeze().numpy()

    def scale_action(self, action):
        '''Scale continuous actions from tanh range'''
        action_space = self.body.action_space
        low, high = torch.from_numpy(action_space.low), torch.from_numpy(action_space.high)
        return action * (high - low) / 2 + (low + high) / 2

    def guard_q_actions(self, actions):
        '''Guard to convert actions to one-hot for input to Q-network'''
        if self.body.is_discrete:
            # TODO support multi-discrete actions
            actions = F.one_hot(actions.long(), self.body.action_dim).float()
        return actions

    def calc_log_prob_action(self, action_pd, reparam=False):
        '''Calculate log_probs and actions with option to reparametrize from paper eq. 11'''
        samples = action_pd.rsample() if reparam else action_pd.sample()
        if self.body.is_discrete:  # this is straightforward using GumbelSoftmax
            actions = samples
            log_probs = action_pd.log_prob(actions)
        else:
            mus = samples
            actions = self.scale_action(torch.tanh(mus))
            if actions.dim() == 1:  # handle shape consistency for single actions
                actions = actions.unsqueeze(dim=-1)
            # paper Appendix C. Enforcing Action Bounds for continuous actions
            log_probs = (action_pd.log_prob(mus) - torch.log(1 - actions.pow(2) + 1e-6).sum(1))
        return log_probs, actions

    def calc_q(self, state, action, net):
        '''Forward-pass to calculate the predicted state-action-value from q1_net.'''
        if not self.body.is_discrete and action.dim() == 1:  # handle shape consistency for single continuous action
            action = action.unsqueeze(dim=-1)
        q_pred = net(state, action).view(-1)
        return q_pred

    def calc_q_targets(self, batch):
        '''Q_tar = r + gamma * (target_Q(s', a') - alpha * log pi(a'|s'))'''
        next_states = batch['next_states']
        with torch.no_grad():
            pdparams = self.calc_pdparam(next_states)
            action_pd = policy_util.init_action_pd(self.body.ActionPD, pdparams)
            next_log_probs, next_actions = self.calc_log_prob_action(action_pd)
            next_actions = self.guard_q_actions(next_actions)  # non-reparam discrete actions need to be converted into one-hot

            next_target_q1_preds = self.calc_q(next_states, next_actions, self.target_q1_net)
            next_target_q2_preds = self.calc_q(next_states, next_actions, self.target_q2_net)
            next_target_q_preds = torch.min(next_target_q1_preds, next_target_q2_preds)
            q_targets = batch['rewards'] + self.gamma * (1 - batch['dones']) * (next_target_q_preds - self.alpha * next_log_probs)
        return q_targets

    def calc_reg_loss(self, preds, targets):
        '''Calculate the regression loss for V and Q values, using the same loss function from net_spec'''
        assert preds.shape == targets.shape, f'{preds.shape} != {targets.shape}'
        reg_loss = self.net.loss_fn(preds, targets)
        self.body.tb_tracker['reg_loss'] = reg_loss
        return reg_loss

    def calc_policy_loss(self, batch, log_probs, reparam_actions):
        '''policy_loss = alpha * log pi(f(a)|s) - Q1(s, f(a)), where f(a) = reparametrized action'''
        states = batch['states']
        q1_preds = self.calc_q(states, reparam_actions, self.q1_net)
        q2_preds = self.calc_q(states, reparam_actions, self.q2_net)
        q_preds = torch.min(q1_preds, q2_preds)
        policy_loss = (self.alpha * log_probs - q_preds).mean()
        self.body.tb_tracker['policy_loss'] = policy_loss
        return policy_loss

    def calc_alpha_loss(self, log_probs):
        alpha_loss = - (self.log_alpha * (log_probs.detach() + self.target_entropy)).mean()
        self.body.tb_tracker['alpha_loss'] = alpha_loss
        return alpha_loss

    def try_update_per(self, q_preds, q_targets):
        if 'Prioritized' in util.get_class_name(self.body.memory):  # PER
            with torch.no_grad():
                errors = (q_preds - q_targets).abs().cpu().numpy()
            self.body.memory.update_priorities(errors)

    def train_alpha(self, alpha_loss):
        '''Custom method to train the alpha variable'''
        self.alpha_lr_scheduler.step(epoch=self.body.env.clock.frame)
        self.alpha_optim.zero_grad()
        alpha_loss.backward()
        self.alpha_optim.step()
        # clamp for safety
        self.log_alpha = self.log_alpha.clamp(-5, 4)
        self.alpha = self.log_alpha.detach().exp()
        self.body.tb_tracker['alpha'] = self.alpha

    def train(self):
        '''Train actor critic by computing the loss in batch efficiently'''
        if util.in_eval_lab_modes():
            return np.nan
        clock = self.body.env.clock
        if self.to_train == 1:
            for _ in range(self.training_iter):
                batch = self.sample()
                clock.set_batch_size(len(batch))

                states = batch['states']
                actions = self.guard_q_actions(batch['actions'])
                q_targets = self.calc_q_targets(batch)
                # Q-value loss for both Q nets
                q1_preds = self.calc_q(states, actions, self.q1_net)
                q1_loss = self.calc_reg_loss(q1_preds, q_targets)
                self.q1_net.train_step(q1_loss, self.q1_optim, self.q1_lr_scheduler, clock=clock, global_net=self.global_q1_net)

                q2_preds = self.calc_q(states, actions, self.q2_net)
                q2_loss = self.calc_reg_loss(q2_preds, q_targets)
                self.q2_net.train_step(q2_loss, self.q2_optim, self.q2_lr_scheduler, clock=clock, global_net=self.global_q2_net)

                # policy loss
                action_pd = policy_util.init_action_pd(self.body.ActionPD, self.calc_pdparam(states))
                log_probs, reparam_actions = self.calc_log_prob_action(action_pd, reparam=True)
                policy_loss = self.calc_policy_loss(batch, log_probs, reparam_actions)
                self.net.train_step(policy_loss, self.optim, self.lr_scheduler, clock=clock, global_net=self.global_net)

                # alpha loss
                alpha_loss = self.calc_alpha_loss(log_probs)
                self.train_alpha(alpha_loss)

                loss = q1_loss + q2_loss + policy_loss + alpha_loss
                # update target networks
                self.update_nets()
                # update PER priorities if availalbe
                self.try_update_per(torch.min(q1_preds, q2_preds), q_targets)

            # reset
            self.to_train = 0
            logger.debug(f'Trained {self.name} at epi: {clock.epi}, frame: {clock.frame}, t: {clock.t}, total_reward so far: {self.body.env.total_reward}, loss: {loss:g}')
            return loss.item()
        else:
            return np.nan

    def update_nets(self):
        '''Update target networks'''
        if util.frame_mod(self.body.env.clock.frame, self.q1_net.update_frequency, self.body.env.num_envs):
            if self.q1_net.update_type == 'replace':
                net_util.copy(self.q1_net, self.target_q1_net)
                net_util.copy(self.q2_net, self.target_q2_net)
            elif self.q1_net.update_type == 'polyak':
                net_util.polyak_update(self.q1_net, self.target_q1_net, self.q1_net.polyak_coef)
                net_util.polyak_update(self.q2_net, self.target_q2_net, self.q2_net.polyak_coef)
            else:
                raise ValueError('Unknown q1_net.update_type. Should be "replace" or "polyak". Exiting.')

    @lab_api
    def update(self):
        '''Override parent method to do nothing'''
        return self.body.explore_var
<|endoftext|>


<commit_before># pylint: disable=R0903

from __future__ import unicode_literals, print_function

import re

from pypeg2 import Keyword, maybe_some, some, optional, Literal, omit

class Rule(object):

    def __init__(self, args):
        self.children = args

    def accept(self, visitor):
        return visitor.visit(self, *[arg.accept(visitor) for arg in self.children])

    def __eq__(self, other):
        return type(self) == type(other) and self.children == other.args

    def __repr__(self):
        return "%s(%s)" % (self.__class__.__name__, repr(self.children))

class UnaryRule(object):

    def __init__(self, child):
        self.child = child

    def accept(self, visitor):
        return visitor.visit(self, self.child.accept(visitor))

    def __eq__(self, other):
        return type(self) == type(other) and self.child == other.child

    def __repr__(self):
        return "%s(%s)" % (self.__class__.__name__, repr(self.child))

class LeafRule(object):

    def __init__(self, value):
        self.value = value

    def accept(self, visitor):
        return visitor.visit(self, self.value)

    def __eq__(self, other):
        return type(self) == type(other) and self.value == other.value

    def __repr__(self):
        return '%s(%s)' % (self.__class__.__name__, repr(self.value))

class EmptyLeafRule(object):

    def accept(self, visitor):
        return visitor.visit(self)

    def __eq__(self, other):
        return type(self) == type(other)

    def __repr__(self):
        return '%s' % self.__class__.__name__

"""
Misc.
"""

class Content(LeafRule):
    grammar = re.compile(r'[^{]+')

class CommentStatement(LeafRule):
    grammar = re.compile(r"{\*.*?\*}", re.S)

class LiteralStatement(LeafRule):
    grammar = re.compile("{literal}.*?{/literal}", re.S)

class Whitespace(object):
    grammar = maybe_some([Literal(' '), Literal('\n'), Literal('\t')])

_ = omit(Whitespace)

class Identifier(LeafRule):
    grammar = re.compile(r'[\w\-\+\*\/]*\w')

"""
Logical Operators.
"""

class AndOperator(EmptyLeafRule):
    grammar = [Literal('and'), Literal('&&')]

class OrOperator(EmptyLeafRule):
    grammar = [Literal('or'), Literal('||')]

class EqOperator(EmptyLeafRule):
    grammar = [Literal('=='), Literal('eq')]

class NeOperator(EmptyLeafRule):
    grammar = [Literal('!='), Literal('ne'), Literal('neq')]

class GtOperator(EmptyLeafRule):
    grammar = [Literal('>'), Literal('gt')]

class LtOperator(EmptyLeafRule):
    grammar = [Literal('<'), Literal('lt')]

class LteOperator(EmptyLeafRule):
    grammar = Literal('<=')

class GteOperator(EmptyLeafRule):
    grammar = Literal('>=')

class IsOperator(EmptyLeafRule):
    grammar = Literal('instanceof')

class RightParen(EmptyLeafRule):
    grammar = Literal(')')

class LeftParen(EmptyLeafRule):
    grammar = Literal('(')

class Operator(UnaryRule):
    grammar = [AndOperator, EqOperator, GteOperator, LteOperator,
               LtOperator, GtOperator, NeOperator, OrOperator, IsOperator]

"""
Smarty variables.
"""

class SingleQuotedString(LeafRule):
    grammar = "'", re.compile(r"([^']|\\.)*"), "'"

class DoubleQuotedString(LeafRule):
    grammar = '"', re.compile(r'([^"\$]|\\.)*'), '"'

class String(UnaryRule):
    grammar = [SingleQuotedString, DoubleQuotedString]

class Text(LeafRule):
    grammar = some([re.compile(r'[^$`"\\]+'), re.compile(r'\\.')])

class NotOperator(EmptyLeafRule):
    grammar = Literal('!')

class AtOperator(EmptyLeafRule):
    grammar = Literal('@')

class EmptyOperator(EmptyLeafRule):
    grammar = Literal('')

class Variable(UnaryRule):
    grammar = Literal('$'), Identifier

class Symbol(Rule):
    grammar = ([NotOperator, AtOperator, EmptyOperator],
               [Variable, Identifier])

class DollarSymbol(Rule):
    grammar = [NotOperator, AtOperator, EmptyOperator], Variable

class Expression(UnaryRule):
    pass

class Array(Rule):
    grammar = Symbol, "[", optional(Expression), "]"

class ObjectDereference(Rule):
    grammar = [Array, Symbol], ['.', '->'], Expression

class VariableString(Rule):
    grammar = '"', some([Text, ('`', Expression, '`'), ('$', Expression)]), '"'

class ExpNoModifier(UnaryRule):
    grammar = Literal(':'), [ObjectDereference, Array, Symbol, VariableString, String]

class ModifierParameters(Rule):
    grammar = some(ExpNoModifier)

class ModifierElement(Rule):
    grammar = ('|',
               [AtOperator, EmptyOperator],
               Identifier,
               optional(ModifierParameters))

class ModifierRight(Rule):
    grammar = some(ModifierElement)

class Modifier(Rule):
    grammar = [ObjectDereference, Array, Symbol,
               VariableString, String], ModifierRight

class FuncParams(Rule):
    grammar = Expression, some((',', _, Expression))

class FuncCall(Rule):
    grammar = Identifier, omit(LeftParen), optional([FuncParams, Expression]), omit(RightParen)

Expression.grammar = [FuncCall, Modifier, ObjectDereference,
                      Array, Symbol, String, VariableString]

"""
Smarty Statements.
"""

class SmartyLanguage(Rule):
    pass

class ElseStatement(UnaryRule):
    grammar = '{', Keyword('else'), '}', SmartyLanguage

class ForeachelseStatement(UnaryRule):
    grammar = '{', Keyword('foreachelse'), '}', SmartyLanguage

class NoFilter(EmptyLeafRule):
    grammar = Literal('nofilter')

class PrintStatement(Rule):
    grammar = ('{', _,
               optional('e '),
               [FuncCall, Modifier, ObjectDereference,
                Array, DollarSymbol, String, VariableString],
               optional(_, NoFilter),
               _, '}')

class FunctionParameter(Rule):
    grammar = Symbol, '=', Expression

class FunctionStatement(Rule):
    grammar = '{', _, Symbol, some(_, FunctionParameter), _, '}'

class ForFrom(UnaryRule):
    grammar = Keyword('from'), '=', omit(optional(['"', '\''])), Expression, omit(optional(['"', '\'']))

class ForItem(UnaryRule):
    grammar = Keyword('item'), '=', omit(optional(['"', '\''])), Symbol, omit(optional(['"', '\'']))

class ForName(UnaryRule):
    grammar = Keyword('name'), '=', omit(optional(['"', '\''])), Symbol, omit(optional(['"', '\'']))

class ForKey(UnaryRule):
    grammar = Keyword('key'), '=', omit(optional(['"', '\''])), Symbol, omit(optional(['"', '\'']))

class IfCondition(Rule):
    grammar = maybe_some(LeftParen), _, Expression, _, maybe_some(RightParen)

class IfConditionList(Rule):
    grammar = IfCondition, some(_, Operator, _, IfCondition)

class ElseifStatement(Rule):
    grammar = ('{', Keyword('elseif'), _, [IfConditionList, IfCondition], _, '}', SmartyLanguage)

class IfMoreStatement(Rule):
    grammar = some([ElseStatement, ElseifStatement])

class IfStatement(Rule):
    grammar = ('{', _, Keyword('if'), _, [IfConditionList, IfCondition], _, '}',
               SmartyLanguage, optional(IfMoreStatement), '{/', Keyword('if'), '}')

class ForeachArray(Rule):
    grammar = _, Expression, _, 'as', _, Symbol

class AddOperator(EmptyLeafRule):
    grammar = '+'

class SubOperator(EmptyLeafRule):
    grammar = '-'

class MultOperator(EmptyLeafRule):
    grammar = '*'

class DivOperator(EmptyLeafRule):
    grammar = '/'

class ArithmeticOperator(UnaryRule):
    grammar = [AddOperator, SubOperator, MultOperator, DivOperator]

class Number(LeafRule):
    grammar = re.compile(r'\d+')

class ForExpression(Rule):
    grammar = ArithmeticOperator, Number

class ForVariableIdentifier(LeafRule):
    grammar = re.compile(r'\w+')

class ForVariable(Rule):
    grammar = '{', _, Variable, '@', ForVariableIdentifier, optional(_, ForExpression), _, '}'

class ForeachParameters(Rule):
    grammar = some(_, [ForFrom, ForItem, ForName, ForKey])

class ForContent(Rule):
    grammar = some([ForVariable, SmartyLanguage])

class ForStatement(Rule):
    grammar = ('{', _, Keyword('foreach'), [ForeachParameters, ForeachArray], _, '}',
               ForContent, optional(ForeachelseStatement),
               '{/', Keyword('foreach'), '}')

class IsLink(LeafRule):
    grammar = Literal('quoted='), re.compile('true|false')

class TranslationStatement(Rule):
    grammar = ('{', _, Keyword('t'), _,
               Literal('id='), Expression,
               optional(_, IsLink),
               _, '}')

class AssignStatement(Rule):
    grammar = ('{', _, Keyword('assign'), _,
               Literal('var='), Identifier, _,
               Literal('value='), Expression,
               _, '}')

class LeftDelimTag(EmptyLeafRule):
    grammar = '{ldelim}'

class RightDelimTag(EmptyLeafRule):
    grammar = '{rdelim}'

class LeftDelim(EmptyLeafRule):
    grammar = '{'

class IncludeStatement(UnaryRule):
    grammar = '{', _, Keyword('include'), _, Literal('file='), DoubleQuotedString, _, '}'

"""
Finally, the actual language description.
"""

SmartyLanguage.grammar = some([LiteralStatement, TranslationStatement,
                              IfStatement, ForStatement, IncludeStatement,
                              AssignStatement,
                              FunctionStatement, CommentStatement,
                              PrintStatement, Content,
                              LeftDelimTag, RightDelimTag])

class SmartyLanguageMain(Rule):
    grammar = some([LiteralStatement, TranslationStatement,
                    IfStatement, ForStatement, IncludeStatement,
                    AssignStatement,
                    FunctionStatement, CommentStatement,
                    PrintStatement, Content,
                    LeftDelimTag, RightDelimTag, LeftDelim])

class SmartyLanguageMainOrEmpty(UnaryRule):
    grammar = [SmartyLanguageMain, EmptyOperator]
<commit_msg>Handles function calls without parameters<commit_after># pylint: disable=R0903

from __future__ import unicode_literals, print_function

import re

from pypeg2 import Keyword, maybe_some, some, optional, Literal, omit

class Rule(object):

    def __init__(self, args):
        self.children = args

    def accept(self, visitor):
        return visitor.visit(self, *[arg.accept(visitor) for arg in self.children])

    def __eq__(self, other):
        return type(self) == type(other) and self.children == other.args

    def __repr__(self):
        return "%s(%s)" % (self.__class__.__name__, repr(self.children))

class UnaryRule(object):

    def __init__(self, child):
        self.child = child

    def accept(self, visitor):
        return visitor.visit(self, self.child.accept(visitor))

    def __eq__(self, other):
        return type(self) == type(other) and self.child == other.child

    def __repr__(self):
        return "%s(%s)" % (self.__class__.__name__, repr(self.child))

class LeafRule(object):

    def __init__(self, value):
        self.value = value

    def accept(self, visitor):
        return visitor.visit(self, self.value)

    def __eq__(self, other):
        return type(self) == type(other) and self.value == other.value

    def __repr__(self):
        return '%s(%s)' % (self.__class__.__name__, repr(self.value))

class EmptyLeafRule(object):

    def accept(self, visitor):
        return visitor.visit(self)

    def __eq__(self, other):
        return type(self) == type(other)

    def __repr__(self):
        return '%s' % self.__class__.__name__

"""
Misc.
"""

class Content(LeafRule):
    grammar = re.compile(r'[^{]+')

class CommentStatement(LeafRule):
    grammar = re.compile(r"{\*.*?\*}", re.S)

class LiteralStatement(LeafRule):
    grammar = re.compile("{literal}.*?{/literal}", re.S)

class Whitespace(object):
    grammar = maybe_some([Literal(' '), Literal('\n'), Literal('\t')])

_ = omit(Whitespace)

class Identifier(LeafRule):
    grammar = re.compile(r'[\w\-\+\*\/]*\w')

"""
Logical Operators.
"""

class AndOperator(EmptyLeafRule):
    grammar = [Literal('and'), Literal('&&')]

class OrOperator(EmptyLeafRule):
    grammar = [Literal('or'), Literal('||')]

class EqOperator(EmptyLeafRule):
    grammar = [Literal('=='), Literal('eq')]

class NeOperator(EmptyLeafRule):
    grammar = [Literal('!='), Literal('ne'), Literal('neq')]

class GtOperator(EmptyLeafRule):
    grammar = [Literal('>'), Literal('gt')]

class LtOperator(EmptyLeafRule):
    grammar = [Literal('<'), Literal('lt')]

class LteOperator(EmptyLeafRule):
    grammar = Literal('<=')

class GteOperator(EmptyLeafRule):
    grammar = Literal('>=')

class IsOperator(EmptyLeafRule):
    grammar = Literal('instanceof')

class RightParen(EmptyLeafRule):
    grammar = Literal(')')

class LeftParen(EmptyLeafRule):
    grammar = Literal('(')

class Operator(UnaryRule):
    grammar = [AndOperator, EqOperator, GteOperator, LteOperator,
               LtOperator, GtOperator, NeOperator, OrOperator, IsOperator]

"""
Smarty variables.
"""

class SingleQuotedString(LeafRule):
    grammar = "'", re.compile(r"([^']|\\.)*"), "'"

class DoubleQuotedString(LeafRule):
    grammar = '"', re.compile(r'([^"\$]|\\.)*'), '"'

class String(UnaryRule):
    grammar = [SingleQuotedString, DoubleQuotedString]

class Text(LeafRule):
    grammar = some([re.compile(r'[^$`"\\]+'), re.compile(r'\\.')])

class NotOperator(EmptyLeafRule):
    grammar = Literal('!')

class AtOperator(EmptyLeafRule):
    grammar = Literal('@')

class EmptyOperator(EmptyLeafRule):
    grammar = Literal('')

class Variable(UnaryRule):
    grammar = Literal('$'), Identifier

class Symbol(Rule):
    grammar = ([NotOperator, AtOperator, EmptyOperator],
               [Variable, Identifier])

class DollarSymbol(Rule):
    grammar = [NotOperator, AtOperator, EmptyOperator], Variable

class Expression(UnaryRule):
    pass

class Array(Rule):
    grammar = Symbol, "[", optional(Expression), "]"

class ObjectDereference(Rule):
    grammar = [Array, Symbol], ['.', '->'], Expression

class VariableString(Rule):
    grammar = '"', some([Text, ('`', Expression, '`'), ('$', Expression)]), '"'

class ExpNoModifier(UnaryRule):
    grammar = Literal(':'), [ObjectDereference, Array, Symbol, VariableString, String]

class ModifierParameters(Rule):
    grammar = some(ExpNoModifier)

class ModifierElement(Rule):
    grammar = ('|',
               [AtOperator, EmptyOperator],
               Identifier,
               optional(ModifierParameters))

class ModifierRight(Rule):
    grammar = some(ModifierElement)

class Modifier(Rule):
    grammar = [ObjectDereference, Array, Symbol,
               VariableString, String], ModifierRight

class FuncParams(Rule):
    grammar = Expression, some((',', _, Expression))

class FuncCall(Rule):
    grammar = Identifier, optional(omit(LeftParen), optional([FuncParams, Expression]), omit(RightParen))

Expression.grammar = [FuncCall, Modifier, ObjectDereference,
                      Array, Symbol, String, VariableString]

"""
Smarty Statements.
"""

class SmartyLanguage(Rule):
    pass

class ElseStatement(UnaryRule):
    grammar = '{', Keyword('else'), '}', SmartyLanguage

class ForeachelseStatement(UnaryRule):
    grammar = '{', Keyword('foreachelse'), '}', SmartyLanguage

class NoFilter(EmptyLeafRule):
    grammar = Literal('nofilter')

class PrintStatement(Rule):
    grammar = ('{', _,
               optional('e '),
               [FuncCall, Modifier, ObjectDereference,
                Array, DollarSymbol, String, VariableString],
               optional(_, NoFilter),
               _, '}')

class FunctionParameter(Rule):
    grammar = Symbol, '=', Expression

class FunctionStatement(Rule):
    grammar = '{', _, Symbol, some(_, FunctionParameter), _, '}'

class ForFrom(UnaryRule):
    grammar = Keyword('from'), '=', omit(optional(['"', '\''])), Expression, omit(optional(['"', '\'']))

class ForItem(UnaryRule):
    grammar = Keyword('item'), '=', omit(optional(['"', '\''])), Symbol, omit(optional(['"', '\'']))

class ForName(UnaryRule):
    grammar = Keyword('name'), '=', omit(optional(['"', '\''])), Symbol, omit(optional(['"', '\'']))

class ForKey(UnaryRule):
    grammar = Keyword('key'), '=', omit(optional(['"', '\''])), Symbol, omit(optional(['"', '\'']))

class IfCondition(Rule):
    grammar = maybe_some(LeftParen), _, Expression, _, maybe_some(RightParen)

class IfConditionList(Rule):
    grammar = IfCondition, some(_, Operator, _, IfCondition)

class ElseifStatement(Rule):
    grammar = ('{', Keyword('elseif'), _, [IfConditionList, IfCondition], _, '}', SmartyLanguage)

class IfMoreStatement(Rule):
    grammar = some([ElseStatement, ElseifStatement])

class IfStatement(Rule):
    grammar = ('{', _, Keyword('if'), _, [IfConditionList, IfCondition], _, '}',
               SmartyLanguage, optional(IfMoreStatement), '{/', Keyword('if'), '}')

class ForeachArray(Rule):
    grammar = _, Expression, _, 'as', _, Symbol

class AddOperator(EmptyLeafRule):
    grammar = '+'

class SubOperator(EmptyLeafRule):
    grammar = '-'

class MultOperator(EmptyLeafRule):
    grammar = '*'

class DivOperator(EmptyLeafRule):
    grammar = '/'

class ArithmeticOperator(UnaryRule):
    grammar = [AddOperator, SubOperator, MultOperator, DivOperator]

class Number(LeafRule):
    grammar = re.compile(r'\d+')

class ForExpression(Rule):
    grammar = ArithmeticOperator, Number

class ForVariableIdentifier(LeafRule):
    grammar = re.compile(r'\w+')

class ForVariable(Rule):
    grammar = '{', _, Variable, '@', ForVariableIdentifier, optional(_, ForExpression), _, '}'

class ForeachParameters(Rule):
    grammar = some(_, [ForFrom, ForItem, ForName, ForKey])

class ForContent(Rule):
    grammar = some([ForVariable, SmartyLanguage])

class ForStatement(Rule):
    grammar = ('{', _, Keyword('foreach'), [ForeachParameters, ForeachArray], _, '}',
               ForContent, optional(ForeachelseStatement),
               '{/', Keyword('foreach'), '}')

class IsLink(LeafRule):
    grammar = Literal('quoted='), re.compile('true|false')

class TranslationStatement(Rule):
    grammar = ('{', _, Keyword('t'), _,
               Literal('id='), Expression,
               optional(_, IsLink),
               _, '}')

class AssignStatement(Rule):
    grammar = ('{', _, Keyword('assign'), _,
               Literal('var='), Identifier, _,
               Literal('value='), Expression,
               _, '}')

class LeftDelimTag(EmptyLeafRule):
    grammar = '{ldelim}'

class RightDelimTag(EmptyLeafRule):
    grammar = '{rdelim}'

class LeftDelim(EmptyLeafRule):
    grammar = '{'

class IncludeStatement(UnaryRule):
    grammar = '{', _, Keyword('include'), _, Literal('file='), DoubleQuotedString, _, '}'

"""
Finally, the actual language description.
"""

SmartyLanguage.grammar = some([LiteralStatement, TranslationStatement,
                              IfStatement, ForStatement, IncludeStatement,
                              AssignStatement,
                              FunctionStatement, CommentStatement,
                              PrintStatement, Content,
                              LeftDelimTag, RightDelimTag])

class SmartyLanguageMain(Rule):
    grammar = some([LiteralStatement, TranslationStatement,
                    IfStatement, ForStatement, IncludeStatement,
                    AssignStatement,
                    FunctionStatement, CommentStatement,
                    PrintStatement, Content,
                    LeftDelimTag, RightDelimTag, LeftDelim])

class SmartyLanguageMainOrEmpty(UnaryRule):
    grammar = [SmartyLanguageMain, EmptyOperator]
<|endoftext|>