<jupyter_start><jupyter_text>Problems:-Need to provide float args to vertex position (x and y)-Need to provide index argument to vertex position-Need to provide index arguments to distance measurement<jupyter_code># DEFINE THE NETWORK

class DQN(nn.Module):
    def __init__(self, outputs):
        super(DQN, self).__init__()

        self.fc1 = nn.Linear(6, 50)  # 6*6 from image dimension
        self.fc2 = nn.Linear(50, 30)
        self.head = nn.Linear(30, outputs)

    # Called with either one element to determine next action, or a batch
    # during optimization
    # Returns a distribution over the possible actions tensor([[left0exp,right0exp]...]).
    def forward(self, x):
        x = x.float()
        #print(x.shape)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.head(x.view(-1, 30)) # self.head(x.view(x.size(0), -1))

"""
batch_size = 5
nb_classes = 2
in_features = 10

model = nn.Linear(in_features, nb_classes)
criterion = nn.CrossEntropyLoss()
"""
# DEFINE TRANSITION OPERATIONS
BATCH_SIZE = 128
GAMMA = 0.5 #0.999
EPS_START = 0.9
EPS_END = 0.05
EPS_DECAY = 50
TARGET_UPDATE = 10

# Get number of actions from gym action space
n_actions = 9

policy_net = DQN(n_actions).to(device)
target_net = DQN(n_actions).to(device)
target_net.load_state_dict(policy_net.state_dict())
target_net.eval()

# convert to double?
policy_net = DQN(n_actions).float()
target_net = DQN(n_actions).float()

optimizer = optim.RMSprop(policy_net.parameters())
memory = ReplayMemory(10000)


steps_done = 0


def select_action(state):
    global steps_done
    sample = random.random()
    eps_threshold = EPS_END + (EPS_START - EPS_END) * \
        math.exp(-1. * steps_done / EPS_DECAY)
    steps_done += 1
    if sample > eps_threshold:
        with torch.no_grad():
            # t.max(1) will return largest column value of each row.
            # second column on max result is index of where max element was
            # found, so we pick action with the larger expected reward.
            """
            best_action = policy_net(state)
            best_action = best_action.max(1)
            best_action = best_action[1]
            best_action.view(1, 1)
            """
            return  policy_net(state).max(1)[1].view(1, 1)
    else:
        random_decision = random.randint(0, 2)
        if random_decision == 0:
            return torch.tensor([[8]], device=device, dtype=torch.long)    
        return torch.tensor([[random.randrange(n_actions)]], device=device, dtype=torch.long)


episode_durations = []


def plot_durations():
    """
    plt.figure(2)
    plt.clf()
    durations_t = torch.tensor(episode_durations, dtype=torch.float)
    plt.title('Training...')
    plt.xlabel('Episode')
    plt.ylabel('Duration')
    plt.plot(durations_t.numpy())
    # Take 100 episode averages and plot them too
    if len(durations_t) >= 100:
        means = durations_t.unfold(0, 100, 1).mean(1).view(-1)
        means = torch.cat((torch.zeros(99), means))
        plt.plot(means.numpy())

    plt.pause(0.001)  # pause a bit so that plots are updated
    if is_ipython:
        display.clear_output(wait=True)
        display.display(plt.gcf())
    """


# TRAINING FUNCTION
def optimize_model():
    if len(memory) < BATCH_SIZE:
        return
    transitions = memory.sample(BATCH_SIZE)
    # Transpose the batch (see https://stackoverflow.com/a/19343/3343043 for
    # detailed explanation). This converts batch-array of Transitions
    # to Transition of batch-arrays.
    batch = Transition(*zip(*transitions))
    #print(batch)

    # Compute a mask of non-final states and concatenate the batch elements
    # (a final state would've been the one after which simulation ended)
    non_final_mask = torch.tensor(tuple(map(lambda s: s is not None,
                                          batch.next_state)), device=device, dtype=torch.bool)
    non_final_next_states = torch.cat([s for s in batch.next_state
                                                if s is not None])
    state_batch = torch.cat(batch.state)
    action_batch = torch.cat(batch.action)
    reward_batch = torch.cat(batch.reward)

    # Compute Q(s_t, a) - the model computes Q(s_t), then we select the
    # columns of actions taken. These are the actions which would've been taken
    # for each batch state according to policy_net
    state_batch = state_batch.view(BATCH_SIZE, 6)
    state_action_values = policy_net(state_batch).gather(1, action_batch)

    # Compute V(s_{t+1}) for all next states.
    # Expected values of actions for non_final_next_states are computed based
    # on the "older" target_net; selecting their best reward with max(1)[0].
    # This is merged based on the mask, such that we'll have either the expected
    # state value or 0 in case the state was final.
    non_final_next_states = non_final_next_states.view(non_final_next_states.shape[0] // 6, 6)

    next_state_values = torch.zeros(BATCH_SIZE, device=device)
    next_state_values[non_final_mask] = target_net(non_final_next_states).max(1)[0].detach()
    # Compute the expected Q values
    expected_state_action_values = (next_state_values * GAMMA) + reward_batch

    # Compute Huber loss
    state_action_values = state_action_values.double()
    expected_state_action_values = expected_state_action_values.double()
    loss = F.smooth_l1_loss(state_action_values, expected_state_action_values.unsqueeze(1))

    # Optimize the model
    optimizer.zero_grad()
    loss = loss.double()
    loss.backward()
    for param in policy_net.parameters():
        param.grad.data.clamp_(-1, 1)
    optimizer.step()
# TRAINING LOOP
num_episodes = 2000
max_count = 999
for i_episode in range(num_episodes):
    # Initialize the environment and state
    print("Episode: ", i_episode)
    env.reset()
    """
    if i_episode < int(num_episodes / 2):
        env.reset(get_complex_reward=False)
    else:
        env.reset(get_complex_reward=True)
    """
    state = env.get_state()
    for t in count():
        #print(t)
        #print(state)
        # Select and perform an action
        action = select_action(state)
        reward, done = env.step(action.item())
        reward = torch.tensor([reward], device=device).double()

        # Observe new state
        if not done:
            next_state = env.get_state()
        else:
            next_state = None

        # Store the transition in memory
        memory.push(state, action, next_state, reward)

        # Move to the next state
        state = next_state

        # Perform one step of the optimization (on the target network)
        optimize_model()
        if done or t > max_count:
            episode_durations.append(t + 1)
            plot_durations()
            print("Took ", t, " steps to complete task")
            break
    # Update the target network, copying all weights and biases in DQN
    if i_episode % TARGET_UPDATE == 0:
        target_net.load_state_dict(policy_net.state_dict())

print('Complete')
env.render()
#env.close()
plt.ioff()
plt.show()
test_states = []
test_actions = []

env.reset()
env.set_state(random.randint(1000, 10000), random.randint(1000, 10000))
state = env.get_state()
for t in count():
    print(t)
    #print(state)
    # Select and perform an action
    action = select_action(state)
    reward, done = env.step(action.item())
    reward = torch.tensor([reward], device=device)

    test_states.append(state)
    test_actions.append(action)

    # Observe new state
    if not done:
        next_state = env.get_state()
    else:
        next_state = None


    # Move to the next state
    state = next_state

    if done or t > max_count:
        episode_durations.append(t + 1)
        plot_durations()
        break
import matplotlib.animation as animation
from IPython.display import HTML

measurements = np.array(test_actions) == 8
#=========================================
# Animate Images

fig = plt.figure(figsize=(10, 10))
ims = []
for i in range(len(test_states)):
    x = np.array(test_states[i])[2:5:2]
    y = np.array(test_states[i])[3:6:2]
    if (measurements[i]):
        ims.append([plt.scatter(x, y, animated=True)])
        im, = plt.plot(x, y, 'ro-', animated=True)
    else:
        im = plt.scatter(x, y, animated=True)
    ims.append([im])

ani = animation.ArtistAnimation(fig, ims, interval=300, blit=True, repeat_delay=1000)
plt.close()

# Show the animation
HTML(ani.to_jshtml())
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 unused import

import matplotlib.pyplot as plt
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import numpy as np

def test_range(x, y):
    X, Y = np.meshgrid(x, y)

    step_vals = []

    def compute_steps(x_val, y_val):
        #print("x: ", x_val, "y :", y_val)
        env.reset()
        env.set_state(x_val, y_val)
        state = env.get_state()
        for t in count():
            #print(state)
            # Select and perform an action
            action = select_action(state)
            reward, done = env.step(action.item())
            reward = torch.tensor([reward], device=device)

            test_states.append(state)
            test_actions.append(action)

            # Observe new state
            if not done:
                next_state = env.get_state()
            else:
                next_state = None


            # Move to the next state
            state = next_state

            if done or t > max_count:
                #print("Took steps: ", t)
                return t
                

    vectorized_steps = np.vectorize(compute_steps)

    Z = vectorized_steps(X, Y)

    plt.pcolormesh(X, Y, Z)
    plt.colorbar()
    plt.show()

    optimal = Z < 3
    plt.pcolormesh(X, Y, optimal)
    plt.show()
x_simple = y_simple = np.linspace(1, 100, 20)
test_range(x_simple, y_simple)
x_simple = y_simple = np.linspace(1, 1000, 20)
test_range(x_simple, y_simple)
x_simple = y_simple = np.linspace(1, 100000, 20)
test_range(x_simple, y_simple)
x_simple = y_simple = np.linspace(1, 1000000000, 20)
test_range(x_simple, y_simple)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11146182679029394, 'fraction_numerical': 0.019826395738804497, 'mean_word_length': 2.8030757689422354, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 10, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 1, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '322426', 'license': 'MIT', 'n_tokens_mistral': 3461, 'n_tokens_neox': 3253, 'n_words': 835, 'path': 'raw_implementations/simple_DQN_pythagoras.ipynb', 'repo_name': 'Whisky-Jack/Euclid', 'size': '10138'}