Calculate loss of the supplied rollout
330579
def calculate_gradient(self, batch_info, device, model, rollout):
        
        assert isinstance(rollout, Trajectories), "ACER algorithm requires trajectory input"

        local_epsilon = 1e-6

        evaluator = model.evaluate(rollout)

        actions = evaluator.get('rollout:actions')
        rollout_probabilities = torch.exp(evaluator.get('rollout:logprobs'))

        # We calculate the trust-region update with respect to the average model
        if self.trust_region:
            self.update_average_model(model)

        logprobs = evaluator.get('model:logprobs')
        q = evaluator.get('model:q')

        # Selected action values
        action_logprobs = select_indices(logprobs, actions)
        action_q = select_indices(q, actions)

        # We only want to propagate gradients through specific variables
        with torch.no_grad():
            model_probabilities = torch.exp(logprobs)

            # Importance sampling correction - we must find the quotient of probabilities
            rho = model_probabilities / (rollout_probabilities + local_epsilon)

            # Probability quotient only for selected actions
            actions_rho = select_indices(rho, actions)

            # Calculate policy state values
            model_state_values = (model_probabilities * q).sum(dim=1)

            trajectory_rewards = rollout.transition_tensors['rewards']
            trajectory_dones = rollout.transition_tensors['dones']

            q_retraced = self.retrace(
                trajectory_rewards,
                trajectory_dones,
                action_q.reshape(trajectory_rewards.size()),
                model_state_values.reshape(trajectory_rewards.size()),
                actions_rho.reshape(trajectory_rewards.size()),
                rollout.rollout_tensors['final_values']
            ).flatten()

            advantages = q_retraced - model_state_values
            importance_sampling_coefficient = torch.min(actions_rho, self.rho_cap * torch.ones_like(actions_rho))

            explained_variance = 1 - torch.var(q_retraced - action_q) / torch.var(q_retraced)

        # Entropy of the policy distribution
        policy_entropy = torch.mean(model.entropy(logprobs))
        policy_gradient_loss = -torch.mean(advantages * importance_sampling_coefficient * action_logprobs)

        # Policy gradient bias correction
        with torch.no_grad():
            advantages_bias_correction = q - model_state_values.view(model_probabilities.size(0), 1)
            bias_correction_coefficient = F.relu(1.0 - self.rho_cap / (rho + local_epsilon))

        # This sum is an expectation with respect to action probabilities according to model policy
        policy_gradient_bias_correction_gain = torch.sum(
            logprobs * bias_correction_coefficient * advantages_bias_correction * model_probabilities,
            dim=1
        )

        policy_gradient_bias_correction_loss = - torch.mean(policy_gradient_bias_correction_gain)

        policy_loss = policy_gradient_loss + policy_gradient_bias_correction_loss

        q_function_loss = 0.5 * F.mse_loss(action_q, q_retraced)

        if self.trust_region:
            with torch.no_grad():
                average_evaluator = self.average_model.evaluate(rollout)
                average_action_logits = average_evaluator.get('model:logprobs')

            actor_loss = policy_loss - self.entropy_coefficient * policy_entropy
            q_loss = self.q_coefficient * q_function_loss

            actor_gradient = torch.autograd.grad(-actor_loss, logprobs, retain_graph=True)[0]

            # kl_divergence = model.kl_divergence(average_action_logits, action_logits).mean()
            # kl_divergence_grad = torch.autograd.grad(kl_divergence, action_logits, retain_graph=True)

            # Analytically calculated derivative of KL divergence on logits
            # That makes it hardcoded for discrete action spaces
            kl_divergence_grad_symbolic = - torch.exp(average_action_logits) / logprobs.size(0)

            k_dot_g = (actor_gradient * kl_divergence_grad_symbolic).sum(dim=-1)
            k_dot_k = (kl_divergence_grad_symbolic ** 2).sum(dim=-1)

            adjustment = (k_dot_g - self.trust_region_delta) / k_dot_k
            adjustment_clipped = adjustment.clamp(min=0.0)

            actor_gradient_updated = actor_gradient - adjustment_clipped.view(adjustment_clipped.size(0), 1)

            # Populate gradient from the newly updated fn
            logprobs.backward(gradient=-actor_gradient_updated, retain_graph=True)
            q_loss.backward(retain_graph=True)
        else:
            # Just populate gradient from the loss
            loss = policy_loss + self.q_coefficient * q_function_loss - self.entropy_coefficient * policy_entropy

            loss.backward()

        return {
            'policy_loss': policy_loss.item(),
            'policy_gradient_loss': policy_gradient_loss.item(),
            'policy_gradient_bias_correction': policy_gradient_bias_correction_loss.item(),
            'avg_q_selected': action_q.mean().item(),
            'avg_q_retraced': q_retraced.mean().item(),
            'q_loss': q_function_loss.item(),
            'policy_entropy': policy_entropy.item(),
            'advantage_norm': torch.norm(advantages).item(),
            'explained_variance': explained_variance.item(),
            'model_prob_std': model_probabilities.std().item(),
            'rollout_prob_std': rollout_probabilities.std().item()
        }
557098
def rollback(self, date):
        """"""
        if self.onOffset(date):
            return date
        else:
            return date - QuarterEnd(month=self.month)
459802
def combined_loss(new_params,
                  old_params,
                  policy_and_value_net_apply,
                  padded_observations,
                  padded_actions,
                  padded_rewards,
                  reward_mask,
                  gamma=0.99,
                  lambda_=0.95,
                  epsilon=0.2,
                  c1=1.0,
                  c2=0.01):
  """"""
  log_probab_actions_new, value_predictions = policy_and_value_net_apply(
      padded_observations, new_params)

  log_probab_actions_old, _ = policy_and_value_net_apply(
      padded_observations, old_params)

  # (combined_loss, ppo_loss, value_loss, entropy_bonus)
  return combined_loss_given_predictions(log_probab_actions_new,
                                         log_probab_actions_old,
                                         value_predictions,
                                         padded_actions,
                                         padded_rewards,
                                         reward_mask,
                                         c1=c1,
                                         c2=c2,
                                         gamma=gamma,
                                         lambda_=lambda_,
                                         epsilon=epsilon)
616655
def get_loss(self, logits: mx.sym.Symbol, labels: mx.sym.Symbol) -> mx.sym.Symbol:
        
        if self.loss_config.normalization_type == C.LOSS_NORM_VALID:
            normalization = "valid"
        elif self.loss_config.normalization_type == C.LOSS_NORM_BATCH:
            normalization = "null"
        else:
            raise ValueError("Unknown loss normalization type: %s" % self.loss_config.normalization_type)
        return mx.sym.SoftmaxOutput(data=logits,
                                    label=labels,
                                    ignore_label=self.ignore_label,
                                    use_ignore=True,
                                    normalization=normalization,
                                    smooth_alpha=self.loss_config.label_smoothing,
                                    name=self.name)
737409
def rollout(self, batch_info: BatchInfo, model: RlModel, number_of_steps: int) -> Rollout:
        
        assert not model.is_recurrent, "Replay env roller does not support recurrent models"

        accumulator = TensorAccumulator()
        episode_information = []  # List of dictionaries with episode information

        for step_idx in range(number_of_steps):
            step = model.step(self.last_observation)

            replay_extra_information = {}

            accumulator.add('observations', self.last_observation_cpu)

            # Add step to the tensor accumulator
            for name, tensor in step.items():
                tensor_cpu = tensor.cpu()
                accumulator.add(name, tensor_cpu)

                if name != 'actions':
                    replay_extra_information[name] = tensor_cpu.numpy()

            actions_numpy = step['actions'].detach().cpu().numpy()
            new_obs, new_rewards, new_dones, new_infos = self.environment.step(actions_numpy)

            # Store rollout in the experience replay buffer
            self.replay_buffer.store_transition(
                frame=self.last_observation_cpu.numpy(),
                action=actions_numpy,
                reward=new_rewards,
                done=new_dones,
                extra_info=replay_extra_information
            )

            # Done is flagged true when the episode has ended AND the frame we see is already a first frame from the
            # next episode

            dones_tensor = torch.from_numpy(new_dones.astype(np.float32)).clone()
            accumulator.add('dones', dones_tensor)

            self.last_observation_cpu = torch.from_numpy(new_obs).clone()
            self.last_observation = self.last_observation_cpu.to(self.device)
            accumulator.add('rewards', torch.from_numpy(new_rewards.astype(np.float32)).clone())

            episode_information.append(new_infos)

        accumulated_tensors = accumulator.result()

        return Trajectories(
            num_steps=accumulated_tensors['observations'].size(0),
            num_envs=accumulated_tensors['observations'].size(1),
            environment_information=episode_information,
            transition_tensors=accumulated_tensors,
            rollout_tensors={
                'final_values': model.value(self.last_observation).cpu()
            }
        )
252205
def extract_src_loss_table(dstore, loss_type):
    
    oq = dstore['oqparam']
    li = oq.lti[loss_type]
    source_ids = dstore['source_info']['source_id']
    idxs = dstore['ruptures'].value[['srcidx', 'grp_id']]
    losses = dstore['rup_loss_table'][:, li]
    slt = numpy.zeros(len(source_ids), [('grp_id', U32), (loss_type, F32)])
    for loss, (srcidx, grp_id) in zip(losses, idxs):
        slt[srcidx][loss_type] += loss
        slt[srcidx]['grp_id'] = grp_id
    slt = util.compose_arrays(source_ids, slt, 'source_id')
    slt.sort(order=loss_type)
    return slt[::-1]
719306
def rollout(env, acts):
    
    total_rew = 0
    env.reset()
    steps = 0
    for act in acts:
        _obs, rew, done, _info = env.step(act)
        steps += 1
        total_rew += rew
        if done:
            break

    return steps, total_rew
294955
def on_loss_begin(self, last_output:Tuple[Tensor,Tensor,Tensor], **kwargs):
        ""
        self.raw_out,self.out = last_output[1],last_output[2]
        return {'last_output': last_output[0]}
1250509
def rollout(self, batch_info: BatchInfo, model: Model, number_of_steps: int) -> Rollout:
        
        accumulator = TensorAccumulator()
        episode_information = []  # List of dictionaries with episode information

        if self.hidden_state is None and model.is_recurrent:
            self.hidden_state = model.zero_state(self.last_observation.size(0)).to(self.device)

        # Remember rollout initial state, we'll use that for training as well
        initial_hidden_state = self.hidden_state

        for step_idx in range(number_of_steps):
            if model.is_recurrent:
                step = model.step(self.last_observation.to(self.device), state=self.hidden_state)
                self.hidden_state = step['state']
            else:
                step = model.step(self.last_observation.to(self.device))

            # Add step to the tensor accumulator
            for name, tensor in step.items():
                accumulator.add(name, tensor.cpu())

            accumulator.add('observations', self.last_observation)

            actions_numpy = step['actions'].detach().cpu().numpy()
            new_obs, new_rewards, new_dones, new_infos = self.environment.step(actions_numpy)

            # Done is flagged true when the episode has ended AND the frame we see is already a first frame from the
            # next episode
            dones_tensor = torch.from_numpy(new_dones.astype(np.float32)).clone()
            accumulator.add('dones', dones_tensor)

            self.last_observation = torch.from_numpy(new_obs).clone()

            if model.is_recurrent:
                # Zero out state in environments that have finished
                self.hidden_state = self.hidden_state * (1.0 - dones_tensor.unsqueeze(-1)).to(self.device)

            accumulator.add('rewards', torch.from_numpy(new_rewards.astype(np.float32)).clone())

            episode_information.append(new_infos)

        if model.is_recurrent:
            final_values = model.value(self.last_observation.to(self.device), state=self.hidden_state).cpu()
        else:
            final_values = model.value(self.last_observation.to(self.device)).cpu()

        accumulated_tensors = accumulator.result()

        return Trajectories(
            num_steps=accumulated_tensors['observations'].size(0),
            num_envs=accumulated_tensors['observations'].size(1),
            environment_information=episode_information,
            transition_tensors=accumulated_tensors,
            rollout_tensors={
                'initial_hidden_state': initial_hidden_state.cpu() if initial_hidden_state is not None else None,
                'final_values': final_values
            }
        )
1097200
def compute_loss_curves_maps(filename, builder, rlzi, monitor):
    
    with datastore.read(filename) as dstore:
        rlzs = dstore['losses_by_event']['rlzi']
        losses = dstore['losses_by_event'][rlzs == rlzi]['loss']
    return rlzi, builder.build_curves_maps(losses, rlzi)
1350233
def rollback             (self, n=None):
		""""""
		
		if n is None:
			if self.haveSnapshots: return self.fromSnapshot(self.latestLink)
			else:                  return self.fromScratch()
		elif isinstance(n, int):
			loadSnapshotPath = self.getFullPathToSnapshot(n)
			assert(os.path.isdir(loadSnapshotPath))
			return self.__markLatest(n).fromSnapshot(loadSnapshotPath)
		else:
			raise ValueError("n must be int, or None!")
1196548
def value_loss(self, model, observations, discounted_rewards):
        
        value_outputs = model.value(observations)
        value_loss = 0.5 * F.mse_loss(value_outputs, discounted_rewards)
        return value_loss
1030987
def export_agg_losses(ekey, dstore):
    
    dskey = ekey[0]
    oq = dstore['oqparam']
    dt = oq.loss_dt()
    name, value, tags = _get_data(dstore, dskey, oq.hazard_stats().items())
    writer = writers.CsvWriter(fmt=writers.FIVEDIGITS)
    expvalue = dstore['exposed_value'].value  # shape (T1, T2, ..., L)
    tagcol = dstore['assetcol/tagcol']
    tagnames = tuple(dstore['oqparam'].aggregate_by)
    header = ('loss_type',) + tagnames + (
        'loss_value', 'exposed_value', 'loss_ratio')
    for r, tag in enumerate(tags):
        rows = []
        for multi_idx, loss in numpy.ndenumerate(value[:, r]):
            l, *tagidxs = multi_idx
            evalue = expvalue[tuple(tagidxs) + (l,)]
            row = tagcol.get_tagvalues(tagnames, tagidxs) + (
                loss, evalue, loss / evalue)
            rows.append((dt.names[l],) + row)
        dest = dstore.build_fname(name, tag, 'csv')
        writer.save(rows, dest, header)
    return writer.getsaved()
1354771
def get_rollout_from_id(self, rollout_id):
    
    layer = self.rollout_id_map.get(rollout_id)

    if layer:
      return layer

    self.logger.error('Rollout with ID "%s" is not in datafile.' % rollout_id)
    return None
1022756
function rollback(client, done, error, callback) {
  client.query('ROLLBACK', rollbackError => {
    done(rollbackError);
    return callback(rollbackError || error, null);
  });
}
1294384
def calculate_gradient(self, batch_info, device, model, rollout):
        
        evaluator = model.evaluate(rollout)
        batch_size = rollout.frames()

        dones_tensor = evaluator.get('rollout:dones')
        rewards_tensor = evaluator.get('rollout:rewards')

        assert dones_tensor.dtype == torch.float32

        with torch.no_grad():
            target_evaluator = self.target_model.evaluate(rollout)

            if self.double_dqn:
                # DOUBLE DQN
                # Histogram gets returned as logits initially, we need to exp it before projection
                target_value_histogram_for_all_actions = target_evaluator.get('model:q_dist_next').exp()
                model_value_histogram_for_all_actions = evaluator.get('model:q_dist_next').exp()

                atoms_aligned = self.support_atoms.view(1, 1, self.num_atoms)

                selected_action_indices = (
                    (atoms_aligned * model_value_histogram_for_all_actions).sum(dim=-1).argmax(dim=1)
                )

                # Select largest 'target' value based on action that 'model' selects
                next_value_histograms = (
                    target_value_histogram_for_all_actions[range(batch_size), selected_action_indices]
                )
            else:
                # REGULAR DQN
                # Histogram gets returned as logits initially, we need to exp it before projection
                target_value_histogram_for_all_actions = target_evaluator.get('model:q_dist_next').exp()

                atoms_aligned = self.support_atoms.view(1, 1, self.num_atoms)

                selected_action_indices = (
                    (atoms_aligned * target_value_histogram_for_all_actions).sum(dim=-1).argmax(dim=1)
                )

                next_value_histograms = (
                    target_value_histogram_for_all_actions[range(batch_size), selected_action_indices]
                )

            # HISTOGRAM PROJECTION CODE
            forward_steps = rollout.extra_data.get('forward_steps', 1)

            atoms_projected = (
                rewards_tensor.unsqueeze(1) +
                (self.discount_factor ** forward_steps) *
                (1 - dones_tensor).unsqueeze(1) * self.support_atoms.unsqueeze(0)
            )

            atoms_projected = atoms_projected.clamp(min=self.vmin, max=self.vmax)
            projection_indices = (atoms_projected - self.vmin) / self.atom_delta

            index_floor = projection_indices.floor().long()
            index_ceil = projection_indices.ceil().long()

            # Fix corner case when index_floor == index_ceil
            index_floor[(index_ceil > 0) * (index_floor == index_ceil)] -= 1
            index_ceil[(index_floor < (self.num_atoms - 1)) * (index_floor == index_ceil)] += 1

            value_histogram_projected = torch.zeros_like(next_value_histograms)

            # Following part will be a bit convoluted, in an effort to fully vectorize projection operation

            # Special offset index tensor
            offsets = (
                torch.arange(0, batch_size * self.num_atoms, self.num_atoms)
                .unsqueeze(1)
                .expand(batch_size, self.num_atoms)
                .contiguous().view(-1).to(device)
            )

            # Linearize all the buffers
            value_histogram_projected = value_histogram_projected.view(-1)
            index_ceil = index_ceil.view(-1)
            index_floor = index_floor.view(-1)
            projection_indices = projection_indices.view(-1)

            value_histogram_projected.index_add_(
                0,
                index_floor+offsets,
                (next_value_histograms.view(-1) * (index_ceil.float() - projection_indices))
            )

            value_histogram_projected.index_add_(
                0,
                index_ceil+offsets,
                (next_value_histograms.view(-1) * (projection_indices - index_floor.float()))
            )

            value_histogram_projected = value_histogram_projected.reshape(next_value_histograms.shape)

        q_log_histogram_selected = evaluator.get('model:action:q_dist')

        # Cross-entropy loss as usual
        original_losses = -(value_histogram_projected * q_log_histogram_selected).sum(dim=1)

        if evaluator.is_provided('rollout:weights'):
            weights = evaluator.get('rollout:weights')
        else:
            weights = torch.ones_like(rewards_tensor)

        loss_value = torch.mean(weights * original_losses)
        loss_value.backward()

        with torch.no_grad():
            mean_q_model = (self.support_atoms.unsqueeze(0) * torch.exp(q_log_histogram_selected)).sum(dim=1).mean()
            mean_q_target = (self.support_atoms.unsqueeze(0) * value_histogram_projected).sum(dim=1).mean()

        return {
            'loss': loss_value.item(),
            # We need it to update priorities in the replay buffer:
            'errors': original_losses.detach().cpu().numpy(),
            'average_q_selected': mean_q_model.item(),
            'average_q_target': mean_q_target.item()
        }
1243241
def optimizer_step(self, batch_info, device, model, rollout):
        
        rollout = rollout.to_transitions()

        # This algorithm makes quote strong assumptions about how does the model look
        # so it does not make that much sense to switch to the evaluator interface
        # As it would be more of a problem than actual benefit

        observations = rollout.batch_tensor('observations')
        returns = rollout.batch_tensor('returns')

        # Evaluate model on the observations
        policy_params = model.policy(observations)
        policy_entropy = torch.mean(model.entropy(policy_params))

        policy_loss = self.calc_policy_loss(model, policy_params, policy_entropy, rollout)
        policy_grad = p2v(autograd.grad(policy_loss, model.policy_parameters(), retain_graph=True)).detach()

        # Calculate gradient of KL divergence of model with fixed version of itself
        # Value of kl_divergence will be 0, but what we need is the gradient, actually the 2nd derivarive
        kl_divergence = torch.mean(model.kl_divergence(policy_params.detach(), policy_params))
        kl_divergence_gradient = p2v(torch.autograd.grad(kl_divergence, model.policy_parameters(), create_graph=True))

        step_direction = conjugate_gradient_method(
            matrix_vector_operator=lambda x: self.fisher_vector_product(x, kl_divergence_gradient, model),
            # Because we want to decrease the loss, we want to go into the direction of -gradient
            loss_gradient=-policy_grad,
            nsteps=self.cg_iters
        )

        shs = 0.5 * step_direction @ self.fisher_vector_product(step_direction, kl_divergence_gradient, model)
        lm = torch.sqrt(shs / self.mak_kl)
        full_step = step_direction / lm

        # Because we want to decrease the loss, we want to go into the direction of -gradient
        expected_improvement = (-policy_grad) @ full_step
        original_parameter_vec = p2v(model.policy_parameters()).detach_()

        policy_optimization_success, ratio, policy_loss_improvement, new_policy_loss, kl_divergence_step = self.line_search(
            model, rollout, policy_loss, policy_params, original_parameter_vec, full_step, expected_improvement
        )

        gradient_norms = []

        for i in range(self.vf_iters):
            batch_info.optimizer.zero_grad()
            value_loss = self.value_loss(model, observations, returns)

            value_loss.backward()

            # Gradient clipping
            if self.max_grad_norm is not None:
                grad_norm = torch.nn.utils.clip_grad_norm_(
                    filter(lambda p: p.requires_grad, model.parameters()),
                    max_norm=self.max_grad_norm
                )

                gradient_norms.append(grad_norm)

            batch_info.optimizer.step(closure=None)

        if gradient_norms:
            gradient_norm = np.mean(gradient_norms)
        else:
            gradient_norm = 0.0

        # noinspection PyUnboundLocalVariable
        return {
            'new_policy_loss': new_policy_loss.item(),
            'policy_entropy': policy_entropy.item(),
            'value_loss': value_loss.item(),
            'policy_optimization_success': float(policy_optimization_success),
            'policy_improvement_ratio': ratio.item(),
            'kl_divergence_step': kl_divergence_step.item(),
            'policy_loss_improvement': policy_loss_improvement.item(),
            'grad_norm': gradient_norm,
            'advantage_norm': torch.norm(rollout.batch_tensor('advantages')).item(),
            'explained_variance': explained_variance(returns, rollout.batch_tensor('values'))
        }