<jupyter_start><jupyter_text>Here, we implement one-step actor-critic with neural networks and test it on a simple 2D reaching environment. We will start with the environment first. The agent is confined in a 2D arena defined by state $S \in [-1, +1]^2$. An example of states is then $[0.5, -0.2]^\top$. Agent has eight actions available that change the state by one of following eight possible increments $f(A)$: $[-b, -b], [-b, 0], [-b, 1], [0, -b], [0, 0], [0, b], [b, -b], [b, 0], \text{ and } [b, b]$, where $b>0$ is a small scalar. After choosing an action $A_t$, the state is changed as follows: $S_{t+1} = S_t + f(A_t)$, except when the state goes outside the arena, in which case it remains within the arena. It is an episodic task. Each episode, the agent is spawned from a uniformly randomly drawn state from the arena. The episode terminates when the agent reaches a small square around $(0, 0)$. The goal of the agent is to reach the terminal state as soon as possible. We formulate this goal by giving $-1$ rewards every time step.<jupyter_code># Actor critic agent
%reset -f
import torch as tor
import matplotlib.pyplot as plt


# Problem
tor.manual_seed(1)
nacts = 8
aval = 0.03*tor.tensor([
    [-1, -1], [-1, 0], [-1, 1], [0, -1], [0, 0], [0, 1], [1, -1], [1, 0], [1, 1]
], dtype=tor.float32)
LB = tor.tensor([-1., -1.]) 
UB = tor.tensor([1., 1.])

# Agent
nhid = 10
alpha = 0.0003
actor_body = tor.nn.Sequential(tor.nn.Linear(2, nhid), tor.nn.ReLU(),
                               tor.nn.Linear(nhid, nhid), tor.nn.ReLU()
                               )
actor_pref = tor.nn.Sequential(tor.nn.Linear(nhid, nacts))
actor_pref[-1].weight.data[:] = 0; actor_pref[-1].bias.data[:] = 0
critic = tor.nn.Sequential(tor.nn.Linear(2, nhid), tor.nn.ReLU(),
                           tor.nn.Linear(nhid, nhid), tor.nn.ReLU(),
                          tor.nn.Linear(nhid, 1))
popt = tor.optim.Adam(list(actor_body.parameters())+list(actor_pref.parameters()),lr=alpha)
copt = tor.optim.Adam(critic.parameters(), lr=10*alpha)

# Experiment
EP = 600
rets = []
Slogs = []
i = 0
for ep in range(EP):
    Slogs.append([])
    S = tor.rand((1, 2))*(UB-LB) + LB
    Slogs[-1].append(S)
    ret = 0
    while True:
        # Take action
        feat = actor_body(S)
        pref = actor_pref(feat)
        try:
            pol = tor.distributions.Categorical(logits=pref)
        except:
            print("E", pref)
        Ai = pol.sample()

        # Receive reward and next state
        noise = tor.rand(2)*0.06-0.03
        SP = tor.clamp(S + aval[Ai] + noise, LB, UB)
        R = -0.01

        # Learning
        done = tor.allclose(SP, tor.zeros(2), atol=0.2)
        vs = critic(S); vsp = critic(SP)
        pobj = pol.log_prob(Ai)*(R + (1-done)*vsp - vs).detach()
        ploss = -pobj
        closs = (R + (1-done)*vsp.detach() - vs)**2
        popt.zero_grad()
        ploss.backward()
        popt.step()
        copt.zero_grad()
        closs.backward()
        copt.step()

        # Log
        Slogs[-1].append(SP)
        ret += R

        # Termination
        if done:
            print(vsp)
            rets.append(ret)
            i += 1
            print(i, len(Slogs[-1]))
            break
        S = SP

# Plotting
plt.plot(-100*tor.tensor(rets))
plt.figure()

colors = ["tab:blue", "tab:green", "tab:orange", "tab:purple", "tab:red", "tab:brown"]
for i in range(-min(30, EP), 0):
    color = colors[i%len(colors)]
    Slog = tor.cat(Slogs[i])
    for i in range(Slog.shape[0]-1):
        plt.plot(Slog[i:i+2,0], Slog[i:i+2,1], alpha=(i+1)/Slog.shape[0], color=color)
plt.xlim([LB[0], UB[0]])
plt.ylim([LB[1], UB[1]])
plt.gca().set_aspect('equal', adjustable='box')
plt.grid()
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16817939135077417, 'fraction_numerical': 0.029631607047517352, 'mean_word_length': 2.5282485875706215, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '503224', 'license': 'Apache-2.0', 'n_tokens_mistral': 1427, 'n_tokens_neox': 1347, 'n_words': 383, 'path': 'discrete_2d.ipynb', 'repo_name': 'jchzener/Algorithms-Notebooks', 'size': '3746'}