Some initial changes to introduce Pytest into the code

.github/workflows/test.yml

@@ -24,4 +24,4 @@ jobs:
           pip install torch
           pip install .
       - name: Execute tests
-        run: python -m unittest tests/*.py tests/**/*.py
+        run: python -m pytest tests

setup.py

@@ -13,6 +13,7 @@
             'gym==0.21',
             'griddly',
             'networkx',
+            'pytest',
             'box2d',
             'tensorboard',
         ]

tests/action_sampler_test.py

@@ -0,0 +1,60 @@
+import logging
+
+import pytest
+import numpy as np
+import torch
+from torch import rand
+from gym.spaces import Discrete, MultiDiscrete, Box
+from watts.models.categorical_action_sampler import ActionSampler
+
+
+
+def test_multi_discrete_sampling():
+    action_space = MultiDiscrete([2, 5])
+    sampler = ActionSampler(action_space=action_space)
+    # obs_space = Box(0.0, 255.0, (5, 5, 6), np.float64)
+    # state = torch.FloatTensor([obs_space.sample()])
+    # get NN and run the state through the NN to get real logits
+    random_logits = rand(1, sum(action_space.nvec))
+    actions, logp, entropy = sampler.sample(random_logits)
+    logging.info(actions)
+    logging.info(logp)
+    logging.info(entropy)
+
+def test_discrete_sampling():
+    action_space = Discrete(5)
+    sampler = ActionSampler(action_space=action_space)
+    # obs_space = Box(0.0, 255.0, (5, 5, 6), np.float64)
+    # state = torch.FloatTensor([obs_space.sample()])
+    # get NN and run the state through the NN to get real logits
+    random_logits = rand(32, action_space.n)
+    actions, logp, entropy = sampler.sample(random_logits)
+    logging.info(actions)
+    logging.info(logp)
+    logging.info(entropy)
+
+def test_discrete_max_sampling():
+    action_space = Discrete(5)
+    sampler = ActionSampler(action_space=action_space)
+    # obs_space = Box(0.0, 255.0, (5, 5, 6), np.float64)
+    # state = torch.FloatTensor([obs_space.sample()])
+    # get NN and run the state through the NN to get real logits
+    logits = torch.FloatTensor([[0.01, 0.3, 0.29, 0.2, 0.2]])
+    # rand(1, action_space.n)
+    actions, logp, entropy = sampler.sample(logits, max=True)
+    assert(actions == 1)
+
+def test_multi_discrete_max_sampling():
+    action_space = MultiDiscrete([2, 5])
+    sampler = ActionSampler(action_space=action_space)
+    # obs_space = Box(0.0, 255.0, (5, 5, 6), np.float64)
+    # state = torch.FloatTensor([obs_space.sample()])
+    # get NN and run the state through the NN to get real logits
+    random_logits = torch.FloatTensor([[0.6, 0.4, 0.2, 0.4, 0.2, 0.1, 0.1]])
+    actions, logp, entropy = sampler.sample(random_logits, max=True)
+    assert(actions[0][0] == 0, actions[0][1] == 1)
+
+def test_invalid_sampling():
+    action_space = Box(-2.0, 2.0, (1,), np.float32)
+    with pytest.raises(ValueError):
+        sampler = ActionSampler(action_space=action_space)

tests/enigma_logs/test/POET_test_0.eebmi2pf/params.json

@@ -0,0 +1,140 @@
+{
+  "_fake_gpus": false,
+  "action_space": null,
+  "actions_in_input_normalized": false,
+  "batch_mode": "truncate_episodes",
+  "callbacks": "<class 'ray.rllib.agents.callbacks.DefaultCallbacks'>",
+  "clip_actions": false,
+  "clip_param": 0.3,
+  "clip_rewards": null,
+  "collect_metrics_timeout": 180,
+  "compress_observations": false,
+  "create_env_on_driver": false,
+  "custom_eval_function": null,
+  "custom_resources_per_worker": {},
+  "eager_tracing": false,
+  "entropy_coeff": 0.0,
+  "entropy_coeff_schedule": null,
+  "env": "limited_zelda_custom",
+  "env_config": {
+    "environment_name": "limited_zelda_custom",
+    "global_observer_type": "ObserverType.VECTOR",
+    "level": 0,
+    "max_steps": 500,
+    "player_observer_type": "ObserverType.VECTOR",
+    "random_level_on_reset": false,
+    "record_video_config": {
+      "directory": "./videos",
+      "frequency": 5000
+    },
+    "yaml_file": "/home/rohindasari/research/watts_package/tests/solvers/../../example_levels/limited_zelda.yaml"
+  },
+  "env_task_fn": null,
+  "evaluation_config": {},
+  "evaluation_interval": null,
+  "evaluation_num_episodes": 10,
+  "evaluation_num_workers": 0,
+  "evaluation_parallel_to_training": false,
+  "exploration_config": {
+    "type": "StochasticSampling"
+  },
+  "explore": true,
+  "extra_python_environs_for_driver": {},
+  "extra_python_environs_for_worker": {},
+  "fake_sampler": false,
+  "framework": "torch",
+  "gamma": 0.99,
+  "grad_clip": null,
+  "horizon": null,
+  "ignore_worker_failures": false,
+  "in_evaluation": false,
+  "input": "sampler",
+  "input_config": {},
+  "input_evaluation": [
+    "is",
+    "wis"
+  ],
+  "kl_coeff": 0.2,
+  "kl_target": 0.01,
+  "lambda": 1.0,
+  "local_tf_session_args": {
+    "inter_op_parallelism_threads": 8,
+    "intra_op_parallelism_threads": 8
+  },
+  "log_level": "WARN",
+  "log_sys_usage": true,
+  "logger_config": null,
+  "lr": 5e-05,
+  "lr_schedule": null,
+  "metrics_smoothing_episodes": 100,
+  "min_iter_time_s": 0,
+  "model": {
+    "custom_model": "AIIDE_PINSKY_MODEL",
+    "custom_model_config": {}
+  },
+  "monitor": -1,
+  "multiagent": {
+    "count_steps_by": "env_steps",
+    "observation_fn": null,
+    "policies": {},
+    "policies_to_train": null,
+    "policy_map_cache": null,
+    "policy_map_capacity": 100,
+    "policy_mapping_fn": null,
+    "replay_mode": "independent"
+  },
+  "no_done_at_end": false,
+  "normalize_actions": true,
+  "num_cpus_for_driver": 1,
+  "num_cpus_per_worker": 1,
+  "num_envs_per_worker": 2,
+  "num_gpus": 0,
+  "num_gpus_per_worker": 0,
+  "num_sgd_iter": 30,
+  "num_workers": 1,
+  "observation_filter": "NoFilter",
+  "observation_space": null,
+  "optimizer": {},
+  "output": null,
+  "output_compress_columns": [
+    "obs",
+    "new_obs"
+  ],
+  "output_max_file_size": 67108864,
+  "placement_strategy": "PACK",
+  "postprocess_inputs": false,
+  "preprocessor_pref": "deepmind",
+  "record_env": false,
+  "remote_env_batch_wait_ms": 0,
+  "remote_worker_envs": false,
+  "render_env": false,
+  "rollout_fragment_length": 200,
+  "sample_async": false,
+  "sample_collector": "<class 'ray.rllib.evaluation.collectors.simple_list_collector.SimpleListCollector'>",
+  "seed": null,
+  "sgd_minibatch_size": 128,
+  "shuffle_buffer_size": 0,
+  "shuffle_sequences": true,
+  "simple_optimizer": true,
+  "soft_horizon": false,
+  "synchronize_filters": true,
+  "tf_session_args": {
+    "allow_soft_placement": true,
+    "device_count": {
+      "CPU": 1
+    },
+    "gpu_options": {
+      "allow_growth": true
+    },
+    "inter_op_parallelism_threads": 2,
+    "intra_op_parallelism_threads": 2,
+    "log_device_placement": false
+  },
+  "timesteps_per_iteration": 0,
+  "train_batch_size": 4000,
+  "use_critic": true,
+  "use_gae": true,
+  "vf_clip_param": 10.0,
+  "vf_loss_coeff": 1.0,
+  "vf_share_layers": -1
+}

tests/enigma_logs/test/POET_test_0.lvx863q6/params.json

@@ -0,0 +1,140 @@
+{
+  "_fake_gpus": false,
+  "action_space": null,
+  "actions_in_input_normalized": false,
+  "batch_mode": "truncate_episodes",
+  "callbacks": "<class 'ray.rllib.agents.callbacks.DefaultCallbacks'>",
+  "clip_actions": false,
+  "clip_param": 0.3,
+  "clip_rewards": null,
+  "collect_metrics_timeout": 180,
+  "compress_observations": false,
+  "create_env_on_driver": false,
+  "custom_eval_function": null,
+  "custom_resources_per_worker": {},
+  "eager_tracing": false,
+  "entropy_coeff": 0.0,
+  "entropy_coeff_schedule": null,
+  "env": "limited_zelda_custom",
+  "env_config": {
+    "environment_name": "limited_zelda_custom",
+    "global_observer_type": "ObserverType.VECTOR",
+    "level": 0,
+    "max_steps": 500,
+    "player_observer_type": "ObserverType.VECTOR",
+    "random_level_on_reset": false,
+    "record_video_config": {
+      "directory": "./videos",
+      "frequency": 5000
+    },
+    "yaml_file": "/home/rohindasari/research/watts_package/tests/solvers/../../example_levels/limited_zelda.yaml"
+  },
+  "env_task_fn": null,
+  "evaluation_config": {},
+  "evaluation_interval": null,
+  "evaluation_num_episodes": 10,
+  "evaluation_num_workers": 0,
+  "evaluation_parallel_to_training": false,
+  "exploration_config": {
+    "type": "StochasticSampling"
+  },
+  "explore": true,
+  "extra_python_environs_for_driver": {},
+  "extra_python_environs_for_worker": {},
+  "fake_sampler": false,
+  "framework": "torch",
+  "gamma": 0.99,
+  "grad_clip": null,
+  "horizon": null,
+  "ignore_worker_failures": false,
+  "in_evaluation": false,
+  "input": "sampler",
+  "input_config": {},
+  "input_evaluation": [
+    "is",
+    "wis"
+  ],
+  "kl_coeff": 0.2,
+  "kl_target": 0.01,
+  "lambda": 1.0,
+  "local_tf_session_args": {
+    "inter_op_parallelism_threads": 8,
+    "intra_op_parallelism_threads": 8
+  },
+  "log_level": "WARN",
+  "log_sys_usage": true,
+  "logger_config": null,
+  "lr": 5e-05,
+  "lr_schedule": null,
+  "metrics_smoothing_episodes": 100,
+  "min_iter_time_s": 0,
+  "model": {
+    "custom_model": "AIIDE_PINSKY_MODEL",
+    "custom_model_config": {}
+  },
+  "monitor": -1,
+  "multiagent": {
+    "count_steps_by": "env_steps",
+    "observation_fn": null,
+    "policies": {},
+    "policies_to_train": null,
+    "policy_map_cache": null,
+    "policy_map_capacity": 100,
+    "policy_mapping_fn": null,
+    "replay_mode": "independent"
+  },
+  "no_done_at_end": false,
+  "normalize_actions": true,
+  "num_cpus_for_driver": 1,
+  "num_cpus_per_worker": 1,
+  "num_envs_per_worker": 2,
+  "num_gpus": 0,
+  "num_gpus_per_worker": 0,
+  "num_sgd_iter": 30,
+  "num_workers": 1,
+  "observation_filter": "NoFilter",
+  "observation_space": null,
+  "optimizer": {},
+  "output": null,
+  "output_compress_columns": [
+    "obs",
+    "new_obs"
+  ],
+  "output_max_file_size": 67108864,
+  "placement_strategy": "PACK",
+  "postprocess_inputs": false,
+  "preprocessor_pref": "deepmind",
+  "record_env": false,
+  "remote_env_batch_wait_ms": 0,
+  "remote_worker_envs": false,
+  "render_env": false,
+  "rollout_fragment_length": 200,
+  "sample_async": false,
+  "sample_collector": "<class 'ray.rllib.evaluation.collectors.simple_list_collector.SimpleListCollector'>",
+  "seed": null,
+  "sgd_minibatch_size": 128,
+  "shuffle_buffer_size": 0,
+  "shuffle_sequences": true,
+  "simple_optimizer": true,
+  "soft_horizon": false,
+  "synchronize_filters": true,
+  "tf_session_args": {
+    "allow_soft_placement": true,
+    "device_count": {
+      "CPU": 1
+    },
+    "gpu_options": {
+      "allow_growth": true
+    },
+    "inter_op_parallelism_threads": 2,
+    "intra_op_parallelism_threads": 2,
+    "log_device_placement": false
+  },
+  "timesteps_per_iteration": 0,
+  "train_batch_size": 4000,
+  "use_critic": true,
+  "use_gae": true,
+  "vf_clip_param": 10.0,
+  "vf_loss_coeff": 1.0,
+  "vf_share_layers": -1
+}