SAC actor crashes in training (train.py) : TanhNormalProjectionNetwork.call() got unexpected network_state

Open Beginner friendly
  • Is something wrong?
#194 5 comments 0 reactions 0 assignees View on GitHub

@Sidharth1743 is already working on this.

Since Oct 9, 2026.

  • #195 by @Sidharth1743 — closed without merging
  • #199 by @Sidharth1743 — open

Assessment

Difficulty
1/5
Estimated time
Under an hour
Newbie friendliness
85/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
python, tensorflow
Domain
machine-learning

Research direction

Fix lives in _TanhNormalProjectionNetworkWrapper.call() in smart_control/reinforcement_learning/agents/networks/sac_networks.py: it must accept network_state (and strip step_type) without forwarding it to TanhNormalProjectionNetwork.call(), exactly as the wrapper in smart_control/notebooks/SAC_Demo.ipynb and CQL_Demo.ipynb already does. Reproduce by running the train module with --agent-type sac after building the starter buffer; done means the SAC agent constructs and training starts.

Written by the indexing model from the issue text.

Description

summary

poetry run python -m smart_control.reinforcement_learning.scripts.train --agent-type sac dies while constructing the SAC agent. The environment and starter buffer load successfully. The failure is in _TanhNormalProjectionNetworkWrapper in smart_control/reinforcement_learning/agents/networks/sac_networks.py.

TF-Agents 0.18.0 NestMap always calls the child layer with network_state. The script wrapper accepts only inputs, **kwargs, so network_state is forwarded into TanhNormalProjectionNetwork.call(), which does not take that argument.

reference

The same wrapper in smart_control/notebooks/SAC_Demo.ipynb (and CQL_Demo.ipynb) already names the argument and does not forward it:

def call(self, inputs, network_state=(), **kwargs):
  kwargs["outer_rank"] = self.predefined_outer_rank
  if "step_type" in kwargs:
    del kwargs["step_type"]
  return super().call(inputs, **kwargs)

Environment

  • Python 3.11
  • Linux
  • tensorflow==2.15.1
  • tf-agents==0.18.0 (as pinned in pyproject.toml)
  • Keras 2.15.0

Steps to reproduce

  1. Install with Poetry (poetry install --with dev) and generate the starter buffer:

    poetry run python -m smart_control.reinforcement_learning.scripts.populate_starter_buffer \
  --buffer-name starter

  2. Train SAC:

    poetry run python -m smart_control.reinforcement_learning.scripts.train \
  --starter-buffer-path smart_control/reinforcement_learning/data/starter_buffers/starter_seqlen2_exp500 \
  --agent-type sac \
  --experiment-name sac_run1

Actual behavior

SacAgent.__init__ calls actor_network.create_variables(), which calls the wrapper. Traceback ends with:

TypeError: Exception encountered when calling layer 'TanhNormalProjectionNetwork' (type _TanhNormalProjectionNetworkWrapper).

TanhNormalProjectionNetwork.call() got an unexpected keyword argument 'network_state'

Call arguments received by layer 'TanhNormalProjectionNetwork' (type _TanhNormalProjectionNetworkWrapper):
  • inputs=tf.Tensor(shape=(1, 256), dtype=float32)
  • kwargs={'network_state': (), 'step_type': '...', 'training': 'None'}
Dominant language
Python
Stars
132
Forks
64
No merged PRs in 30d

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from google/sbsim

All issues in google/sbsim

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.