X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=problems.py;h=d7dbc542aa14b0c77999b7b7a55ece9394e57c90;hb=732349f7c16e43ff84380d28e021d671f2c56492;hp=b8fcdb34b32a07e8fd46900c88d889c3cf9761ea;hpb=76e62a5782fc2509ce989fcfc0d0aedc17322b3a;p=picoclvr.git diff --git a/problems.py b/problems.py index b8fcdb3..d7dbc54 100755 --- a/problems.py +++ b/problems.py @@ -24,6 +24,8 @@ class Problem: #################### + + class ProblemDegradation(Problem): def __init__(self, nb_state_tokens=5, nb_time_steps=12, value_max=25, hard=False): assert value_max // nb_state_tokens >= 2 @@ -108,6 +110,48 @@ class ProblemDegradation(Problem): #################### +class ProblemMemory(Problem): + def __init__(self, len_total=25): + self.len_total = len_total + self.max_len_pattern = 5 + self.nb_noise_tokens = 10 + self.start_pattern_token = 0 + self.end_pattern_token = 1 + self.start_result_token = 2 + self.end_result_token = 3 + self.token_string = "[]<>" + "".join( + [chr(ord("a") + k) for k in range(self.nb_noise_tokens)] + ) + + def generate_sequences(self, nb): + sequences = ( + torch.randint(self.nb_noise_tokens, (nb, self.len_total)) + + self.end_result_token + + 1 + ) + len_patterns = torch.randint(self.max_len_pattern, (nb,)) + 1 + pattern_positions = torch.randint( + self.len_total - (5 + 2 * self.max_len_pattern), (nb,) + ) + k = self.len_total - (3 + self.max_len_pattern) + for i in range(nb): + l = len_patterns[i] + j = pattern_positions[i] + sequences[i, j] = self.start_pattern_token + sequences[i, j + l + 2] = self.end_pattern_token + sequences[i, k] = self.start_result_token + sequences[i, k + l + 2] = self.end_result_token + sequences[i, k + 1 : k + 2 + l] = sequences[i, j + 1 : j + 2 + l] + + j = torch.arange(self.len_total)[None, :] + ar_mask = (j > k).long() * (j <= k + 1 + len_patterns[:, None]).long() + + return sequences, ar_mask + + def seq2str(self, seq): + return "".join(self.token_string[x.item()] for x in seq) + + class ProblemTwoTargets(Problem): def __init__(self, len_total=10, len_targets=3): assert len_targets >= 3 @@ -287,44 +331,74 @@ class ProblemAddition(Problem): class ProblemMixing(Problem): - def __init__(self, height=4, width=4, nb_time_steps=9, hard=False): + def __init__( + self, height=4, width=4, nb_time_steps=9, hard=False, random_start=True + ): self.height = height self.width = width self.nb_time_steps = nb_time_steps self.hard = hard + self.random_start = random_start def start_random(self, nb): y = torch.arange(self.height * self.width).reshape(1, -1).expand(nb, -1) - # m = (torch.rand(y.size()).sort(dim=-1).indices < y.size(1) // 2).long() + if self.random_start: + i = ( + torch.arange(self.height) + .reshape(1, -1, 1) + .expand(nb, self.height, self.width) + ) + j = ( + torch.arange(self.width) + .reshape(1, 1, -1) + .expand(nb, self.height, self.width) + ) - i = torch.arange(self.height).reshape(1,-1,1).expand(nb,self.height,self.width) - j = torch.arange(self.width).reshape(1,1,-1).expand(nb,self.height,self.width) + ri = torch.randint(self.height, (nb,)).reshape(nb, 1, 1) + rj = torch.randint(self.width, (nb,)).reshape(nb, 1, 1) - ri = torch.randint(self.height, (nb,)).reshape(nb,1,1) - rj = torch.randint(self.width, (nb,)).reshape(nb,1,1) + m = 1 - torch.logical_or(i == ri, j == rj).long().flatten(1) - m = 1 - torch.logical_or(i==ri,j==rj).long().flatten(1) + y = y * m + self.height * self.width * (1 - m) - y = (y * m + self.height * self.width * (1 - m)).reshape( - nb, self.height, self.width - ) + y = y.reshape(nb, self.height, self.width) return y def start_error(self, x): - i = torch.arange(self.height, device=x.device).reshape(1,-1,1).expand_as(x) - j = torch.arange(self.width, device=x.device).reshape(1,1,-1).expand_as(x) - - ri = (x == self.height * self.width).long().sum(dim=-1).argmax(-1).view(-1,1,1) - rj = (x == self.height * self.width).long().sum(dim=-2).argmax(-1).view(-1,1,1) + if self.random_start: + i = ( + torch.arange(self.height, device=x.device) + .reshape(1, -1, 1) + .expand_as(x) + ) + j = torch.arange(self.width, device=x.device).reshape(1, 1, -1).expand_as(x) + + ri = ( + (x == self.height * self.width) + .long() + .sum(dim=-1) + .argmax(-1) + .view(-1, 1, 1) + ) + rj = ( + (x == self.height * self.width) + .long() + .sum(dim=-2) + .argmax(-1) + .view(-1, 1, 1) + ) - m = 1 - torch.logical_or(i==ri,j==rj).long().flatten(1) + m = 1 - torch.logical_or(i == ri, j == rj).long().flatten(1) + else: + m = 1 x = x.flatten(1) - u = torch.arange(self.height * self.width, device = x.device).reshape(1, -1) + u = torch.arange(self.height * self.width, device=x.device).reshape(1, -1) d = (x - (m * u + (1 - m) * self.height * self.width)).abs().sum(-1) + return d def moves(self, x): @@ -390,7 +464,15 @@ class ProblemMixing(Problem): return " | ".join( [ " ".join( - ["-".join([f"{x:02d}" if x < self.height * self.width else "**" for x in s]) for s in r.split(self.width)] + [ + "-".join( + [ + f"{x:02d}" if x < self.height * self.width else "**" + for x in s + ] + ) + for s in r.split(self.width) + ] ) for r in seq.split(self.height * self.width) ] @@ -400,7 +482,8 @@ class ProblemMixing(Problem): #################### if __name__ == "__main__": - p = ProblemMixing() + p = ProblemMixing(height=3, width=3, random_start=False) + s, m = p.generate_sequences(10000) for x in s[:5]: print(p.seq2str(x))