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])