def __init__(self): super().__init__() base = MLP([nn.ReLU(), nn.ReLU()], [2, 16, 16]) head0 = MLP([nn.ReLU(), nn.Tanh()], [16, 8, 1]) head1 = MLP([nn.ReLU(), nn.Softplus()], [16, 8, 1]) self.dualhead = Dualhead(base, head0, head1)
def set_up_dualhead_from_pixels(encoder: nn.Module, encoder_out_dim: int, middle: Optional[nn.Module] = None, head0: Optional[nn.Module] = None, head1: Optional[nn.Module] = None): # if encoder is None: # encoder = ... if feed_forward is None: feed_forward = MLP([nn.ReLU(), nn.ReLU()], [encoder_out_dim, 128, 64]) if head0 is None: head0 = MLP([nn.ReLU(), nn.ReLU()], [64, 32, 1]) if head1 is None: head1 = MLP([nn.ReLU(), nn.ReLU()], [64, 32, 1]) base = nn.Sequential( encoder, feed_forward, ) return Dualhead(base, head0, head1)
def set_up_repr_dualhead_from_pixels( encoder: nn.Module, encoder_out_dim: int, middle: Optional[nn.Module] = None, head0: Optional[nn.Module] = None, head1: Optional[nn.Module] = None, out_dim: Optional[int] = 1, ): # if encoder is None: # encoder = ... if middle is None: middle = MLP([nn.ReLU()], [encoder_out_dim, 64]) if head0 is None: head0 = MLP([nn.ReLU(), nn.Sigmoid()], [64, 32, out_dim]) if head1 is None: head1 = MLP([nn.ReLU(), nn.Softplus()], [64, 32, out_dim]) # if head0 is None: # head0 = MLP([nn.ReLU(), nn.Tanh(), lambda x: x], [64, 32, out_dim, out_dim]) # if head1 is None: # head1 = MLP([nn.ReLU(), nn.Tanh(), lambda x: x], [64, 32, out_dim, out_dim]) return ReprDualhead(encoder, middle, head0, head1)
def __init__(self): super().__init__() self.mlp = MLP([nn.ReLU(), nn.ReLU()], [2, 64, 64]) self.linear = nn.Linear(64, 3)
def __init__(self): super().__init__() self.actor = MLP([nn.ReLU(), nn.ReLU(), None], [4, 64, 64, 2]) self.critic = MLP([nn.ReLU(), nn.ReLU(), None], [4, 64, 64, 1])