I’ve been trying to simulate some traffic in a roundabout connected to a t-intersection and it keeps coming up with the error, “unexpected keyword argument”. I’ve included my road graph and the simulation code along with the entire error message, can anyone help? (Furthermore, this is done on google colab using “Waymax” simulation software.

import dataclasses
import jax
from jax import numpy as jnp
from matplotlib import pyplot as plt

from waymax import config as _config
from waymax import dataloader
from waymax import datatypes
from waymax import visualization

import numpy as np
def roundabout_t_intersection_model(inputs):
  # Extract the state from the inputs.
  state = inputs['state']

  # Get the roadgraph from the state.
  roadgraph = state.roadgraph

  # Define the roundabout center.
  roundabout_center = jnp.array([0.0, 0.0])

  # Define the roundabout radius.
  roundabout_radius = 10.0

  # Define the number of lanes in the roundabout.
  num_lanes = 2

  # Define the lane width.
  lane_width = 3.5

  # Create a list to store the lanes.
  lanes = []

  # Iterate over the number of lanes.
  for i in range(num_lanes):
    # Define the lane offset.
    lane_offset = i * lane_width

    # Define the lane center.
    lane_center = roundabout_center + lane_offset

    # Define the lane points.
    lane_points = jnp.array([
      [lane_center[0] + roundabout_radius, lane_center[1]],
      [lane_center[0] - roundabout_radius, lane_center[1]]
    ])

    # Add the lane to the list.
    lanes.append(lane_points)

  # Define the t-intersection center.
  t_intersection_center = jnp.array([10.0, 10.0])

  # Define the t-intersection road width.
  t_intersection_road_width = 10.0

  # Define the t-intersection lane width.
  t_intersection_lane_width = 3.5

  # Define the t-intersection lanes.
  t_intersection_lanes = [
    # Northbound lane.
    np.array([
      [t_intersection_center[0], t_intersection_center[1] + t_intersection_road_width / 2],
      [t_intersection_center[0], t_intersection_center[1] + t_intersection_lane_width / 2]
    ]),
    # Eastbound lane.
    np.array([
      [t_intersection_center[0] + t_intersection_road_width / 2, t_intersection_center[1]],
      [t_intersection_center[0] + t_intersection_lane_width / 2, t_intersection_center[1]]
    ]),
    # Southbound lane.
    np.array([
      [t_intersection_center[0], t_intersection_center[1] - t_intersection_road_width / 2],
      [t_intersection_center[0], t_intersection_center[1] - t_intersection_lane_width / 2]
    ]),
    # Westbound lane.
    np.array([
      [t_intersection_center[0] - t_intersection_road_width / 2, t_intersection_center[1]],
      [t_intersection_center[0] - t_intersection_lane_width / 2, t_intersection_center[1]]
    ])
  ]

  # Combine the roundabout and t-intersection lanes.
  lanes.extend(t_intersection_lanes)

  # Return the lanes.
  return lanes

# Define the network model.
network_model = roundabout_t_intersection_model

# Define the simulation configuration.
sim_config = dataclasses.replace(
    _config.WOD_1_1_0_TRAINING,
    max_num_objects=64,
    network_model=network_model
)
# Set the network model for the simulation configuration.
sim_config.network_model = network_model

# Create a simulator.
simulator = dataloader.simulator(config=sim_config)

# Generate a scenario.
scenario = next(simulator)

# Plot the scenario.
img = visualization.plot_simulator_state(scenario, use_log_traj=True)
mediapy.show_image(img)

# Simulate the scenario.
for _ in tqdm(range(scenario.remaining_timesteps)):
  scenario = simulator.step(scenario)

# Plot the final state of the scenario.
img = visualization.plot_simulator_state(scenario, use_log_traj=True)
mediapy.show_image(img)

# Error message.
TypeError                                 Traceback (most recent call last)
<ipython-input-21-9affebfac1bd> in <cell line: 14>()
     12 
     13 # Define the simulation configuration.
---> 14 sim_config = dataclasses.replace(
     15     _config.WOD_1_1_0_TRAINING,
     16     max_num_objects=64,

/usr/lib/python3.10/dataclasses.py in replace(obj, **changes)
   1451     # changes that aren't fields, this will correctly raise a
   1452     # TypeError.
-> 1453     return obj.__class__(**changes)

TypeError: DatasetConfig.__init__() got an unexpected keyword argument 'network_model'

This should be producing a small animation of vehicles on the created roadgraph.