Seeking Guidance to Add Edge Computing Task Scheduling Environment in Mava

[AminAvan]

Thank you for providing this excellent GitHub repository!

I am planning to test an experiment by adding a new environment to Mava that simulates task scheduling in edge computing. As I am new to the domain of MARL & Mava, I was wondering if anyone would be open to mentoring or guiding me through this process.

Thank you in advance!

[sash-a]

Hi @AminAvan I assume that you can boil your problem down to a scheduling problem, in that case you can look at the Job Shop Scheduling environment from Jumanji. We don't directly support it, but we have a lot of envs from Jumanji so there are lots of examples.

Unfortunately I don't think any of us have the time to mentor you, but you are welcome to ask specific questions here and we can offer advice when we have the time 😄

[AminAvan]

Dear @sash-a,
Thank you very much for your prompt & kind response. I appreciate your guidance regarding the Job Shop Scheduling environment. I’ve reviewed it, and as you mentioned, it seems approximately suitable for my use case, though it may require some modifications or customization. Respectfully, I have two questions:

1. Regarding your note that "we don't directly support it" — is it possible for me to integrate this environment with Mava using the existing wrapper implementations?

2. You mentioned that “we have a lot of envs from Jumanji, so there are lots of examples.” Are any of these examples specifically related to scheduling? If so, would you kindly point me to where I can find them?

I sincerely appreciate your support and the helpful insights from you and the community.

[sash-a]

1. Yes it shouldn't bee to hard, you will need to write your own wrapper and maybe a custom network also. I'd encourage you to look through the jumanji training code for inspiration of how they format the network and you can try do something similar in Mava, or you can use jumanji directly if you don't need a multi-agent approach.

2. Nope, we don't have any scheduling envs in Mava as it's not area we've explored. You can see this file for examples of existing jumanji wrappers in Mava.

[AminAvan]

Thank you so much. I would appreciate your advice on the following points:

Regarding points (1) and (2), I would like to use MARL to solve the JobShop environment (Jumanji). However, I need some clarification regarding the following assumption (please correct me if I am mistaken):
"Since JobShop is a Jumanji environment—similar to Cleaner, Connector, LevelBasedForaging, and RobotWarehouse—I should be able to use JobShop with the existing wrapper in mava/wrappers/jumanji.py by simply adding from jumanji.environments.packing.job_shop import JobShop to mava/wrappers/jumanji.py."

from jumanji.env import Environment
from jumanji.environments.routing.cleaner import Cleaner
from jumanji.environments.routing.cleaner.constants import DIRTY, WALL
from jumanji.environments.routing.connector import Connector
from jumanji.environments.routing.connector.constants import (
    AGENT_INITIAL_VALUE,
    EMPTY,
    PATH,
    POSITION,
    TARGET,
)

My understanding is that if JobShop is indeed a Jumanji environment, then it should be accessible and integrated with Mava using the existing wrapper located at mava/wrappers/jumanji.py. However, if a new wrapper is required, I would appreciate clarification on the intended functionality of the current mava/wrappers/jumanji.py file. In that case, should I copy its content into a new file named mava/wrappers/job_shop_wrapper.py and then proceed step-by-step to modify it accordingly?

P.S.: I developed a specific wrapper for JobShop in Mava, available here, and I would appreciate any guidance on it.

Thank you again for your guidance.

[AminAvan]

Could you please clarify whether the value of time_limit: 500 in mava/configs/env/rware.yaml is derived from the value of num_epochs: 500 in jumanji/training/configs/env/robot_warehouse.yaml?

[AminAvan]

Dear @sash-a,
I have added the files I think are necessary to integrate JobShop into Mava for testing MARL-based algorithms. The files are listed below:

1. job_shop_wrapper.py that I make it inspired by the jumanji.py
2. Updating __init__.py in mava/wrappers/__init__.py.
3. I created job_shop.yaml in mava/configs/env/job_shop.yaml to have a environment file.
4. Finally, I created a scenario named jobshop-test-scenario.yaml in mava/configs/env/scenario/jobshop-test-scenario.yaml.
   However, when I run python mava/systems/ppo/anakin/ff_ippo.py env=job_shop in terminal, I encountered with the following errors:

An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
Error executing job with overrides: ['env=job_shop']
Traceback (most recent call last):
  File "/mnt/c/Users/100807003/PycharmProjects/Mava/mava/systems/ppo/anakin/ff_ippo.py", line 537, in hydra_entry_point
    eval_performance = run_experiment(cfg)
                       ^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/100807003/PycharmProjects/Mava/mava/systems/ppo/anakin/ff_ippo.py", line 409, in run_experiment
    env, eval_env = environments.make(config)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/100807003/PycharmProjects/Mava/mava/utils/make_env.py", line 292, in make
    return make_jumanji_env(config, add_global_state)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/100807003/PycharmProjects/Mava/mava/utils/make_env.py", line 134, in make_jumanji_env
    train_env = jumanji.make(config.env.scenario.name, generator=generator, **env_config)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/100807003/PycharmProjects/Mava/.venv/lib/python3.12/site-packages/jumanji/registration.py", line 172, in make
    return env_fn(*args, **env_fn_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: JobShop.__init__() got an unexpected keyword argument 'time_limit'

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.

Could you please advise me on this issue, any guidance is appreciated.
Thank you!

[AminAvan]

Update: I now understand that the wrappers in mava/wrappers/... do not convert a single-agent scenario into a multi-agent one. Therefore, a custom wrapper needs to be created for JobShop (based on Jumanji) to support a multi-agent structure.

[sash-a]

I developed a specific wrapper for JobShop in Mava, available here, and I would appreciate any guidance on it.

Sorry I don't have the time to review your wrapper

Could you please clarify whether the value of time_limit: 500 in mava/configs/env/rware.yaml is derived from the value of num_epochs: 500 in jumanji/training/configs/env/robot_warehouse.yaml?

No epochs is the number of training iterations and time limit is the number of steps in an episode

Could you please advise me on this issue, any guidance is appreciated.

Remove env_kwargs/time_limit from you config file, jobshop doesn't take that argument

[AminAvan]

Dear @sash-a, Thank you for your response and explanations.
Currently, I am working on making a proper 'network' for JobShop for smooth integration into Mava. However, I have encountered with the following errors in the last 3-days, would you please advise me about a proper approach to its solution? appreciate it.

Error executing job with overrides: ['env=job_shop']
Traceback (most recent call last):
  File "/mnt/c/Users/100807003/PycharmProjects/Mava/mava/systems/ppo/anakin/ff_ippo.py", line 537, in hydra_entry_point
    eval_performance = run_experiment(cfg)
                       ^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/100807003/PycharmProjects/Mava/mava/systems/ppo/anakin/ff_ippo.py", line 417, in run_experiment
    learn, actor_network, learner_state = learner_setup(
                                          ^^^^^^^^^^^^^^
  File "/mnt/c/Users/100807003/PycharmProjects/Mava/mava/systems/ppo/anakin/ff_ippo.py", line 335, in learner_setup
    actor_params = actor_network.init(actor_net_key, init_x)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/100807003/PycharmProjects/Mava/mava/networks/base.py", line 39, in __call__
    obs_embedding = self.torso(observation.agents_view)
                               ^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'Observation' object has no attribute 'agents_view'

Thank you for your time & attention.

[sash-a]

Ye so this is a tricky one, you're going to need to do a bit of work to get it working. The way we designed Mava was to be a library that always expects the Observation object to have an agents_view attribute. You can see the expected spec in types.py.

Problem is job shop has lots of different attributes in it's observation and so it doesn't really make sense to concat them all into an agents_view attribute, which is why we don't explicitly support jobshop. We plan to slowly start refactoring Mava so that it is easier to do this in the future, but I'd guess that won't be merged for ~6 months, maybe more.

For you to do this I'd suggest making a custom encoder, you can take inspiration from jumanji's jobshop networks. To actually add this to mava is a bit tricky, firstly you'd need to change the hydra instantiation of the networks and use your own ones with the custom encoder and you'd need to remove anywhere that has a reference to agents_view - you likely couldn't use our agent ID wrapper for example and will need a custom one, there may be other areas you'd also have to change which I can't think of right now

[AminAvan]

Dear @sash-a, Thank you for your kind response and explanations.
I have observed that the Mava printout displays the CONFIG when I execute the code. However, over the past few days, I have encountered an error (as indicated below) while attempting to modify sections job_shop_network.py, job_shop_wrapper.py, job_shop_network.yaml, and job_shop.yaml. Could you kindly provide guidance on an appropriate approach to resolve these issues? I greatly appreciate your assistance.

(id-mava) mywsl@UOFT23-MP1:/mnt/c/Users/MyStud/PycharmProjects/Mava$ python mava/systems/ppo/anakin/ff_ippo.py env=job_shop
An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
INFO:absl:Created `ArrayHandler` with primary_host=0, replica_id=0
init_x type: <class 'mava.types.Observation'>
Has agents_view: True
init_x attributes: ['__add__', '__annotations__', '__class__', '__class_getitem__', '__contains__', '__delattr__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__getitem__', '__getnewargs__', '__getstate__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__iter__', '__le__', '__len__', '__lt__', '__match_args__', '__module__', '__mul__', '__ne__', '__new__', '__orig_bases__', '__reduce__', '__reduce_ex__', '__repr__', '__rmul__', '__setattr__', '__sizeof__', '__slots__', '__str__', '__subclasshook__', '_asdict', '_field_defaults', '_fields', '_make', '_replace', 'action_mask', 'agents_view', 'count', 'index', 'step_count']
CONFIG
{
│   'logger': {
│   │   'base_exp_path': 'results/',
│   │   'loggers': {
│   │   │   'console': {'_target_': 'mava.utils.logger.ConsoleLogger', 'enabled': True},
│   │   │   'neptune': {
│   │   │   │   '_target_': 'mava.utils.logger.NeptuneLogger',
│   │   │   │   'enabled': False,
│   │   │   │   'run_id': None,
│   │   │   │   'project': 'Instadeep/mava-benchmark',
│   │   │   │   'tag': ['delete'],
│   │   │   │   'group_tag': ['delete'],
│   │   │   │   'detailed_logging': False,
│   │   │   │   'architecture_name': 'anakin',
│   │   │   │   'upload_json_data': False
│   │   │   },
│   │   │   'json': {'_target_': 'mava.utils.logger.JsonLogger', 'enabled': False, 'path': None, 'task_name': 'jobshop-test-scenario', 'env_name': 'JobShop', 'seed': 42},
│   │   │   'tensorboard': {'_target_': 'mava.utils.logger.TensorboardLogger', 'enabled': False}
│   │   },
│   │   'checkpointing': {
│   │   │   'save_model': False,
│   │   │   'save_args': {'save_interval_steps': 1, 'max_to_keep': 1, 'keep_period': None, 'checkpoint_uid': None, 'rel_dir': 'checkpoints'},
│   │   │   'load_model': False,
│   │   │   'load_args': {'checkpoint_uid': '', 'rel_dir': 'checkpoints'}
│   │   },
│   │   'system_name': 'ff_ippo'
│   },
│   'arch': {
│   │   'architecture_name': 'anakin',
│   │   'num_envs': 16,
│   │   'evaluation_greedy': False,
│   │   'num_evaluation': 122,
│   │   'num_eval_episodes': 32,
│   │   'num_absolute_metric_eval_episodes': 320,
│   │   'absolute_metric': True
│   },
│   'system': {
│   │   'total_timesteps': 4096000,
│   │   'num_updates': 1000,
│   │   'seed': 42,
│   │   'add_agent_id': 0,
│   │   'actor_lr': 0.00025,
│   │   'critic_lr': 0.00025,
│   │   'update_batch_size': 2,
│   │   'rollout_length': 128,
│   │   'ppo_epochs': 4,
│   │   'num_minibatches': 2,
│   │   'gamma': 0.99,
│   │   'gae_lambda': 0.95,
│   │   'clip_eps': 0.2,
│   │   'ent_coef': 0.01,
│   │   'vf_coef': 0.5,
│   │   'max_grad_norm': 0.5,
│   │   'decay_learning_rates': False,
│   │   'num_agents': 4,
│   │   'num_updates_per_eval': 8
│   },
│   'network': {
│   │   '_name_': 'job_shop_network',
│   │   'num_layers': 1,
│   │   'embed_dim': 64,
│   │   'num_heads': 1,
│   │   'use_swiglu': False,
│   │   'actor_network': {
│   │   │   '_target_': 'mava.networks.job_shop_network.JobShopActor',
│   │   │   'pre_torso': {'_target_': 'mava.networks.torsos.TransformerTorso', 'num_layers': 1, 'embed_dim': 64, 'num_heads': 1, 'input_dim': 128, 'use_swiglu': False},
│   │   │   'post_torso': {'_target_': 'mava.networks.torsos.MLPTorso', 'layer_sizes': [128], 'use_layer_norm': False, 'activation': 'relu'}
│   │   },
│   │   'critic_network': {
│   │   │   '_target_': 'mava.networks.job_shop_network.JobShopCritic',
│   │   │   'pre_torso': {'_target_': 'mava.networks.torsos.TransformerTorso', 'num_layers': 1, 'embed_dim': 64, 'num_heads': 1, 'input_dim': 128, 'use_swiglu': False},
│   │   │   'post_torso': {'_target_': 'mava.networks.torsos.MLPTorso', 'layer_sizes': [128], 'use_layer_norm': False, 'activation': 'relu'}
│   │   },
│   │   'q_network': {
│   │   │   'pre_torso': {'_target_': 'mava.networks.torsos.TransformerTorso', 'num_layers': 1, 'embed_dim': 64, 'num_heads': 1, 'input_dim': 128, 'use_swiglu': False},
│   │   │   'post_torso': {'_target_': 'mava.networks.torsos.MLPTorso', 'layer_sizes': [128], 'use_layer_norm': False, 'activation': 'relu'}
│   │   }
│   },
│   'env': {
│   │   'env_name': 'JobShop',
│   │   'eval_metric': 'episode_return',
│   │   'implicit_agent_id': True,
│   │   'log_win_rate': True,
│   │   'kwargs': {},
│   │   'scenario': {
│   │   │   'name': 'JobShop-v0',
│   │   │   'task_name': 'jobshop-test-scenario',
│   │   │   'task_config': {'num_jobs': 5, 'num_machines': 4, 'max_num_ops': 4, 'max_op_duration': 4},
│   │   │   'env_kwargs': {},
│   │   │   'network': {
│   │   │   │   'num_layers_machines': 1,
│   │   │   │   'num_layers_operations': 1,
│   │   │   │   'num_layers_joint_machines_jobs': 2,
│   │   │   │   'transformer_num_heads': 8,
│   │   │   │   'transformer_key_size': 16,
│   │   │   │   'transformer_mlp_units': [512]
│   │   │   },
│   │   │   'training': {'num_epochs': 65, 'num_learner_steps_per_epoch': 200, 'n_steps': 10, 'total_batch_size': 128},
│   │   │   'evaluation': {'eval_total_batch_size': 512, 'greedy_eval_total_batch_size': 512},
│   │   │   'a2c': {'normalize_advantage': False, 'discount_factor': 0.99, 'bootstrapping_factor': 0.95, 'l_pg': 1.0, 'l_td': 1.0, 'l_en': 0.5, 'learning_rate': 0.0001}
│   │   }
│   }
}
Error executing job with overrides: ['env=job_shop']
Traceback (most recent call last):
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/.venv/lib/python3.12/site-packages/jax/_src/lax/lax.py", line 150, in broadcast_shapes
    return _broadcast_shapes_cached(*shapes)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/.venv/lib/python3.12/site-packages/jax/_src/util.py", line 311, in wrapper
    return cached(config.trace_context(), *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/.venv/lib/python3.12/site-packages/jax/_src/util.py", line 304, in cached
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/.venv/lib/python3.12/site-packages/jax/_src/lax/lax.py", line 156, in _broadcast_shapes_cached
    return _broadcast_shapes_uncached(*shapes)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/.venv/lib/python3.12/site-packages/jax/_src/lax/lax.py", line 172, in _broadcast_shapes_uncached
    raise ValueError(f"Incompatible shapes for broadcasting: shapes={list(shapes)}")
ValueError: Incompatible shapes for broadcasting: shapes=[(16, 4, 6), (16, 6), ()]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/mava/systems/ppo/anakin/ff_ippo.py", line 552, in hydra_entry_point
    eval_performance = run_experiment(cfg)
                       ^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/mava/systems/ppo/anakin/ff_ippo.py", line 477, in run_experiment
    learner_output = learn(learner_state)
                     ^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/mava/systems/ppo/anakin/ff_ippo.py", line 284, in learner_fn
    learner_state, (episode_info, loss_info) = jax.lax.scan(
                                               ^^^^^^^^^^^^^
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/mava/systems/ppo/anakin/ff_ippo.py", line 103, in _update_step
    learner_state, (traj_batch, episode_metrics) = jax.lax.scan(
                                                   ^^^^^^^^^^^^^
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/mava/systems/ppo/anakin/ff_ippo.py", line 85, in _env_step
    actor_policy = actor_apply_fn(params.actor_params, last_timestep.observation)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/mava/networks/base.py", line 41, in __call__
    return self.action_head(obs_embedding, observation.action_mask)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/mava/networks/heads.py", line 55, in __call__
    masked_logits = jnp.where(
                    ^^^^^^^^^^
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/.venv/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py", line 1956, in where
    return util._where(condition, x, y)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/.venv/lib/python3.12/site-packages/jax/_src/numpy/util.py", line 450, in _where
    condition_arr, x_arr, y_arr = _broadcast_arrays(condition, x, y)
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/MyStud/PycharmProjects/Mava/.venv/lib/python3.12/site-packages/jax/_src/numpy/util.py", line 407, in _broadcast_arrays
    result_shape = lax.broadcast_shapes(*shapes)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: Incompatible shapes for broadcasting: shapes=[(16, 4, 6), (16, 6), ()]

Thank you once again for your time and consideration.

[sash-a]

There is an issue with your action masking. Looks like your logits third dim is num agents when it should be the second.

[AminAvan]

Dear @sash-a ,
Thank you for your explanations. I have two points to raise—some good news and a question.
Good news: I’m happy to report that the integration of JobShop into Mava appears to be successful. The training process ran smoothly, as shown below:

(id-mava) aminavan@UOFT23-MP1:/mnt/c/Users/MyUser/PycharmProjects/Mava$ python mava/systems/ppo/anakin/ff_ippo.py env=job_shop
An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
init_x type: <class 'mava.types.Observation'>
Has agents_view: True
init_x attributes: ['__add__', '__annotations__', '__class__', '__class_getitem__', '__contains__', '__delattr__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__getitem__', '__getnewargs__', '__getstate__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__iter__', '__le__', '__len__', '__lt__', '__match_args__', '__module__', '__mul__', '__ne__', '__new__', '__orig_bases__', '__reduce__', '__reduce_ex__', '__repr__', '__rmul__', '__setattr__', '__sizeof__', '__slots__', '__str__', '__subclasshook__', '_asdict', '_field_defaults', '_fields', '_make', '_replace', 'action_mask', 'agents_view', 'count', 'index', 'step_count']
CONFIG
{
│   'logger': {
│   │   'base_exp_path': 'results/',
│   │   'loggers': {
│   │   │   'console': {'_target_': 'mava.utils.logger.ConsoleLogger', 'enabled': True},
│   │   │   'neptune': {
│   │   │   │   '_target_': 'mava.utils.logger.NeptuneLogger',
│   │   │   │   'enabled': False,
│   │   │   │   'run_id': None,
│   │   │   │   'project': 'Instadeep/mava-benchmark',
│   │   │   │   'tag': ['delete'],
│   │   │   │   'group_tag': ['delete'],
│   │   │   │   'detailed_logging': False,
│   │   │   │   'architecture_name': 'anakin',
│   │   │   │   'upload_json_data': False
│   │   │   },
│   │   │   'json': {'_target_': 'mava.utils.logger.JsonLogger', 'enabled': False, 'path': None, 'task_name': 'jobshop-test-scenario', 'env_name': 'JobShop', 'seed': 42},
│   │   │   'tensorboard': {'_target_': 'mava.utils.logger.TensorboardLogger', 'enabled': False}
│   │   },
│   │   'checkpointing': {
│   │   │   'save_model': False,
│   │   │   'save_args': {'save_interval_steps': 1, 'max_to_keep': 1, 'keep_period': None, 'checkpoint_uid': None, 'rel_dir': 'checkpoints'},
│   │   │   'load_model': False,
│   │   │   'load_args': {'checkpoint_uid': '', 'rel_dir': 'checkpoints'}
│   │   },
│   │   'system_name': 'ff_ippo'
│   },
│   'arch': {
│   │   'architecture_name': 'anakin',
│   │   'num_envs': 16,
│   │   'evaluation_greedy': False,
│   │   'num_evaluation': 122,
│   │   'num_eval_episodes': 32,
│   │   'num_absolute_metric_eval_episodes': 320,
│   │   'absolute_metric': True
│   },
│   'system': {
│   │   'total_timesteps': 4096000,
│   │   'num_updates': 1000,
│   │   'seed': 42,
│   │   'add_agent_id': 0,
│   │   'actor_lr': 0.00025,
│   │   'critic_lr': 0.00025,
│   │   'update_batch_size': 2,
│   │   'rollout_length': 128,
│   │   'ppo_epochs': 4,
│   │   'num_minibatches': 2,
│   │   'gamma': 0.99,
│   │   'gae_lambda': 0.95,
│   │   'clip_eps': 0.2,
│   │   'ent_coef': 0.01,
│   │   'vf_coef': 0.5,
│   │   'max_grad_norm': 0.5,
│   │   'decay_learning_rates': False,
│   │   'num_agents': 4,
│   │   'num_updates_per_eval': 8
│   },
│   'network': {
│   │   '_name_': 'job_shop_network',
│   │   'num_jobs': 5,
│   │   'num_machines': 4,
│   │   'max_num_ops': 4,
│   │   'hidden_dim': 256,
│   │   'num_layers': 3,
│   │   'num_heads': 4,
│   │   'action_dim': 24,
│   │   'actor_network': {
│   │   │   'pre_torso': {'_target_': 'mava.networks.job_shop_network.JobShopTorso', 'num_jobs': 5, 'num_machines': 4, 'max_num_ops': 4, 'hidden_dim': 256, 'num_layers': 3},
│   │   │   'post_torso': {'_target_': 'mava.networks.torsos.MLPTorso', 'layer_sizes': [256, 128], 'use_layer_norm': False, 'activation': 'relu'}
│   │   },
│   │   'critic_network': {
│   │   │   'pre_torso': {'_target_': 'mava.networks.job_shop_network.JobShopTorso', 'num_jobs': 5, 'num_machines': 4, 'max_num_ops': 4, 'hidden_dim': 256, 'num_layers': 3},
│   │   │   'post_torso': {'_target_': 'mava.networks.torsos.MLPTorso', 'layer_sizes': [256, 128], 'use_layer_norm': False, 'activation': 'relu'}
│   │   },
│   │   'q_network': {
│   │   │   'pre_torso': {'_target_': 'mava.networks.job_shop_network.JobShopAttentionTorso', 'num_jobs': 5, 'num_machines': 4, 'max_num_ops': 4, 'hidden_dim': 256, 'num_heads': 4},        
│   │   │   'post_torso': {'_target_': 'mava.networks.torsos.MLPTorso', 'layer_sizes': [256, 128], 'use_layer_norm': False, 'activation': 'relu'}
│   │   }
│   },
│   'env': {
│   │   'env_name': 'JobShop',
│   │   'eval_metric': 'episode_return',
│   │   'implicit_agent_id': True,
│   │   'log_win_rate': True,
│   │   'kwargs': {},
│   │   'scenario': {
│   │   │   'name': 'JobShop-v0',
│   │   │   'task_name': 'jobshop-test-scenario',
│   │   │   'task_config': {'num_jobs': 5, 'num_machines': 4, 'max_num_ops': 4, 'max_op_duration': 4},
│   │   │   'action_dim': 24,
│   │   │   'env_kwargs': {},
│   │   │   'network': {
│   │   │   │   'num_layers_machines': 1,
│   │   │   │   'num_layers_operations': 1,
│   │   │   │   'num_layers_joint_machines_jobs': 2,
│   │   │   │   'transformer_num_heads': 8,
│   │   │   │   'transformer_key_size': 16,
│   │   │   │   'transformer_mlp_units': [512]
│   │   │   },
│   │   │   'training': {'num_epochs': 65, 'num_learner_steps_per_epoch': 200, 'n_steps': 10, 'total_batch_size': 128},
│   │   │   'evaluation': {'eval_total_batch_size': 512, 'greedy_eval_total_batch_size': 512},
│   │   │   'a2c': {'normalize_advantage': False, 'discount_factor': 0.99, 'bootstrapping_factor': 0.95, 'l_pg': 1.0, 'l_td': 1.0, 'l_en': 0.5, 'learning_rate': 0.0001}
│   │   }
│   }
}
MISC - Timestep: 32768
ACTOR - Steps per second: 554.251 | Episode length mean: 14.930 | Episode length std: 4.763 | Episode length min: 1 | Episode length max: 32 | Episode return mean: -36.937 | Episode return std: 34.082 | Episode return min: -106.000 | Episode return max: -4.000
TRAINER - Actor loss: -0.003 | Entropy: 0.148 | Total loss: 203.788 | Value loss: 407.583
EVALUATOR - Steps per second: 88.840 | Episode length mean: 14.031 | Episode length std: 4.276 | Episode length min: 1 | Episode length max: 22 | Episode return mean: -53.531 | Episode return std: 38.307 | Episode return min: -98.000 | Episode return max: -12.000
....
...
...

Question: Would it be possible to extend the JobShop machine model with custom parameters to capture machine-specific differences— for example, a communication attribute (e.g., M1.communication = 2, M2.communication = 4)? With such an attribute, reassigning a task from M1 to M2 would incur a +2 overhead, while the reverse assignment would incur -2. Is there an existing mechanism in the framework of Jumanji to support this, or would I need to modify the machine class myself? If I modify the machine definition in Jumanji, will it be difficult to integrate those changes when using Jumanji as an external library in Mava? Or would it be better to add the new parameters directly within Mava— for example, in job_shop_wrapper.py or another JobShop-related module?

Thank you again for your help and time!

[sash-a]

Hi @AminAvan, you definitely need to code this yourself. Whether you do it in the wrapper or env itself is up to you.

As for how to work with external libraries this is more of a python question, I'd encourage you to look it up online. We usually have jumanji installed locally uv pip install -e path/to/jumanji and then are able to edit and run changes easily, but there are a few other ways to do this.