cleanup

MyRemoteVectorEnv.py

@@ -1,130 +0,0 @@
-from typing import Tuple, Callable, Optional
-from collections import OrderedDict
-
-import gym
-import torch
-import numpy as np
-import ray
-from ray.rllib.env.base_env import BaseEnv, ASYNC_RESET_RETURN
-from ray.rllib.utils.annotations import PublicAPI
-from ray.rllib.utils.typing import MultiEnvDict, EnvType, EnvID, MultiAgentDict
-from stable_baselines3.common.vec_env.base_vec_env import VecEnvObs
-from stable_baselines3.common.vec_env.util import obs_space_info, dict_to_obs
-
-from MyDummyVecEnv import MyDummyVecEnv
-
-
-@PublicAPI
-class MyRemoteVectorEnv(BaseEnv):
-    """Vector env that executes envs in remote workers.
-    This provides dynamic batching of inference as observations are returned
-    from the remote simulator actors. Both single and multi-agent child envs
-    are supported, and envs can be stepped synchronously or async.
-    You shouldn't need to instantiate this class directly. It's automatically
-    inserted when you use the `remote_worker_envs` option for Trainers.
-    """
-
-    @property
-    def observation_space(self):
-        return self._observation_space
-
-    def __init__(self, make_env: Callable[[int], EnvType], num_workers: int, env_per_worker: int, observation_space: Optional[gym.spaces.Space], device: torch.device):
-        self.make_local_env = make_env
-        self.num_workers = num_workers
-        self.env_per_worker = env_per_worker
-        self.num_envs = num_workers * env_per_worker
-        self.poll_timeout = None
-
-        self.actors = None  # lazy init
-        self.pending = None  # lazy init
-
-        self.observation_space = observation_space
-        self.keys, shapes, dtypes = obs_space_info(self.observation_space)
-
-        self.device = device
-
-        self.buf_obs = OrderedDict(
-            [(k, torch.zeros((self.num_envs,) + tuple(shapes[k]), dtype=torch.float, device=self.device)) for k in self.keys])
-        self.buf_dones = np.zeros((self.num_envs,), dtype=bool)
-        self.buf_rews = np.zeros((self.num_envs,), dtype=np.float32)
-        self.buf_infos = [{} for _ in range(self.num_envs)]
-
-    def _save_obs(self, env_idx: int, obs: VecEnvObs) -> None:
-        for key in self.keys:
-            self.buf_obs[key][env_idx * self.env_per_worker: (env_idx + 1) * self.env_per_worker] = torch.from_numpy(obs[key]).to(self.device,
-                                                                                           non_blocking=True)
-
-    def poll(self) -> Tuple[MultiEnvDict, MultiEnvDict, MultiEnvDict,
-                            MultiEnvDict, MultiEnvDict]:
-        if self.actors is None:
-
-            def make_remote_env(i):
-                return _RemoteSingleAgentEnv.remote(self.make_local_env, i, self.env_per_worker)
-
-            self.actors = [make_remote_env(i) for i in range(self.num_workers)]
-
-        if self.pending is None:
-            self.pending = {a.reset.remote(): a for a in self.actors}
-
-        # each keyed by env_id in [0, num_remote_envs)
-        ready = []
-
-        # Wait for at least 1 env to be ready here
-        while not ready:
-            ready, _ = ray.wait(
-                list(self.pending),
-                num_returns=len(self.pending),
-                timeout=self.poll_timeout)
-
-        for obj_ref in ready:
-            actor = self.pending.pop(obj_ref)
-            env_id = self.actors.index(actor)
-            ob, rew, done, info = ray.get(obj_ref)
-
-            self._save_obs(env_id, ob)
-            self.buf_rews[env_id * self.env_per_worker: (env_id + 1) * self.env_per_worker] = rew
-            self.buf_dones[env_id * self.env_per_worker: (env_id + 1) * self.env_per_worker] = done
-            self.buf_infos[env_id * self.env_per_worker: (env_id + 1) * self.env_per_worker] = info
-        return (self._obs_from_buf(), self.buf_rews, self.buf_dones, self.buf_infos)
-
-    def _obs_from_buf(self) -> VecEnvObs:
-        return dict_to_obs(self.observation_space, self.buf_obs)
-
-    @PublicAPI
-    def send_actions(self, action_list) -> None:
-        for worker_id in range(self.num_workers):
-            actions = action_list[worker_id * self.env_per_worker: (worker_id + 1) * self.env_per_worker]
-            actor = self.actors[worker_id]
-            obj_ref = actor.step.remote(actions)
-            self.pending[obj_ref] = actor
-
-    @PublicAPI
-    def try_reset(self,
-                  env_id: Optional[EnvID] = None) -> Optional[MultiAgentDict]:
-        actor = self.actors[env_id]
-        obj_ref = actor.reset.remote()
-        self.pending[obj_ref] = actor
-        return ASYNC_RESET_RETURN
-
-    @PublicAPI
-    def stop(self) -> None:
-        if self.actors is not None:
-            for actor in self.actors:
-                actor.__ray_terminate__.remote()
-
-    @observation_space.setter
-    def observation_space(self, value):
-        self._observation_space = value
-
-@ray.remote(num_cpus=1)
-class _RemoteSingleAgentEnv:
-    """Wrapper class for making a gym env a remote actor."""
-
-    def __init__(self, make_env, i, env_per_worker):
-        self.env = MyDummyVecEnv([lambda: make_env((i * env_per_worker) + k) for k in range(env_per_worker)])
-
-    def reset(self):
-        return self.env.reset(), 0, False, {}
-
-    def step(self, actions):
-        return self.env.step(actions)

MyVecEnv.py

@@ -1,47 +0,0 @@
-from typing import Optional, List, Union, Sequence, Type, Any
-
-import gym
-import numpy as np
-from ray.rllib import BaseEnv
-from stable_baselines3.common.vec_env import VecEnv
-from stable_baselines3.common.vec_env.base_vec_env import VecEnvIndices, VecEnvStepReturn, VecEnvObs
-
-from MyRemoteVectorEnv import MyRemoteVectorEnv
-
-
-class WrapperRay(VecEnv):
-
-    def __init__(self, make_env, num_workers, per_worker_env, device):
-        self.one_env = make_env(0)
-        self.remote: BaseEnv = MyRemoteVectorEnv(make_env, num_workers, per_worker_env, self.one_env.observation_space, device)
-        super(WrapperRay, self).__init__(num_workers * per_worker_env, self.one_env.observation_space, self.one_env.action_space)
-
-    def reset(self) -> VecEnvObs:
-        return self.remote.poll()[0]
-
-    def step_async(self, actions: np.ndarray) -> None:
-        self.remote.send_actions(actions)
-
-    def step_wait(self) -> VecEnvStepReturn:
-        return self.remote.poll()
-
-    def close(self) -> None:
-        self.remote.stop()
-
-    def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> List[Any]:
-        pass
-
-    def set_attr(self, attr_name: str, value: Any, indices: VecEnvIndices = None) -> None:
-        pass
-
-    def env_method(self, method_name: str, *method_args, indices: VecEnvIndices = None, **method_kwargs) -> List[Any]:
-        pass
-
-    def env_is_wrapped(self, wrapper_class: Type[gym.Wrapper], indices: VecEnvIndices = None) -> List[bool]:
-        pass
-
-    def get_images(self) -> Sequence[np.ndarray]:
-        pass
-
-    def seed(self, seed: Optional[int] = None) -> List[Union[None, int]]:
-        pass

Network.py

@@ -1,114 +0,0 @@
-import math
-
-import numpy as np
-import torch
-from torch import nn, Tensor
-from torch.distributions import Categorical
-
-
-class PositionalEncoding(nn.Module):
-
-    def __init__(self, d_model: int, max_len: int = 100):
-        super().__init__()
-        position = torch.arange(max_len).unsqueeze(1)
-        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
-        pe = torch.zeros(max_len, d_model)
-        pe[:, 0::2] = torch.sin(position * div_term)
-        pe[:, 1::2] = torch.cos(position * div_term)
-        self.register_buffer('pe', pe)
-
-    def forward(self, positions: Tensor) -> Tensor:
-        return self.pe[positions]
-
-
-class Actor(nn.Module):
-
-    def __init__(self, pos_encoder):
-        super(Actor, self).__init__()
-        self.activation = nn.Tanh()
-        self.project = nn.Linear(4, 8)
-        nn.init.xavier_uniform_(self.project.weight, gain=1.0)
-        nn.init.constant_(self.project.bias, 0)
-        self.pos_encoder = pos_encoder
-
-        self.embedding_fixed = nn.Embedding(2, 1)
-        self.embedding_legal_op = nn.Embedding(2, 1)
-
-        self.tokens_start_end = nn.Embedding(3, 4)
-
-        # self.conv_transform = nn.Conv1d(5, 1, 1)
-        # nn.init.kaiming_normal_(self.conv_transform.weight, mode="fan_out", nonlinearity="relu")
-        # nn.init.constant_(self.conv_transform.bias, 0)
-
-        self.enc1 = nn.TransformerEncoderLayer(8, 1, dim_feedforward=8 * 4, dropout=0.0, batch_first=True,
-                                               norm_first=True)
-        self.enc2 = nn.TransformerEncoderLayer(8, 1, dim_feedforward=8 * 4, dropout=0.0, batch_first=True,
-                                               norm_first=True)
-
-        self.final_tmp = nn.Sequential(
-            layer_init_tanh(nn.Linear(8, 32)),
-            nn.Tanh(),
-            layer_init_tanh(nn.Linear(32, 1), std=0.01)
-        )
-        self.no_op = nn.Sequential(
-            layer_init_tanh(nn.Linear(8, 32)),
-            nn.Tanh(),
-            layer_init_tanh(nn.Linear(32, 1), std=0.01)
-        )
-
-    def forward(self, obs, attention_interval_mask, job_resource, mask, indexes_inter, tokens_start_end):
-        embedded_obs = torch.cat((self.embedding_fixed(obs[:, :, :, 0].long()), obs[:, :, :, 1:3],
-                                  self.embedding_legal_op(obs[:, :, :, 3].long())), dim=3)
-        non_zero_tokens = tokens_start_end != 0
-        t = tokens_start_end[non_zero_tokens].long()
-        embedded_obs[non_zero_tokens] = self.tokens_start_end(t)
-        pos_encoder = self.pos_encoder(indexes_inter.long())
-        pos_encoder[non_zero_tokens] = 0
-        obs = self.project(embedded_obs) + pos_encoder
-
-        transformed_obs = obs.view(-1, obs.shape[2], obs.shape[3])
-        attention_interval_mask = attention_interval_mask.view(-1, attention_interval_mask.shape[-1])
-        transformed_obs = self.enc1(transformed_obs, src_key_padding_mask=attention_interval_mask == 1)
-        transformed_obs = transformed_obs.view(obs.shape)
-        obs = transformed_obs.mean(dim=2)
-
-        job_resource = job_resource[:, :-1, :-1] == 0
-
-        obs_action = self.enc2(obs, src_mask=job_resource) + obs
-
-        logits = torch.cat((self.final_tmp(obs_action).squeeze(2), self.no_op(obs_action).mean(dim=1)), dim=1)
-        return logits.masked_fill(mask == 0, -3.4028234663852886e+38)
-
-
-class Agent(nn.Module):
-    def __init__(self):
-        super(Agent, self).__init__()
-        self.pos_encoder = PositionalEncoding(8)
-        self.actor = Actor(self.pos_encoder)
-
-    def forward(self, data, attention_interval_mask, job_resource_masks, mask, indexes_inter, tokens_start_end,
-                action=None):
-        logits = self.actor(data, attention_interval_mask, job_resource_masks, mask, indexes_inter, tokens_start_end)
-        probs = Categorical(logits=logits)
-        if action is None:
-            probabilities = probs.probs
-            actions = torch.multinomial(probabilities, probabilities.shape[1])
-            return actions, torch.log(probabilities), probs.entropy()
-        else:
-            return logits, probs.log_prob(action), probs.entropy()
-
-    def get_action_only(self, data, attention_interval_mask, job_resource_masks, mask, indexes_inter, tokens_start_end):
-        logits = self.actor(data, attention_interval_mask, job_resource_masks, mask, indexes_inter, tokens_start_end)
-        probs = Categorical(logits=logits)
-        return probs.sample()
-
-    def get_logits_only(self,data, attention_interval_mask, job_resource_masks, mask, indexes_inter, tokens_start_end):
-        logits = self.actor(data, attention_interval_mask, job_resource_masks, mask, indexes_inter, tokens_start_end)
-        return logits
-
-
-def layer_init_tanh(layer, std=np.sqrt(2), bias_const=0.0):
-    torch.nn.init.orthogonal_(layer.weight, std)
-    if layer.bias is not None:
-        torch.nn.init.constant_(layer.bias, bias_const)
-    return layer

app.py

@@ -3,12 +3,10 @@ import os
 import random
 import time
 
-import gym
 import plotly.figure_factory as ff
 import json
 
 import pandas as pd
-import ray
 
 from compiled_jss.CPEnv import CompiledJssEnvCP
 
@@ -19,7 +17,6 @@ import torch
 import numpy as np
 
 from MyDummyVecEnv import MyDummyVecEnv
-from MyVecEnv import WrapperRay
 
 import gradio as gr