Thomas Knoll
2 years ago
2 changed files with 6 additions and 278 deletions
@ -1,275 +0,0 @@ |
|||||
from typing import Dict, Optional, Union |
|
||||
from ray.rllib.env.base_env import BaseEnv |
|
||||
from ray.rllib.evaluation import RolloutWorker |
|
||||
from ray.rllib.evaluation.episode import Episode |
|
||||
from ray.rllib.evaluation.episode_v2 import EpisodeV2 |
|
||||
from ray.rllib.policy import Policy |
|
||||
from ray.rllib.utils.typing import PolicyID |
|
||||
import stormpy |
|
||||
import stormpy.core |
|
||||
import stormpy.simulator |
|
||||
|
|
||||
from collections import deque |
|
||||
|
|
||||
import stormpy.shields |
|
||||
import stormpy.logic |
|
||||
|
|
||||
import stormpy.examples |
|
||||
import stormpy.examples.files |
|
||||
import os |
|
||||
|
|
||||
import gymnasium as gym |
|
||||
|
|
||||
import minigrid |
|
||||
import numpy as np |
|
||||
|
|
||||
import ray |
|
||||
from ray.tune import register_env |
|
||||
from ray.rllib.algorithms.ppo import PPOConfig |
|
||||
from ray.rllib.utils.test_utils import check_learning_achieved, framework_iterator |
|
||||
from ray import tune, air |
|
||||
from ray.rllib.algorithms.callbacks import DefaultCallbacks |
|
||||
from ray.tune.logger import pretty_print |
|
||||
from ray.rllib.utils.numpy import one_hot |
|
||||
from ray.rllib.algorithms import ppo |
|
||||
|
|
||||
from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFC |
|
||||
from ray.rllib.models.tf.fcnet import FullyConnectedNetwork |
|
||||
from ray.rllib.utils.torch_utils import FLOAT_MIN |
|
||||
|
|
||||
from ray.rllib.models.preprocessors import get_preprocessor |
|
||||
|
|
||||
import matplotlib.pyplot as plt |
|
||||
|
|
||||
import argparse |
|
||||
from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 |
|
||||
from ray.rllib.utils.framework import try_import_tf, try_import_torch |
|
||||
|
|
||||
from examples.shields.rl.Wrappers import OneHotShieldingWrapper |
|
||||
|
|
||||
|
|
||||
torch, nn = try_import_torch() |
|
||||
|
|
||||
class TorchActionMaskModel(TorchModelV2, nn.Module): |
|
||||
"""PyTorch version of above ActionMaskingModel.""" |
|
||||
|
|
||||
def __init__( |
|
||||
self, |
|
||||
obs_space, |
|
||||
action_space, |
|
||||
num_outputs, |
|
||||
model_config, |
|
||||
name, |
|
||||
**kwargs, |
|
||||
): |
|
||||
orig_space = getattr(obs_space, "original_space", obs_space) |
|
||||
assert ( |
|
||||
isinstance(orig_space, Dict) |
|
||||
and "action_mask" in orig_space.spaces |
|
||||
and "observations" in orig_space.spaces |
|
||||
) |
|
||||
|
|
||||
TorchModelV2.__init__( |
|
||||
self, obs_space, action_space, num_outputs, model_config, name, **kwargs |
|
||||
) |
|
||||
nn.Module.__init__(self) |
|
||||
|
|
||||
self.internal_model = TorchFC( |
|
||||
orig_space["observations"], |
|
||||
action_space, |
|
||||
num_outputs, |
|
||||
model_config, |
|
||||
name + "_internal", |
|
||||
) |
|
||||
|
|
||||
# disable action masking --> will likely lead to invalid actions |
|
||||
self.no_masking = False |
|
||||
if "no_masking" in model_config["custom_model_config"]: |
|
||||
self.no_masking = model_config["custom_model_config"]["no_masking"] |
|
||||
|
|
||||
def forward(self, input_dict, state, seq_lens): |
|
||||
# Extract the available actions tensor from the observation. |
|
||||
action_mask = input_dict["obs"]["action_mask"] |
|
||||
|
|
||||
# Compute the unmasked logits. |
|
||||
logits, _ = self.internal_model({"obs": input_dict["obs"]["observations"]}) |
|
||||
|
|
||||
# If action masking is disabled, directly return unmasked logits |
|
||||
if self.no_masking: |
|
||||
return logits, state |
|
||||
|
|
||||
# Convert action_mask into a [0.0 || -inf]-type mask. |
|
||||
inf_mask = torch.clamp(torch.log(action_mask), min=FLOAT_MIN) |
|
||||
masked_logits = logits + inf_mask |
|
||||
|
|
||||
# Return masked logits. |
|
||||
return masked_logits, state |
|
||||
|
|
||||
def value_function(self): |
|
||||
return self.internal_model.value_function() |
|
||||
|
|
||||
|
|
||||
|
|
||||
class MyCallbacks(DefaultCallbacks): |
|
||||
def on_episode_start(self, *, worker: RolloutWorker, base_env: BaseEnv, policies: Dict[PolicyID, Policy], episode: Episode | EpisodeV2, env_index: int | None = None, **kwargs) -> None: |
|
||||
# print(F"Epsiode started Environment: {base_env.get_sub_environments()}") |
|
||||
env = base_env.get_sub_environments()[0] |
|
||||
episode.user_data["count"] = 0 |
|
||||
# print(env.printGrid()) |
|
||||
# print(env.action_space.n) |
|
||||
# print(env.actions) |
|
||||
# print(env.mission) |
|
||||
# print(env.observation_space) |
|
||||
# img = env.get_frame() |
|
||||
# plt.imshow(img) |
|
||||
# plt.show() |
|
||||
|
|
||||
def on_episode_step(self, *, worker: RolloutWorker, base_env: BaseEnv, policies: Dict[PolicyID, Policy] | None = None, episode: Episode | EpisodeV2, env_index: int | None = None, **kwargs) -> None: |
|
||||
episode.user_data["count"] = episode.user_data["count"] + 1 |
|
||||
env = base_env.get_sub_environments()[0] |
|
||||
# print(env.env.env.printGrid()) |
|
||||
|
|
||||
def on_episode_end(self, *, worker: RolloutWorker, base_env: BaseEnv, policies: Dict[PolicyID, Policy], episode: Episode | EpisodeV2 | Exception, env_index: int | None = None, **kwargs) -> None: |
|
||||
# print(F"Epsiode end Environment: {base_env.get_sub_environments()}") |
|
||||
env = base_env.get_sub_environments()[0] |
|
||||
# print(env.env.env.printGrid()) |
|
||||
# print(episode.user_data["count"]) |
|
||||
|
|
||||
|
|
||||
|
|
||||
|
|
||||
|
|
||||
def parse_arguments(argparse): |
|
||||
parser = argparse.ArgumentParser() |
|
||||
# parser.add_argument("--env", help="gym environment to load", default="MiniGrid-Empty-8x8-v0") |
|
||||
parser.add_argument("--env", help="gym environment to load", default="MiniGrid-LavaCrossingS9N1-v0") |
|
||||
parser.add_argument("--seed", type=int, help="seed for environment", default=1) |
|
||||
parser.add_argument("--tile_size", type=int, help="size at which to render tiles", default=32) |
|
||||
parser.add_argument("--agent_view", default=False, action="store_true", help="draw the agent sees") |
|
||||
parser.add_argument("--grid_path", default="Grid.txt") |
|
||||
parser.add_argument("--prism_path", default="Grid.PRISM") |
|
||||
|
|
||||
args = parser.parse_args() |
|
||||
|
|
||||
return args |
|
||||
|
|
||||
|
|
||||
def env_creater(config): |
|
||||
name = config.get("name", "MiniGrid-LavaCrossingS9N1-v0") |
|
||||
# name = config.get("name", "MiniGrid-Empty-8x8-v0") |
|
||||
framestack = config.get("framestack", 4) |
|
||||
|
|
||||
env = gym.make(name) |
|
||||
# env = minigrid.wrappers.RGBImgPartialObsWrapper(env) |
|
||||
env = minigrid.wrappers.ImgObsWrapper(env) |
|
||||
env = OneHotShieldingWrapper(env, |
|
||||
config.vector_index if hasattr(config, "vector_index") else 0, |
|
||||
framestack=framestack |
|
||||
) |
|
||||
|
|
||||
|
|
||||
return env |
|
||||
|
|
||||
|
|
||||
def create_shield(grid_file, prism_path): |
|
||||
os.system(F"/home/tknoll/Documents/main -v 'agent' -i {grid_file} -o {prism_path}") |
|
||||
|
|
||||
f = open(prism_path, "a") |
|
||||
f.write("label \"AgentIsInLava\" = AgentIsInLava;") |
|
||||
f.close() |
|
||||
|
|
||||
|
|
||||
program = stormpy.parse_prism_program(prism_path) |
|
||||
formula_str = "Pmax=? [G !\"AgentIsInLavaAndNotDone\"]" |
|
||||
|
|
||||
formulas = stormpy.parse_properties_for_prism_program(formula_str, program) |
|
||||
options = stormpy.BuilderOptions([p.raw_formula for p in formulas]) |
|
||||
options.set_build_state_valuations(True) |
|
||||
options.set_build_choice_labels(True) |
|
||||
options.set_build_all_labels() |
|
||||
model = stormpy.build_sparse_model_with_options(program, options) |
|
||||
|
|
||||
shield_specification = stormpy.logic.ShieldExpression(stormpy.logic.ShieldingType.PRE_SAFETY, stormpy.logic.ShieldComparison.RELATIVE, 0.1) |
|
||||
result = stormpy.model_checking(model, formulas[0], extract_scheduler=True, shield_expression=shield_specification) |
|
||||
|
|
||||
assert result.has_scheduler |
|
||||
assert result.has_shield |
|
||||
shield = result.shield |
|
||||
|
|
||||
stormpy.shields.export_shield(model, shield,"Grid.shield") |
|
||||
|
|
||||
return shield.construct(), model |
|
||||
|
|
||||
def export_grid_to_text(env, grid_file): |
|
||||
f = open(grid_file, "w") |
|
||||
print(env) |
|
||||
f.write(env.printGrid(init=True)) |
|
||||
# f.write(env.pprint_grid()) |
|
||||
f.close() |
|
||||
|
|
||||
def create_environment(args): |
|
||||
env_id= args.env |
|
||||
env = gym.make(env_id) |
|
||||
env.reset() |
|
||||
return env |
|
||||
|
|
||||
|
|
||||
def main(): |
|
||||
args = parse_arguments(argparse) |
|
||||
|
|
||||
env = create_environment(args) |
|
||||
ray.init(num_cpus=3) |
|
||||
|
|
||||
# print(env.pprint_grid()) |
|
||||
# print(env.printGrid(init=False)) |
|
||||
|
|
||||
grid_file = args.grid_path |
|
||||
export_grid_to_text(env, grid_file) |
|
||||
|
|
||||
prism_path = args.prism_path |
|
||||
shield, model = create_shield(grid_file, prism_path) |
|
||||
|
|
||||
for state_id in model.states: |
|
||||
choices = shield.get_choice(state_id) |
|
||||
print(F"Allowed choices in state {state_id}, are {choices.choice_map} ") |
|
||||
|
|
||||
env_name = "mini-grid" |
|
||||
register_env(env_name, env_creater) |
|
||||
|
|
||||
|
|
||||
algo =( |
|
||||
PPOConfig() |
|
||||
.rollouts(num_rollout_workers=1) |
|
||||
.resources(num_gpus=0) |
|
||||
.environment(env="mini-grid") |
|
||||
.framework("torch") |
|
||||
.callbacks(MyCallbacks) |
|
||||
.training(model={ |
|
||||
"fcnet_hiddens": [256,256], |
|
||||
"fcnet_activation": "relu", |
|
||||
|
|
||||
}) |
|
||||
.build() |
|
||||
) |
|
||||
episode_reward = 0 |
|
||||
terminated = truncated = False |
|
||||
obs, info = env.reset() |
|
||||
|
|
||||
# while not terminated and not truncated: |
|
||||
# action = algo.compute_single_action(obs) |
|
||||
# obs, reward, terminated, truncated = env.step(action) |
|
||||
|
|
||||
for i in range(30): |
|
||||
result = algo.train() |
|
||||
print(pretty_print(result)) |
|
||||
|
|
||||
if i % 5 == 0: |
|
||||
checkpoint_dir = algo.save() |
|
||||
print(f"Checkpoint saved in directory {checkpoint_dir}") |
|
||||
|
|
||||
|
|
||||
|
|
||||
ray.shutdown() |
|
||||
|
|
||||
if __name__ == '__main__': |
|
||||
main() |
|
||||
Write
Preview
Loading…
Cancel
Save
Reference in new issue