"""wrapper for actor critic that takes model name 
as input and creates a new network based on the input architecture 
with actor as normal forward and critic with a special function
"""
# TODO add auxiliary tasks to the model code
from .model import Model
from gym.spaces import Box
import torch.nn as nn
from .model import Model
from typing import Type


class ActorCritic(Model):
    def __init__(self, model_type: Type[Model], image_space: Box, n_classes: int) -> None:
        actor_model = model_type(image_space, n_classes)
        critic_model = model_type(image_space, 1)
        encoder = actor_model.get_encoder()
        actor_decoder = actor_model.get_decoder()
        super().__init__(encoder + actor_decoder, len(encoder), n_classes=n_classes)
        self.critic_decoder = nn.Sequential(*critic_model.get_decoder())

    def get_action_critic(self, x):
        """takes observation as input and returns actor/ critic

        Args:
            x (tensor): input tensor

        Returns:
            tuple: actor output, critic output
        """        
        x = x.to(self.dummy_param.device)
        encoder_output, actor_output = self.get_penultimate(x)
        critic_output = self.critic_decoder(encoder_output)
        return actor_output, critic_output


98a993c485a4bbb3caf50e1db4779e3866eb009f
<mask token>


class ActorCritic(Model):
    <mask token>
    <mask token>

<mask token>


class ActorCritic(Model):

    def __init__(self, model_type: Type[Model], image_space: Box, n_classes:
        int) ->None:
        actor_model = model_type(image_space, n_classes)
        critic_model = model_type(image_space, 1)
        encoder = actor_model.get_encoder()
        actor_decoder = actor_model.get_decoder()
        super().__init__(encoder + actor_decoder, len(encoder), n_classes=
            n_classes)
        self.critic_decoder = nn.Sequential(*critic_model.get_decoder())
    <mask token>

<mask token>


class ActorCritic(Model):

    def __init__(self, model_type: Type[Model], image_space: Box, n_classes:
        int) ->None:
        actor_model = model_type(image_space, n_classes)
        critic_model = model_type(image_space, 1)
        encoder = actor_model.get_encoder()
        actor_decoder = actor_model.get_decoder()
        super().__init__(encoder + actor_decoder, len(encoder), n_classes=
            n_classes)
        self.critic_decoder = nn.Sequential(*critic_model.get_decoder())

    def get_action_critic(self, x):
        """takes observation as input and returns actor/ critic

        Args:
            x (tensor): input tensor

        Returns:
            tuple: actor output, critic output
        """
        x = x.to(self.dummy_param.device)
        encoder_output, actor_output = self.get_penultimate(x)
        critic_output = self.critic_decoder(encoder_output)
        return actor_output, critic_output

<mask token>
from .model import Model
from gym.spaces import Box
import torch.nn as nn
from .model import Model
from typing import Type


class ActorCritic(Model):

    def __init__(self, model_type: Type[Model], image_space: Box, n_classes:
        int) ->None:
        actor_model = model_type(image_space, n_classes)
        critic_model = model_type(image_space, 1)
        encoder = actor_model.get_encoder()
        actor_decoder = actor_model.get_decoder()
        super().__init__(encoder + actor_decoder, len(encoder), n_classes=
            n_classes)
        self.critic_decoder = nn.Sequential(*critic_model.get_decoder())

    def get_action_critic(self, x):
        """takes observation as input and returns actor/ critic

        Args:
            x (tensor): input tensor

        Returns:
            tuple: actor output, critic output
        """
        x = x.to(self.dummy_param.device)
        encoder_output, actor_output = self.get_penultimate(x)
        critic_output = self.critic_decoder(encoder_output)
        return actor_output, critic_output

"""wrapper for actor critic that takes model name 
as input and creates a new network based on the input architecture 
with actor as normal forward and critic with a special function
"""
# TODO add auxiliary tasks to the model code
from .model import Model
from gym.spaces import Box
import torch.nn as nn
from .model import Model
from typing import Type


class ActorCritic(Model):
    def __init__(self, model_type: Type[Model], image_space: Box, n_classes: int) -> None:
        actor_model = model_type(image_space, n_classes)
        critic_model = model_type(image_space, 1)
        encoder = actor_model.get_encoder()
        actor_decoder = actor_model.get_decoder()
        super().__init__(encoder + actor_decoder, len(encoder), n_classes=n_classes)
        self.critic_decoder = nn.Sequential(*critic_model.get_decoder())

    def get_action_critic(self, x):
        """takes observation as input and returns actor/ critic

        Args:
            x (tensor): input tensor

        Returns:
            tuple: actor output, critic output
        """        
        x = x.to(self.dummy_param.device)
        encoder_output, actor_output = self.get_penultimate(x)
        critic_output = self.critic_decoder(encoder_output)
        return actor_output, critic_output

