text: import gym
  import itertools as it
  import torch
  
  from vel.api.base import LinearBackboneModel, Model, ModelFactory
  from vel.rl.api import Rollout, Evaluator
  from vel.rl.modules.deterministic_action_head import DeterministicActionHead
  from vel.rl.modules.deterministic_critic_head import DeterministicCriticHead
  
  
  class DeterministicPolicyEvaluator(Evaluator):
      """ Evaluator for DeterministicPolicyModel """
  
      def __init__(self, model: 'DeterministicPolicyModel', rollout: Rollout):
          super().__init__(rollout)
  
          self.model = model
  
      @Evaluator.provides('model:estimated_values_next')
      def model_estimated_values_next(self):
          """ Estimate state-value of the transition next state """
          observations = self.get('rollout:observations_next')
          action, value = self.model(observations)
          return value
  
      @Evaluator.provides('model:actions')
      def model_actions(self):
          """ Estimate state-value of the transition next state """
          observations = self.get('rollout:observations')
          model_action = self.model.action(observations)
          return model_action
  
      @Evaluator.provides('model:model_action:q')
      def model_model_action_q(self):
          observations = self.get('rollout:observations')
          model_actions = self.get('model:actions')
          return self.model.value(observations, model_actions)
  
      @Evaluator.provides('model:action:q')
      def model_action_q(self):
          observations = self.get('rollout:observations')
          rollout_actions = self.get('rollout:actions')
          return self.model.value(observations, rollout_actions)
  
  
  class DeterministicPolicyModel(Model):
      """ Deterministic Policy Gradient - model """
  
      def __init__(self, policy_backbone: LinearBackboneModel, value_backbone: LinearBackboneModel,
                   action_space: gym.Space):
          super().__init__()
  
          self.policy_backbone = policy_backbone
          self.value_backbone = value_backbone
  
          self.action_head = DeterministicActionHead(self.policy_backbone.output_dim, action_space)
          self.critic_head = DeterministicCriticHead(self.value_backbone.output_dim)
  
      def reset_weights(self):
          """ Initialize properly model weights """
          self.policy_backbone.reset_weights()
          self.value_backbone.reset_weights()
          self.action_head.reset_weights()
          self.critic_head.reset_weights()
  
      def forward(self, observations, input_actions=None):
          """ Calculate model outputs """
          observations = observations.float()
  
          if input_actions is not None:
              actions = input_actions
  
              value_input = torch.cat([observations, actions], dim=1)
              value_hidden = self.value_backbone(value_input)
  
              values = self.critic_head(value_hidden)
          else:
              policy_hidden = self.policy_backbone(observations)
              actions = self.action_head(policy_hidden)
  
              value_input = torch.cat([observations, actions], dim=1)
              value_hidden = self.value_backbone(value_input)
  
              values = self.critic_head(value_hidden)
  
          return actions, values
  
      def policy_parameters(self):
          """ Parameters of policy """
          return it.chain(self.policy_backbone.parameters(), self.action_head.parameters())
  
      def value_parameters(self):
          """ Parameters of policy """
          return it.chain(self.value_backbone.parameters(), self.critic_head.parameters())
  
      def get_layer_groups(self):
          """ Return layers grouped """
          return [
              [self.policy_backbone, self.action_head],
              [self.value_backbone, [y for (x, y) in self.critic_head.named_parameters() if x.endswith('bias')]],
              # OpenAI regularizes only weight on the last layer. I'm just replicating that
              [[y for (x, y) in self.critic_head.named_parameters() if x.endswith('weight')]]
          ]
  
      def step(self, observation):
          """ Select actions based on model's output """
          action, value = self(observation)
  
          return {
              'actions': action,
              'values': value
          }
  
      def value(self, observation, input_actions=None):
          """ Calculate value for given state """
          action, value = self(observation, input_actions)
          return value
  
      def action(self, observations):
          """ Calculate value for given state """
          observations = observations.float()
          policy_hidden = self.policy_backbone(observations)
          action = self.action_head(policy_hidden)
          return action
  
      def evaluate(self, rollout: Rollout) -> Evaluator:
          """ Evaluate model on a rollout """
          return DeterministicPolicyEvaluator(self, rollout)
  
  
  class DeterministicPolicyModelFactory(ModelFactory):
      """ Factory  class for policy gradient models """
      def __init__(self, policy_backbone: ModelFactory, value_backbone: ModelFactory):
          self.policy_backbone = policy_backbone
          self.value_backbone = value_backbone
  
      def instantiate(self, **extra_args):
          """ Instantiate the model """
          policy_backbone = self.policy_backbone.instantiate(**extra_args)
          value_backbone = self.value_backbone.instantiate(**extra_args)
  
          return DeterministicPolicyModel(
              policy_backbone=policy_backbone,
              value_backbone=value_backbone,
              action_space=extra_args['action_space'],
          )
  
  
  def create(policy_backbone: ModelFactory, value_backbone: ModelFactory):
      """ Vel creation function """
      return DeterministicPolicyModelFactory(
          policy_backbone=policy_backbone, value_backbone=value_backbone
      )
  