moved renderer.py to base, added initial salina experiments

studies/sat_mad.py

@@ -1,29 +1,100 @@
 from environments.factory import make
-import salina
+from salina import Workspace, TAgent
+from salina.agents.gyma import AutoResetGymAgent, GymAgent
+from salina.agents import Agents, TemporalAgent
+from salina.rl.functional import _index
 import torch
-from gym.wrappers import FrameStack
+import torch.nn as nn
+from torch.nn.utils import spectral_norm
+import torch.optim as optim
+from torch.distributions import Categorical
 
 
-class MyAgent(salina.TAgent):
-    def __init__(self):
-        super(MyAgent, self).__init__()
+class A2CAgent(TAgent):
+    def __init__(self, observation_size, hidden_size, n_actions):
+        super().__init__()
+        self.model = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(observation_size, hidden_size),
+            nn.ELU(),
+            nn.Linear(hidden_size, hidden_size),
+            nn.ELU(),
+            nn.Linear(hidden_size, n_actions),
+        )
+        self.critic_model = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(observation_size, hidden_size),
+            nn.ELU(),
+            spectral_norm(nn.Linear(hidden_size, 1)),
+        )
 
-    def forward(self, t, **kwargs):
-        self.set(('timer', t), torch.tensor([t]))
+    def forward(self, t, stochastic, **kwargs):
+        observation = self.get(("env/env_obs", t))
+        scores = self.model(observation)
+        probs = torch.softmax(scores, dim=-1)
+        critic = self.critic_model(observation).squeeze(-1)
+        if stochastic:
+            action = torch.distributions.Categorical(probs).sample()
+        else:
+            action = probs.argmax(1)
+
+        self.set(("action", t), action)
+        self.set(("action_probs", t), probs)
+        self.set(("critic", t), critic)
 
 
 if __name__ == '__main__':
-    n_agents = 1
-    env = make('DirtyFactory-v0', n_agents=n_agents)
-    env = FrameStack(env, num_stack=3)
-    env.reset()
-    agent = MyAgent()
-    workspace = salina.Workspace()
-    agent(workspace, t=0, n_steps=10)
+    # Setup agents and workspace
+    env_agent = AutoResetGymAgent(make, dict(env_str='DirtyFactory-v0'), n_envs=1)
+    a2c_agent = A2CAgent(3*4*5*5, 96, 10)
+    workspace = Workspace()
 
-    print(workspace)
+    eval_agent = Agents(GymAgent(make, dict(env_str='DirtyFactory-v0'), n_envs=1), a2c_agent)
+    for i in range(100):
+        eval_agent(workspace, t=i, save_render=True, stochastic=True)
 
+    assert False
+    # combine agents
+    acquisition_agent = TemporalAgent(Agents(env_agent, a2c_agent))
+    acquisition_agent.seed(0)
 
-    for i in range(1000):
-        state, *_ = env.step([env.unwrapped.action_space.sample() for _ in range(n_agents)])
-        #env.render()
+    # optimizers & other parameters
+    optimizer = optim.Adam(a2c_agent.parameters(), lr=1e-3)
+    n_timesteps = 10
+
+    # Decision making loop
+    for epoch in range(200000):
+        workspace.zero_grad()
+        if epoch > 0:
+            workspace.copy_n_last_steps(1)
+            acquisition_agent(workspace, t=1, n_steps=n_timesteps-1, stochastic=True)
+        else:
+            acquisition_agent(workspace, t=0, n_steps=n_timesteps, stochastic=True)
+        #for k in workspace.keys():
+        #    print(f'{k} ==> {workspace[k].size()}')
+        critic, done, action_probs, reward, action = workspace[
+            "critic", "env/done", "action_probs", "env/reward", "action"
+        ]
+
+        target = reward[1:] + 0.99 * critic[1:].detach() * (1 - done[1:].float())
+        td = target - critic[:-1]
+        td_error = td ** 2
+        critic_loss = td_error.mean()
+        entropy_loss = Categorical(action_probs).entropy().mean()
+        action_logp = _index(action_probs, action).log()
+        a2c_loss = action_logp[:-1] * td.detach()
+        a2c_loss = a2c_loss.mean()
+        loss = (
+            -0.001 * entropy_loss
+            + 1.0 * critic_loss
+            - 0.1 * a2c_loss
+        )
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # Compute the cumulated reward on final_state
+        creward = workspace["env/cumulated_reward"]
+        creward = creward[done]
+        if creward.size()[0] > 0:
+            print(f"Cumulative reward at A2C step #{(1+epoch)*n_timesteps}: {creward.mean().item()}")