Repository: stepjam/RLBench
Issue: Bug Report. `task.reset()` will not properly reset to the inital state
Root cause: Validation Gap

The `task.reset()` and `task.reset_to_demo(demo)` functions will not reset to the initial state, which leading to a degraded performance both in the training (using `test.reset_to_demo()`) and testing set (using `task.reset()`).

The following code would reproduce the error, with the saved model weights can be downloaded [here](https://www.dropbox.com/s/b7fo72hdoyjq9jm/model_26999_best_70.pth?dl=0).

```python
from rlbench.environment import Environment
from rlbench.action_modes import ArmActionMode, ActionMode
from rlbench.observation_config import ObservationConfig
from rlbench.tasks import ReachTarget

import torch.nn as nn
import torch.nn.functional as F
import torch
import numpy as np


# define a simple plain network
class PlainNet(nn.Module):
    def __init__(self, input_dim=12,  action_size=7):
        super(PlainNet, self).__init__()
        filter_list = [32, 'P', 64, 'P', 128, 'P', 256, 'P', 256]

        conv_layer = []
        in_channels = input_dim
        for v in filter_list:
            if v == 'P':
                conv_layer += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                conv_layer += [nn.Conv2d(in_channels, v, kernel_size=3, padding=1),
                               nn.ReLU(inplace=True)]
                in_channels = v

        self.conv_encoder = nn.Sequential(*conv_layer)
        self.fc_encoder = nn.Sequential(nn.Linear(filter_list[-1], filter_list[-1]),
                                        nn.ReLU(),
                                        nn.Linear(filter_list[-1], action_size))

    def forward(self, x):
        x = self.conv_encoder(x)
        x = F.adaptive_avg_pool2d(x, 1)
        x = self.fc_encoder(x.view(x.shape[0], x.shape[1]))
        return x


# load network
device = torch.device("cuda:0")
model = PlainNet(input_dim=12, action_size=7).to(device)

# load pre-trained network
model.load_state_dict(torch.load('model_26999_best_70.pth'))

# load demo (reach target)
obs_config = ObservationConfig()
obs_config.set_all(True)
action_mode = ActionMode(ArmActionMode.ABS_JOINT_VELOCITY)
env = Environment(action_mode, '', obs_config, False)

env.launch()
task = env.get_task(ReachTarget)
episode_length = 100
obs = None
model.eval()
success_rate = 0
np.random.seed(0)
with torch.no_grad():
    for i in range(50):
        for s in range(episode_length):
            if s == 0:
                descriptions, obs = task.reset()
                # env._scene._robot.arm.set_joint_positions(env._scene._start_arm_joint_pos, disable_dynamics=True)

            inputs = np.concatenate([obs.front_rgb,
                                     obs.left_shoulder_rgb,
                                     obs.right_shoulder_rgb,
                                     obs.wrist_rgb], axis=-1)
            inputs = torch.from_numpy(inputs).permute(2, 0, 1).unsqueeze(0).to(device)
            action = model(inputs)
            action = torch.cat([action[0], torch.tensor([1.0]).to(device)])
            action_np = action.detach().cpu().numpy()
            obs, reward, terminate = task.step(action_np)
            if reward == 1:
                success_rate += 1 / 50
                break
    env.shutdown()
print('Training Success Rate: {}'.format(success_rate))
```

The above code would produce a 26% success rate on reach_target task, adding this commened woud alleviate the issue ` env._scene._robot.arm.set_joint_positions(env._scene._start_arm_joint_pos, disable_dynamics=True)`  improving the performance to 32% success rate.

The training performance was affected more prominently, improving from 55% to 90% success rate, but still not **identical** to the saved training demo, which makes sense since they were in a simulation, and the noise is inevitable.

One potential reason which might lead to this issue is that the data format in the saved demo and `task.reset()` are different: saved demo uses uint8 format (from 0 to 255) in obs.images, but `task.reset()` uses float32 (from 0 to 1).








Hey Shikun @lorenmt 
Finally getting around to these. **Is this issue still relevant?** Am I remembering correctly in that this issue was related to your data loading, and not RLBench? Or was that a different issue I was thinking of?

Hi Stephen,

I think you can ignore this. It's still from the RLBench side, but I believe the performance drop is mainly due to the poorly designed training method. I believe we should expect some uncertainty occurred from the environment resetting. Maybe you should look into more on the yellow background training issue in the headless mode... which is more important compared to this.

Hey @lorenmt @stepjam , was the `task.reset_to_demo(demo)` fixed? Is it now reloading the scene exactly as it was recorded in the demo or is there some inevitable randomness? Could #279 be caused by this?